From 0c635c667740cc1d0f6a99db8d36ca5d999c1dbc Mon Sep 17 00:00:00 2001 From: Alem Tuzlak Date: Sat, 17 Oct 2026 16:50:25 +0200 Subject: [PATCH 01/17] fix(ai-harness, ai-sandbox): coding commands get an env allowlist, one workspace per thread, and real paths checked under the lock --- .../src/first-party/coding/backend.ts | 71 ++- .../src/first-party/coding/index.ts | 4 +- .../src/first-party/coding/search.ts | 16 +- .../src/first-party/coding/snapshots.ts | 15 +- .../src/first-party/coding/workspace.ts | 412 ++++++++++++++---- .../ai-harness/tests/coding-backend.test.ts | 59 ++- packages/ai-harness/tests/coding-bash.test.ts | 2 + packages/ai-harness/tests/coding-edit.test.ts | 2 + .../ai-harness/tests/coding-formatter.test.ts | 2 + .../ai-harness/tests/coding-patch.test.ts | 2 + packages/ai-harness/tests/coding-read.test.ts | 2 + .../ai-harness/tests/coding-search.test.ts | 22 +- .../ai-harness/tests/coding-snapshots.test.ts | 32 +- .../ai-harness/tests/workspace-tools.test.ts | 191 +++++++- packages/ai-sandbox/src/journal.ts | 2 +- packages/ai-sandbox/src/workspace-backend.ts | 77 +++- .../tests/workspace-backend.test.ts | 105 ++++- 17 files changed, 881 insertions(+), 135 deletions(-) diff --git a/packages/ai-harness/src/first-party/coding/backend.ts b/packages/ai-harness/src/first-party/coding/backend.ts index 545ab72801..f3759cbb5d 100644 --- a/packages/ai-harness/src/first-party/coding/backend.ts +++ b/packages/ai-harness/src/first-party/coding/backend.ts @@ -55,10 +55,10 @@ export interface WorkspaceBackend { /** * The real path of `path`, with every link resolved, or `undefined` when * nothing is at `path`. Throws for a link to a target that is not there. - * The tools use it to refuse a link that leads out of the workspace. - * Optional: without it, the tools do not check links. + * The tools use it to refuse a link that leads out of the workspace. A + * backend that cannot resolve a path throws, and the tool call fails. */ - realpath?: (path: string) => Promise + realpath: (path: string) => Promise /** * The entries of a folder. A symbolic link (or a Windows junction) has * the type `'link'`, whatever it points to. The tools do not follow it. @@ -68,7 +68,9 @@ export interface WorkspaceBackend { ) => Promise> /** * Run a shell command and wait for it to end. `env` is added to the - * environment. A command that fails resolves with its exit code. + * environment. A command that fails resolves with its exit code. On + * {@link hostBackend}, a command gets only a few variables of this + * process (see {@link commandEnv}), plus `env`. */ exec: ( command: string, @@ -93,6 +95,13 @@ export interface WorkspaceBackend { } } +/** + * A workspace that `workspace.bind` of `workspaceTools()` gives: a backend, + * and its workspace folder `root`, an absolute path. `root` is the boundary + * of every tool. + */ +export type BoundWorkspace = WorkspaceBackend & { root: string } + /** What each workspace tool gets from `workspaceTools()`. */ export interface ToolEnv { backend: WorkspaceBackend @@ -100,12 +109,19 @@ export interface ToolEnv { root: string /** The hooks that plugins added. Read them when the tool runs. */ hooks: () => ReadonlyArray - /** Run `fn` when the work before it on `path` is done, one at a time. */ + /** + * Run `fn` when the work before it on `path` is done, one at a time. + * `path` comes from `reach`. Before `fn`, it checks again that `path` has + * no links, so a link made after `reach` cannot move the work. Else it + * throws, and `fn` does not run. + */ lock: (path: string, fn: () => Promise) => Promise /** - * The absolute path of `path`, a file or a folder. A path outside the - * workspace is refused, or needs the user's yes for its folder. The check - * uses the real path, so a link that leads out counts as outside. + * The absolute path of `path`, a file or a folder, with every link in the + * workspace resolved. Use it for the work, so the work goes to the file + * that was checked. A path outside the workspace is refused, or needs the + * user's yes for its folder. The check uses the real path, so a link that + * leads out counts as outside. */ reach: ( path: string, @@ -142,19 +158,42 @@ export function pathsOf(backend: WorkspaceBackend) { } /** - * The environment of a command: the environment of this process, plus - * `env`. Without `env` off Windows: `undefined`, so the command gets the - * environment of this process. On Windows it sets - * `NoDefaultCurrentDirectoryInExePath`. Then `cmd.exe` takes a program - * name like `rg` from the PATH only, and not from the current folder, where - * a cloned repo can put an `rg.cmd`. + * The variables of this process that a command gets. The others, like + * provider keys and `DATABASE_URL`, stay in the server. + */ +const KEPT_ENV = ['PATH', 'HOME', 'LANG', 'LC_ALL', 'TMPDIR', 'TEMP', 'TMP'] +/** On Windows, a command also gets these, so programs can start. */ +const KEPT_WINDOWS_ENV = [ + 'SystemRoot', + 'ComSpec', + 'PATHEXT', + 'WINDIR', + 'USERPROFILE', +] + +/** + * The environment of a command: the variables of this process in a short + * list (`PATH`, `HOME`, `LANG`, `LC_ALL`, `TMPDIR`, `TEMP`, `TMP`, and on + * Windows `SystemRoot`, `ComSpec`, `PATHEXT`, `WINDIR`, `USERPROFILE`), plus + * `env`. No other variable of this process goes to the command. On Windows + * it sets `NoDefaultCurrentDirectoryInExePath`. Then `cmd.exe` takes a + * program name like `rg` from the PATH only, and not from the current + * folder, where a cloned repo can put an `rg.cmd`. */ export function commandEnv( env: Record | undefined, platform: NodeJS.Platform, ) { - if (platform !== 'win32') return env && { ...process.env, ...env } - return { ...process.env, ...env, NoDefaultCurrentDirectoryInExePath: '1' } + const isWindows = platform === 'win32' + const names = isWindows ? [...KEPT_ENV, ...KEPT_WINDOWS_ENV] : KEPT_ENV + const kept: Record = {} + for (const name of names) { + // On Windows, `process.env` finds a name in any letter case. + const value = process.env[name] + if (value !== undefined) kept[name] = value + } + if (!isWindows) return { ...kept, ...env } + return { ...kept, ...env, NoDefaultCurrentDirectoryInExePath: '1' } } /** True for the error of a path that is not there. */ diff --git a/packages/ai-harness/src/first-party/coding/index.ts b/packages/ai-harness/src/first-party/coding/index.ts index a8adcfdcfe..b3bc11ecd8 100644 --- a/packages/ai-harness/src/first-party/coding/index.ts +++ b/packages/ai-harness/src/first-party/coding/index.ts @@ -1,9 +1,9 @@ // First-party plugins for a coding agent. This entry is Node only: the // plugins use the file system and the shell. export { workspaceTools } from './workspace' -export type { WorkspaceToolsOptions } from './workspace' +export type { WorkspaceBinder, WorkspaceToolsOptions } from './workspace' export { hostBackend } from './backend' -export type { WorkspaceBackend } from './backend' +export type { BoundWorkspace, WorkspaceBackend } from './backend' export { WorkspaceHooks } from '../workspace-hooks' export type { SearchProvider, WebToolsOptions } from './web' export { snapshots } from './snapshots' diff --git a/packages/ai-harness/src/first-party/coding/search.ts b/packages/ai-harness/src/first-party/coding/search.ts index 210ac24a87..f9c52f3c54 100644 --- a/packages/ai-harness/src/first-party/coding/search.ts +++ b/packages/ai-harness/src/first-party/coding/search.ts @@ -31,8 +31,16 @@ const FILES_NOTE = const MATCHES_NOTE = '[Not all matches are shown. Use a narrower pattern, glob, or folder.]' +/** + * `git` with the settings that run a program turned off: no fsmonitor and no + * hooks. The model can write the `.git/config` of the workspace, and git + * reads it. Run it with {@link GIT_ENV}. + */ +export const GIT = 'git -c core.fsmonitor= -c core.hooksPath=' +/** The env of every git command: git does not read the system config. */ +export const GIT_ENV = { GIT_CONFIG_NOSYSTEM: '1' } /** `git` applies the `.gitignore` files. `-z` ends each path with NUL. */ -const GIT_FILES = 'git ls-files -z --cached --others --exclude-standard' +const GIT_FILES = `${GIT} ls-files -z --cached --others --exclude-standard` /** rg flags for grep. One match more than shown tells that there are more. */ const RG_GREP = `--line-number --no-heading --color never --max-columns ${MAX_LINE} --max-columns-preview --max-filesize 1M --max-count ${MAX_MATCHES + 1}` @@ -89,8 +97,9 @@ async function bundledRg() { /** The command that starts rg on `backend`, or `undefined` without rg. */ async function findRg(backend: WorkspaceBackend) { - // The bundled binary is a host file. A sandbox cannot run it. - if (backend === hostBackend) { + // The bundled binary is a host file. A sandbox cannot run it. A bound + // workspace is a copy of `hostBackend` with a root, so compare `exec`. + if (backend.exec === hostBackend.exec) { const bundled = await bundledRg() if (bundled !== undefined) return quoteArg(bundled, process.platform) } @@ -338,6 +347,7 @@ export function searchTools(env: ToolEnv) { } const git = await backend.exec(GIT_FILES, { cwd: base, + env: GIT_ENV, timeoutMs: TIMEOUT_MS, }) if (git.exitCode === 0) { diff --git a/packages/ai-harness/src/first-party/coding/snapshots.ts b/packages/ai-harness/src/first-party/coding/snapshots.ts index 389b973923..6e1abcf1dd 100644 --- a/packages/ai-harness/src/first-party/coding/snapshots.ts +++ b/packages/ai-harness/src/first-party/coding/snapshots.ts @@ -4,7 +4,7 @@ import { defineCommand } from '../../commands' import { RevertFiles, RevertStanding, definePlugin } from '../../plugins' import { hostBackend } from './backend' import { isRecord } from '../../utils' -import { quoteArg, shellPlatform } from './search' +import { GIT, GIT_ENV, quoteArg, shellPlatform } from './search' import type { ModelMessage } from '@tanstack/ai' import type { WorkspaceBackend } from './backend' @@ -71,9 +71,13 @@ function shadowRepo(options: SnapshotsOptions) { const gitDir = `${options.dataDir}/${hash}` // The env wins over a GIT_DIR or GIT_INDEX_FILE from a git hook, so git // never writes to the project's repository. - const env = { GIT_DIR: gitDir, GIT_INDEX_FILE: `${gitDir}/index` } + const env = { + ...GIT_ENV, + GIT_DIR: gitDir, + GIT_INDEX_FILE: `${gitDir}/index`, + } // `core.autocrlf=false` keeps the line endings of every file. - const git = `git -c core.autocrlf=false --git-dir ${quote(gitDir)} --work-tree ${quote(options.root)}` + const git = `${GIT} -c core.autocrlf=false --git-dir ${quote(gitDir)} --work-tree ${quote(options.root)}` /** Run a command in the workspace. Throws with git's message when it fails. */ const exec = async (command: string) => { @@ -116,12 +120,13 @@ function shadowRepo(options: SnapshotsOptions) { /** Create the repository. `false` when the backend has no git. */ init: () => serial(async () => { - const probe = await backend.exec('git --version', { + const probe = await backend.exec(`${GIT} --version`, { + env: GIT_ENV, timeoutMs: TIMEOUT_MS, }) if (probe.exitCode !== 0) return false // A bare repository: it has no work tree of its own. - await exec(`git init --bare -q ${quote(gitDir)}`) + await exec(`${GIT} init --bare -q ${quote(gitDir)}`) // Keep the bytes of every file: no line ending change, no filter. await backend.writeFile( `${gitDir}/info/attributes`, diff --git a/packages/ai-harness/src/first-party/coding/workspace.ts b/packages/ai-harness/src/first-party/coding/workspace.ts index 37dfbd8fbb..d70acc7f06 100644 --- a/packages/ai-harness/src/first-party/coding/workspace.ts +++ b/packages/ai-harness/src/first-party/coding/workspace.ts @@ -1,4 +1,5 @@ import { BackgroundJobs, SessionSignal, definePlugin } from '../../plugins' +import { isRecord } from '../../utils' import { PERMISSION_MODES, PermissionDecisionCapability, @@ -13,9 +14,10 @@ import { patchPaths, patchTools } from './patch' import { readTools } from './read' import { searchTools } from './search' import { webTools } from './web' -import type { AnyTool } from '@tanstack/ai' +import type { AnyTool, ToolExecutionContext } from '@tanstack/ai' +import type { RunContext } from '../../types' import type { PermissionDecision, ToolResources } from '../permissions' -import type { ToolEnv, WorkspaceBackend } from './backend' +import type { BoundWorkspace, ToolEnv, WorkspaceBackend } from './backend' import type { WebToolsOptions } from './web' /** True for `true`, `y`, or `yes` (any case). */ @@ -23,8 +25,34 @@ function isYes(answer: unknown) { return answer === true || /^y(es)?$/i.test(String(answer).trim()) } -export interface WorkspaceToolsOptions { - root: string +/** Gives the workspace of each thread. */ +export interface WorkspaceBinder { + /** + * The workspace of the thread of `run`: a backend and its `root`. The + * tools call it once for each thread, before the first turn or at the + * first tool call, and keep the result. When it throws, the next turn or + * tool call calls it again. + */ + bind: (input: { run: RunContext }) => Promise +} + +/** Where the workspace is: one folder for every thread, or one per thread. */ +type WorkspaceSource = + | { + /** The workspace folder of every thread. */ + root: string + /** Where the files are and commands run. Default: {@link hostBackend}. */ + backend?: WorkspaceBackend + workspace?: undefined + } + | { + /** Gives the workspace of each thread. Its `root` is the boundary. */ + workspace: WorkspaceBinder + root?: undefined + backend?: undefined + } + +export type WorkspaceToolsOptions = WorkspaceSource & { bashTimeoutMs?: number /** * A path outside `root`. `'deny'` (default) refuses it. `'ask'` asks the @@ -33,8 +61,6 @@ export interface WorkspaceToolsOptions { * without a question. */ outside?: 'deny' | 'ask' - /** Where the files are and commands run. Default: {@link hostBackend}. */ - backend?: WorkspaceBackend /** * The tools that change files. `'auto'` (default): GPT models (an id that * starts with `gpt-`, or `o` and a digit, like `o3`) get `patch` and no @@ -62,7 +88,8 @@ export interface OutsideAccess { mode: () => unknown /** * The working folder of the thread, from `root` (the `cwd` setting of - * `session.configure`). Relative paths start there. Read at each call. + * `session.configure`). Relative paths start there. Read at each call. A + * folder outside the workspace is refused. */ cwd?: () => string | undefined /** @@ -75,9 +102,21 @@ export interface OutsideAccess { | undefined } +/** What the tools get from the session. */ +interface SessionAccess { + access?: OutsideAccess + hooks?: () => ReadonlyArray + note?: (text: string) => Promise + signal?: AbortSignal + jobs?: { started: (jobId: string) => void; ended: (jobId: string) => void } +} + /** A Windows network path (`\\server\share`). Not `\\?\` or `\\.\`. */ const isUnc = (path: string) => /^[\\/]{2}[^\\/?.]/.test(path) +/** An absolute path in any style: `/a`, `\a`, or `C:/a`. */ +const isAbsolutePath = (path: string) => /^([a-z]:)?[\\/]/i.test(path) + /** Is `full` the folder `dir`, or a path in it? `sep` is the separator. */ function within(dir: string, full: string, sep: string) { return full === dir || full.startsWith(dir.endsWith(sep) ? dir : dir + sep) @@ -111,25 +150,37 @@ function prefersPatch(model: string) { } /** - * The tools of {@link workspaceTools} for one session, and what each call - * touches for `PermissionResources`. `access` asks the user about a path - * outside the workspace when `options.outside` is `'ask'`. `hooks` are the - * hooks that plugins added. `note` tells the model that a background `bash` - * job ended, and `signal` kills the jobs that still run. + * True for the run context that the harness gives a tool call. It checks the + * fields that the tools read. The harness sets the others. */ -export function createWorkspaceTools( +function isRun(value: unknown): value is RunContext { + return ( + isRecord(value) && + typeof value.logId === 'string' && + typeof value.threadId === 'string' + ) +} + +/** The run of a tool call: `ctx.context.run`, or `undefined` without one. */ +function runOf(context: unknown) { + if (!isRecord(context) || !isRecord(context.context)) return undefined + const { run } = context.context + return isRun(run) ? run : undefined +} + +/** + * The tools for one workspace: `backend`, with the boundary `root`. Each + * call gets the paths through `reach`, which resolves every link, and works + * on the resolved path under `lock`. + */ +function bindTools( + workspace: { backend: WorkspaceBackend; root: string }, options: WorkspaceToolsOptions, - session: { - access?: OutsideAccess - hooks?: () => ReadonlyArray - note?: (text: string) => Promise - signal?: AbortSignal - jobs?: { started: (jobId: string) => void; ended: (jobId: string) => void } - } = {}, + session: SessionAccess, ) { - const backend = options.backend ?? hostBackend + const { backend } = workspace const paths = pathsOf(backend) - const root = paths.resolve(options.root) + const root = paths.resolve(workspace.root) const { access } = session // The real outside folders the user allowed in this session, and the // questions that wait for an answer, by folder, so parallel calls ask once. @@ -138,7 +189,12 @@ export function createWorkspaceTools( /** The working folder: `root`, or the thread's `cwd` in it. */ const here = () => { const cwd = access?.cwd?.() - return cwd ? paths.resolve(root, cwd) : root + if (!cwd) return root + const folder = paths.resolve(root, cwd) + if (!within(root, folder, paths.sep)) { + throw new Error(`The working folder ${cwd} is outside the workspace.`) + } + return folder } /** Ask once about `folder`. A yes allows it for the session. */ @@ -155,6 +211,20 @@ export function createWorkspaceTools( return answer } + /** + * `backend.realpath`. A backend from JavaScript can lack it. Then no path + * can be checked for links, so every call fails. + */ + const realpathOf = (path: string) => { + // ponytail: the type needs `realpath`. This check is for JavaScript. + if (typeof backend.realpath !== 'function') { + throw new Error( + 'The workspace backend has no realpath, so the tools cannot check links.', + ) + } + return backend.realpath(path) + } + /** * The real path of `full`, with every link resolved. For a path that is * not there yet: the real path of the nearest parent that is there, plus @@ -164,7 +234,7 @@ export function createWorkspaceTools( const rest: Array = [] let head = full for (;;) { - const real = backend.realpath ? await backend.realpath(head) : head + const real = await realpathOf(head) if (real !== undefined) return paths.join(real, ...rest) const parent = paths.dirname(head) // Not even the top folder is there. @@ -180,7 +250,7 @@ export function createWorkspaceTools( /** The real path of `root`. */ const realRoot = async () => { if (knownRoot !== undefined) return knownRoot - const real = backend.realpath ? await backend.realpath(root) : root + const real = await realpathOf(root) if (real === undefined) return realPath(root) knownRoot = real return real @@ -198,34 +268,14 @@ export function createWorkspaceTools( : real } + /** `full` with every link resolved, as a path in `root` when it is in it. */ + const resolved = async (full: string) => forRules(await realPath(full)) + const shown = (full: string) => (within(root, full, paths.sep) ? paths.relative(root, full) || '.' : full) .split(paths.sep) .join('/') - // What each call touches, as the tool resolves it, for the permission - // rules with a `resource`. - const file = (input: unknown) => [ - paths.resolve(here(), stringArg(input, 'path')), - ] - const folder = (input: unknown) => [ - paths.resolve(here(), optionalString(input, 'path') ?? '.'), - ] - const resources: Record = { - read_file: { paths: file }, - write_file: { paths: file }, - edit_file: { paths: file }, - list_files: { paths: folder }, - grep: { paths: folder }, - patch: { - paths: (input) => - patchPaths(stringArg(input, 'patch')).map((each) => - paths.resolve(here(), each), - ), - }, - bash: { commands: bashResources }, - } - /** * The rules see the path that the call names. A link can lead to a path * with a stricter rule, like `notes` to `.env`. So the rules check the @@ -235,12 +285,11 @@ export function createWorkspaceTools( path: string, tool: string, full: string, - realFull: string, + real: string, ) => { const decide = access?.rules?.() // `bash` touches commands. Its rules do not match paths. - if (!access || !decide || !resources[tool]?.paths) return - const real = await forRules(realFull) + if (!access || !decide || !PATH_TOOLS.has(tool)) return if (paths.relative(full, real) === '') return // `permissions()` checked `full` before the call. Act only when the // real path makes the decision stricter. @@ -286,7 +335,7 @@ export function createWorkspaceTools( } // The boundary check uses real paths, so a link inside the workspace that - // leads out of it counts as outside. + // leads out of it counts as outside. The tools work on the resolved path. const reach: ToolEnv['reach'] = async (path, tool, kind) => { const full = paths.resolve(here(), path) // `realpath` on a UNC path connects to its server, and Windows can send @@ -298,17 +347,30 @@ export function createWorkspaceTools( const folders = [await realRoot(), ...allowed] const isAllowed = folders.some((dir) => within(dir, real, paths.sep)) if (!isAllowed) await allowOutside(path, tool, kind, full, real) - await checkRealPath(path, tool, full, real) - return full + const inRoot = await forRules(real) + await checkRealPath(path, tool, full, inRoot) + return inRoot } + const lockPath = createLock() const env: ToolEnv = { backend, root, hooks: session.hooks ?? (() => []), - // ponytail: one lock per session. Sessions on one folder do not wait - // for each other. - lock: createLock(), + // ponytail: one lock per workspace binding. Sessions on one folder do + // not wait for each other. + lock: (path, fn) => + lockPath(path, async () => { + // `reach` checked `path` before the lock. A link made since then + // would move the work to another file, so check it again here. + const now = await resolved(path) + if (now !== path) { + throw new Error( + `${shown(path)} changed while the tool ran: it now leads to ${shown(now)}. Try again.`, + ) + } + return fn() + }), reach, shown, protectedIn: async (dir) => { @@ -316,7 +378,7 @@ export function createWorkspaceTools( if (!decide) return undefined // rg and the other engines do not follow links in `dir`. So only `dir` // itself can be a link, and its real path covers the files in it. - const real = await forRules(await realPath(dir)) + const real = await resolved(dir) return (path) => decide('read_file', [ paths.resolve(dir, path), @@ -339,11 +401,181 @@ export function createWorkspaceTools( ...(options.web === false ? [] : webTools(options.web)), ] - const prompt = + const intro = options.outside === 'ask' ? `Your workspace is ${root}. Paths are relative to it. A path outside it asks the user first.` : `Your workspace is ${root}. Paths are relative to it.` - return { tools, prompt, resources } + /** The prompt, with the working folder of the thread when it has one. */ + const prompt = () => { + const cwd = access?.cwd?.() + return cwd + ? `${intro} The working folder of this thread is ${paths.resolve(root, cwd)}. Relative paths start there.` + : intro + } + return { + tools, + prompt, + /** `path` as a full path, from the working folder. */ + fromHere: (path: string) => paths.resolve(here(), path), + } +} + +/** The tools whose rules match paths. `bash` rules match commands. */ +const PATH_TOOLS = new Set([ + 'read_file', + 'write_file', + 'edit_file', + 'list_files', + 'grep', + 'patch', +]) + +/** The tools of one workspace, from {@link bindTools}. */ +type BoundTools = ReturnType + +/** + * The tools for the workspace of each thread, from `binder`. A tool call + * reads its run from `ctx.context.run`, and works in the workspace of that + * thread. The first call of a thread binds it. The calls after it reuse + * that workspace. + */ +function threadTools( + binder: WorkspaceBinder, + options: WorkspaceToolsOptions, + session: SessionAccess, +) { + // By the log and the thread of the run. + const bindings = new Map>() + const bound = new Map() + const keyOf = (run: RunContext) => `${run.logId}\0${run.threadId}` + + /** The tools of the workspace of `run`. */ + const bind = (run: RunContext | undefined) => { + if (!run) { + throw new Error( + 'workspace.bind needs the run of the tool call. Call the tool in a harness session.', + ) + } + const key = keyOf(run) + const known = bindings.get(key) + if (known) return known + const binding = binder.bind({ run }).then((workspace) => { + const tools = bindTools( + { backend: workspace, root: workspace.root }, + options, + session, + ) + bound.set(key, tools) + return tools + }) + // A failed bind is not kept, so the next call tries again. + binding.catch(() => { + if (bindings.get(key) === binding) bindings.delete(key) + }) + bindings.set(key, binding) + return binding + } + + // These tools have the names and schemas of the tools of a workspace, and + // call the tools of the workspace of each call. Their own workspace is a + // placeholder: they never use it. + const templates = bindTools( + { backend: hostBackend, root: '/' }, + options, + {}, + ).tools + const tools = templates.map((template): AnyTool => { + const execute = async (args: unknown, context?: ToolExecutionContext) => { + const own = await bind(runOf(context)) + const tool = own.tools.find((each) => each.name === template.name) + if (!tool?.execute) throw new Error(`No tool ${template.name}.`) + return tool.execute(args, context) + } + return { ...template, execute } + }) + + /** The prompt of the workspace of `run`, once it is bound. */ + const prompt = (run: RunContext | undefined) => { + const own = run && bound.get(keyOf(run)) + return own + ? own.prompt() + : 'Paths are relative to the workspace. It is not ready yet.' + } + /** + * The root is not known before the call, so a relative path stays + * relative, from the working folder. `permissions()` reads it from its own + * `root`. + */ + const fromHere = (path: string) => { + const cwd = session.access?.cwd?.() + return cwd && !isAbsolutePath(path) ? `${cwd}/${path}` : path + } + return { tools, prompt, bind, fromHere } +} + +/** + * The tools of {@link workspaceTools} for one session, and what each call + * touches for `PermissionResources`. `access` asks the user about a path + * outside the workspace when `options.outside` is `'ask'`. `hooks` are the + * hooks that plugins added. `note` tells the model that a background `bash` + * job ended, and `signal` kills the jobs that still run. + * + * With `options.workspace`, each tool call works in the workspace of its + * thread. `bind(run)` binds the thread of `run`, and `prompt(run)` gives the + * prompt of its workspace. + */ +export function createWorkspaceTools( + options: WorkspaceToolsOptions, + session: SessionAccess = {}, +) { + const source = + options.workspace === undefined + ? fixedTools(options.root, options, session) + : threadTools(options.workspace, options, session) + const { fromHere } = source + // What each call touches, as the tool resolves it, for the permission + // rules with a `resource`. + const file = (input: unknown) => [fromHere(stringArg(input, 'path'))] + const folder = (input: unknown) => [ + fromHere(optionalString(input, 'path') ?? '.'), + ] + const resources: Record = { + read_file: { paths: file }, + write_file: { paths: file }, + edit_file: { paths: file }, + list_files: { paths: folder }, + grep: { paths: folder }, + patch: { + paths: (input) => patchPaths(stringArg(input, 'patch')).map(fromHere), + }, + bash: { commands: bashResources }, + } + return { + tools: source.tools, + resources, + prompt: source.prompt, + bind: source.bind, + } +} + +/** The tools for one workspace folder `root`, for every thread. */ +function fixedTools( + root: string, + options: WorkspaceToolsOptions, + session: SessionAccess, +) { + const own = bindTools( + { backend: options.backend ?? hostBackend, root }, + options, + session, + ) + return { + tools: own.tools, + prompt: () => own.prompt(), + bind: async () => own, + // A full path, from the working folder. + fromHere: own.fromHere, + } } /** @@ -362,22 +594,36 @@ function withEditStyle( } /** - * File, shell, and web tools for a coding agent, in `root`: `read_file`, - * `write_file`, `edit_file` or `patch` (see `editStyle`), `list_files`, - * `grep`, `bash`, `webfetch`, and `websearch` with a search provider. Edits, - * `bash`, and `webfetch` ask for approval through `permissions()`, and each - * call tells the permission rules which paths or commands it touches. A - * path outside `root` is refused, or with `outside: 'ask'` the user is - * asked first. A link that leads out of `root` counts as outside. Other - * plugins add {@link WorkspaceHooks} to run code after a read or a write. + * File, shell, and web tools for a coding agent: `read_file`, `write_file`, + * `edit_file` or `patch` (see `editStyle`), `list_files`, `grep`, `bash`, + * `webfetch`, and `websearch` with a search provider. Edits, `bash`, and + * `webfetch` ask for approval through `permissions()`, and each call tells + * the permission rules which paths or commands it touches. + * + * The workspace is `root` for every thread, or with `workspace.bind`, a + * workspace for each thread. Its root is the boundary of every tool: a path + * outside it is refused, or with `outside: 'ask'` the user is asked first. A + * link that leads out of the root counts as outside, and the working folder + * must be in the root. Other plugins add {@link WorkspaceHooks} to run code + * after a read or a write. * * The tools use `backend` for files and commands. The default, - * `hostBackend`, runs on this machine with the host's authority. Use a - * sandbox backend for code you do not trust. + * `hostBackend`, runs on this machine with the host's authority. Its + * commands get only a few variables of this process. Use a sandbox backend + * for code you do not trust. * * @example * ```ts * workspaceTools({ root: process.cwd(), outside: 'ask' }) + * // A folder for each thread: + * workspaceTools({ + * workspace: { + * bind: async ({ run }) => ({ + * ...hostBackend, + * root: `/srv/work/${run.threadId}`, + * }), + * }, + * }) * ``` */ export function workspaceTools(options: WorkspaceToolsOptions) { @@ -386,15 +632,14 @@ export function workspaceTools(options: WorkspaceToolsOptions) { optionalRequires: [PermissionDecisionCapability], setup: (ctx) => { const hooks = ctx.collect(WorkspaceHooks) - // The working folder is the thread's `cwd` setting, read at each call. - const cwd = () => ctx.session.settings().cwd - const { tools, prompt, resources } = createWorkspaceTools(options, { + const { tools, prompt, resources, bind } = createWorkspaceTools(options, { access: { // ponytail: D9, PR 7 rebuilds ask. The question goes to the // running turn, not to the run of the tool call that asks. ask: (message) => ctx.session.ask({ message }), mode: () => ctx.config.get('mode'), - cwd, + // The working folder is the thread's `cwd` setting, read at each call. + cwd: () => ctx.session.settings().cwd, // Read at each call, so permissions() can come after this plugin. rules: () => { const decide = ctx.getOptional(PermissionDecisionCapability) @@ -414,16 +659,19 @@ export function workspaceTools(options: WorkspaceToolsOptions) { signal: ctx.getOptional(SessionSignal) ?? ctx.resources.signal, jobs: ctx.getOptional(BackgroundJobs), }) - const paths = pathsOf(options.backend ?? hostBackend) - const folder = () => { - const dir = cwd() - return dir - ? ` The working folder of this thread is ${paths.resolve(options.root, dir)}. Relative paths start there.` - : '' - } return { tools, - prompts: [() => prompt + folder()], + prompts: [(turn) => prompt(turn.run)], + // Bind the workspace of the thread before the turn, so the prompt + // names its root. It adds no tools. + ...(options.workspace + ? { + discoverTools: async (turn) => { + await bind(turn.run) + return [] + }, + } + : {}), prepareTools: (turn) => withEditStyle(turn.tools, options.editStyle ?? 'auto', turn.model), contribute: [ diff --git a/packages/ai-harness/tests/coding-backend.test.ts b/packages/ai-harness/tests/coding-backend.test.ts index 18b1877fe3..a8669a686b 100644 --- a/packages/ai-harness/tests/coding-backend.test.ts +++ b/packages/ai-harness/tests/coding-backend.test.ts @@ -22,6 +22,22 @@ afterEach(async () => { const decoder = new TextDecoder() +/** Run `fn` with `vars` set in the env of this process, then put it back. */ +async function withEnv(vars: Record, fn: () => Promise) { + const before = Object.keys(vars).map( + (name) => [name, process.env[name]] as const, + ) + Object.assign(process.env, vars) + try { + await fn() + } finally { + for (const [name, value] of before) { + if (value === undefined) delete process.env[name] + else process.env[name] = value + } + } +} + describe('hostBackend', () => { it('writes into new folders, then reads, stats, and lists them', async () => { const file = join(dir, 'a', 'b.txt') @@ -118,21 +134,40 @@ describe('hostBackend', () => { }, ) - it('sets NoDefaultCurrentDirectoryInExePath in the env of commands on Windows only', () => { - expect(commandEnv(undefined, 'win32')).toMatchObject({ - NoDefaultCurrentDirectoryInExePath: '1', + it('does not give a variable of the server to a command', async () => { + const print = `node -e "console.log(process.env.SECRET_TOKEN || 'unset')"` + await withEnv({ SECRET_TOKEN: 'hunter2' }, async () => { + const run = await hostBackend.exec(print, { cwd: dir }) + expect(run.stdout).toBe('unset\n') + const job = hostBackend.spawn(print, { cwd: dir }) + await job.wait() + expect(job.output()).toBe('unset\n') }) - expect( - commandEnv( + }) + + it('gives a command only the listed variables of this process, plus env', async () => { + const server = { + SECRET_TOKEN: 'hunter2', + LANG: 'C.UTF-8', + SystemRoot: 'C:\\Windows', + } + await withEnv(server, async () => { + const linux = commandEnv({ AGENT: '1' }, 'linux') + expect(linux).toMatchObject({ LANG: 'C.UTF-8', AGENT: '1' }) + expect(linux).not.toHaveProperty('SECRET_TOKEN') + // The Windows variables and setting are for Windows only. + expect(linux).not.toHaveProperty('SystemRoot') + expect(linux).not.toHaveProperty('NoDefaultCurrentDirectoryInExePath') + const windows = commandEnv( { AGENT: '1', NoDefaultCurrentDirectoryInExePath: '' }, 'win32', - ), - ).toMatchObject({ AGENT: '1', NoDefaultCurrentDirectoryInExePath: '1' }) - // Off Windows, only `env` is added to the environment of this process. - expect(commandEnv(undefined, 'linux')).toBeUndefined() - expect(commandEnv({ AGENT: '1' }, 'linux')).toEqual({ - ...process.env, - AGENT: '1', + ) + expect(windows).toMatchObject({ + SystemRoot: 'C:\\Windows', + AGENT: '1', + NoDefaultCurrentDirectoryInExePath: '1', + }) + expect(windows).not.toHaveProperty('SECRET_TOKEN') }) }) }) diff --git a/packages/ai-harness/tests/coding-bash.test.ts b/packages/ai-harness/tests/coding-bash.test.ts index ed3a310f25..aeeeb655e9 100644 --- a/packages/ai-harness/tests/coding-bash.test.ts +++ b/packages/ai-harness/tests/coding-bash.test.ts @@ -45,6 +45,8 @@ function workspace( files.set(path, data) }, stat: async () => undefined, + // No links: each path is its own real path. + realpath: async (path) => path, readdir: async () => [], exec: async (command, options) => { execs.push({ command, options }) diff --git a/packages/ai-harness/tests/coding-edit.test.ts b/packages/ai-harness/tests/coding-edit.test.ts index 498f5b3b63..f667f320a6 100644 --- a/packages/ai-harness/tests/coding-edit.test.ts +++ b/packages/ai-harness/tests/coding-edit.test.ts @@ -21,6 +21,8 @@ function memoryBackend() { ) }, stat: async () => undefined, + // No links: each path is its own real path. + realpath: async (path) => path, readdir: async () => [], exec: async () => ({ exitCode: 0, stdout: '', stderr: '' }), } diff --git a/packages/ai-harness/tests/coding-formatter.test.ts b/packages/ai-harness/tests/coding-formatter.test.ts index d242e5a37f..0688ef0c4f 100644 --- a/packages/ai-harness/tests/coding-formatter.test.ts +++ b/packages/ai-harness/tests/coding-formatter.test.ts @@ -59,6 +59,8 @@ function project(files: Record, exec: Exec = ok, shell?: 'sh') { const data = store.get(path) return data && { type: 'file', size: data.length, mtimeMs: 0 } }, + // No links: each path is its own real path. + realpath: async (path) => path, readdir: async (dir) => [...store.keys()] .filter((path) => paths.dirname(path) === dir) diff --git a/packages/ai-harness/tests/coding-patch.test.ts b/packages/ai-harness/tests/coding-patch.test.ts index 0b8a373cce..6b5ee0c0aa 100644 --- a/packages/ai-harness/tests/coding-patch.test.ts +++ b/packages/ai-harness/tests/coding-patch.test.ts @@ -59,6 +59,8 @@ function workspace( const data = store.get(path) return data && { type: 'file', size: data.length, mtimeMs: 0 } }, + // No links: each path is its own real path. + realpath: async (path) => path, readdir: async () => [], exec: async () => ({ exitCode: 0, stdout: '', stderr: '' }), } diff --git a/packages/ai-harness/tests/coding-read.test.ts b/packages/ai-harness/tests/coding-read.test.ts index bf3d93bacf..e34e898e30 100644 --- a/packages/ai-harness/tests/coding-read.test.ts +++ b/packages/ai-harness/tests/coding-read.test.ts @@ -30,6 +30,8 @@ function memoryBackend(shell?: 'sh') { if (!data) return undefined return { type: 'file', size: data.length, mtimeMs: 0 } }, + // No links: each path is its own real path. + realpath: async (path) => path, readdir: async (path) => { const prefix = path + separator const paths = [...files.keys()].filter((file) => file.startsWith(prefix)) diff --git a/packages/ai-harness/tests/coding-search.test.ts b/packages/ai-harness/tests/coding-search.test.ts index ffeb512c3c..e3706b6ba7 100644 --- a/packages/ai-harness/tests/coding-search.test.ts +++ b/packages/ai-harness/tests/coding-search.test.ts @@ -38,7 +38,9 @@ interface ExecResult { const ok = (stdout: string): ExecResult => ({ exitCode: 0, stdout, stderr: '' }) const NOT_FOUND: ExecResult = { exitCode: 127, stdout: '', stderr: 'not found' } -const GIT = 'git ls-files -z --cached --others --exclude-standard' +// git runs with no fsmonitor and no hooks: the model can write .git/config. +const GIT = + 'git -c core.fsmonitor= -c core.hooksPath= ls-files -z --cached --others --exclude-standard' const NOT_GIT: ExecResult = { exitCode: 128, stdout: '', @@ -70,6 +72,8 @@ function fakeBackend( ]), ) const commands: Array = [] + /** The env of each command, in the order of `commands`. */ + const envs: Array | undefined> = [] const backend: WorkspaceBackend = { readFile: async (path) => { const data = byPath.get(path) @@ -84,6 +88,8 @@ function fakeBackend( if (!data) return undefined return { type: 'file', size: data.length, mtimeMs: 0 } }, + // No links: each path is its own real path. + realpath: async (path) => path, readdir: async (path) => { const prefix = path + sep const entries = new Map() @@ -97,12 +103,13 @@ function fakeBackend( .sort(([a], [b]) => (a < b ? -1 : 1)) .map(([name, type]) => ({ name, type })) }, - exec: async (command) => { + exec: async (command, options) => { commands.push(command) + envs.push(options?.env) return script(command) ?? NOT_FOUND }, } - return { backend, commands } + return { backend, commands, envs } } /** @@ -216,6 +223,15 @@ describe('list_files and grep engines', () => { }, ) + it('runs git without the system config', async () => { + const fake = fakeBackend(files, (command) => + command === GIT ? ok(gitFiles) : undefined, + ) + await searchOn(fake.backend).listFiles() + const gitEnv = fake.envs[fake.commands.indexOf(GIT)] + expect(gitEnv).toEqual({ GIT_CONFIG_NOSYSTEM: '1' }) + }) + it('runs the bundled rg with the host backend, and does not look on the PATH', async () => { const exec = scriptHostExec(() => ok('./a.ts\0')) expect(await searchOn(hostBackend, tmpdir()).listFiles()).toBe('a.ts') diff --git a/packages/ai-harness/tests/coding-snapshots.test.ts b/packages/ai-harness/tests/coding-snapshots.test.ts index 3dcfaae5eb..eca059c719 100644 --- a/packages/ai-harness/tests/coding-snapshots.test.ts +++ b/packages/ai-harness/tests/coding-snapshots.test.ts @@ -136,6 +136,7 @@ function fakeBackend( readFile: fail, writeFile: async () => undefined, stat: fail, + realpath: fail, readdir: fail, exec, } @@ -704,7 +705,36 @@ describe('snapshots on a cmd.exe backend', () => { await session.send('Hello') // `git --version`, then the init of the shadow repository. - expect(commands[1]).toMatch(/^git init --bare -q \^".+\^"$/) + expect(commands[1]).toMatch( + /^git -c core\.fsmonitor= -c core\.hooksPath= init --bare -q \^".+\^"$/, + ) + await host.close() + }) +}) + +describe('snapshots git safety', () => { + // The model can write the workspace files, so git must not run a program + // that a config names. + it('runs every git command with no fsmonitor, no hooks, and no system config', async () => { + const runs: Array<{ command: string; env?: Record }> = [] + const backend = fakeBackend(async (command, options) => { + runs.push({ command, env: options?.env }) + return { exitCode: 0, stdout: '', stderr: '' } + }) + const { host, session } = await open({ + replies: [() => text('Hi.')], + backend, + }) + + await session.send('Hello') + + expect(runs.length).toBeGreaterThan(1) + for (const run of runs) { + expect(run.command).toMatch( + /^git -c core\.fsmonitor= -c core\.hooksPath= /, + ) + expect(run.env).toMatchObject({ GIT_CONFIG_NOSYSTEM: '1' }) + } await host.close() }) }) diff --git a/packages/ai-harness/tests/workspace-tools.test.ts b/packages/ai-harness/tests/workspace-tools.test.ts index f23657d328..15650de074 100644 --- a/packages/ai-harness/tests/workspace-tools.test.ts +++ b/packages/ai-harness/tests/workspace-tools.test.ts @@ -18,6 +18,7 @@ import type { PermissionRule } from '../src/first-party/permissions' import type { WorkspaceBackend } from '../src/first-party/coding/backend' import type { OutsideAccess, + WorkspaceBinder, WorkspaceToolsOptions, } from '../src/first-party/coding/workspace' @@ -48,6 +49,8 @@ function memoryBackend(extra: Partial = {}) { if (!data) return undefined return { type: 'file', size: data.length, mtimeMs: 0 } }, + // No links: each path is its own real path. + realpath: async (path) => path, readdir: async (path) => { const prefix = path.endsWith('/') ? path : `${path}/` const entries = new Map() @@ -342,7 +345,7 @@ describe.each(fixtures)('workspace tools on the $name backend', (fixture) => { ).toBe('1\tbeta') expect(questions).toHaveLength(1) // The question names the real folder: the temp folder can be a link. - const real = (await backend.realpath?.(docs)) ?? docs + const real = (await backend.realpath(docs)) ?? docs expect(questions[0]).toContain(`Allow ${real} `) }) @@ -369,7 +372,9 @@ describe.each(fixtures)('workspace tools on the $name backend', (fixture) => { it('lists and searches an allowed outside folder with full paths', async () => { const { set } = asking([true]) const folder = paths.join(outside, 'docs') - const shown = folder.split(paths.sep).join('/') + // The tools show the real path: the temp folder can be a link. + const real = (await backend.realpath(folder)) ?? folder + const shown = real.split(paths.sep).join('/') expect(await runIn(set, 'list_files', { path: folder })).toBe( `${shown}/a.md\n${shown}/b.md`, ) @@ -407,7 +412,7 @@ describe('a sandbox backend with POSIX paths', () => { }) describe('the tools of workspaceTools', () => { - const names = (options: Partial = {}) => + const names = (options: Pick = {}) => createWorkspaceTools({ root: '/workspace', backend: memoryBackend().backend, @@ -1061,3 +1066,183 @@ describe('a network path', () => { expect(realpath.mock.calls.flat().join()).not.toContain('attacker') }) }) + +/** The result of the tool call `id`, as the model saw it in its last call. */ +function resultOf(calls: ReturnType['calls'], id: string) { + return JSON.stringify( + calls + .at(-1) + ?.messages.find( + (message: { toolCallId?: string }) => message.toolCallId === id, + )?.content, + ) +} + +describe('a workspace for each thread', () => { + const encoder = new TextEncoder() + + /** + * Workspaces in memory at `/ws/`. Thread `a` has `mine.txt` and a + * link `peek` to the folder of thread `b`, which has `secret.txt`. `binds` + * lists the thread of each `bind` call. With `failFirst`, the first bind + * throws. + */ + function threadWorkspaces(failFirst = false) { + const { backend, files } = memoryBackend() + files.set('/ws/a/mine.txt', encoder.encode('mine')) + files.set('/ws/b/secret.txt', encoder.encode('hunter2')) + const real = (path: string) => + path.replace(/^\/ws\/a\/peek(?=\/|$)/, '/ws/b') + const linked: WorkspaceBackend = { + ...backend, + realpath: async (path) => real(path), + readFile: (path) => backend.readFile(real(path)), + } + const binds: Array = [] + const workspace: WorkspaceBinder = { + bind: async ({ run }) => { + binds.push(run.threadId) + if (failFirst && binds.length === 1) throw new Error('No sandbox yet.') + return { ...linked, root: `/ws/${run.threadId}` } + }, + } + return { files, workspace, binds } + } + + /** Open thread `threadId` on `host`, with the tools of `workspace`. */ + function openThread( + host: ReturnType, + workspace: WorkspaceBinder, + adapter: AnyTextAdapter, + threadId: string, + ) { + return host.open( + defineHarness({ + name: 'test/workspace-bind', + adapter, + plugins: () => [workspaceTools({ workspace, web: false })], + }), + { threadId }, + ) + } + + it('keeps a thread in its root: no path through .. or a link reaches another root', async () => { + const { workspace } = threadWorkspaces() + const { adapter, calls } = mockAdapter([ + () => toolCall('read_file', { path: 'mine.txt' }, 'c0'), + () => toolCall('read_file', { path: '../b/secret.txt' }, 'c1'), + () => toolCall('read_file', { path: 'peek/secret.txt' }, 'c2'), + () => toolCall('read_file', { path: '/ws/b/secret.txt' }, 'c3'), + () => text('done'), + ]) + const host = createHarnessHost({ persistence: memoryPersistence() }) + const session = await openThread(host, workspace, adapter, 'a') + await session.send('read') + expect(resultOf(calls, 'c0')).toBe('"1\\tmine"') + for (const id of ['c1', 'c2', 'c3']) { + expect(resultOf(calls, id)).toContain('outside the workspace') + } + expect(JSON.stringify(calls.at(-1)?.messages)).not.toContain('hunter2') + await host.close() + }) + + it('binds each thread once, names its root in the prompt, and works there', async () => { + const { files, workspace, binds } = threadWorkspaces() + const { adapter, calls } = mockAdapter([ + () => toolCall('write_file', { path: 'a.txt', content: 'A' }, 'c0'), + () => toolCall('write_file', { path: 'a2.txt', content: 'A' }, 'c1'), + () => text('done'), + () => toolCall('write_file', { path: 'b.txt', content: 'B' }, 'c2'), + () => text('done'), + ]) + const host = createHarnessHost({ persistence: memoryPersistence() }) + const a = await openThread(host, workspace, adapter, 'a') + await a.send('write twice') + expect(JSON.stringify(calls[0].systemPrompts)).toContain( + 'Your workspace is /ws/a.', + ) + const b = await openThread(host, workspace, adapter, 'b') + await b.send('write once') + expect(JSON.stringify(calls[3].systemPrompts)).toContain( + 'Your workspace is /ws/b.', + ) + expect(binds).toEqual(['a', 'b']) + expect(files.has('/ws/a/a.txt')).toBe(true) + expect(files.has('/ws/a/a2.txt')).toBe(true) + expect(files.has('/ws/b/b.txt')).toBe(true) + await host.close() + }) + + it('binds again at the tool call when the bind before the turn failed', async () => { + const { workspace, binds } = threadWorkspaces(true) + const { adapter, calls } = mockAdapter([ + () => toolCall('read_file', { path: 'mine.txt' }, 'c0'), + () => text('done'), + ]) + const host = createHarnessHost({ persistence: memoryPersistence() }) + const session = await openThread(host, workspace, adapter, 'a') + await session.send('read') + expect(binds).toEqual(['a', 'a']) + expect(resultOf(calls, 'c0')).toBe('"1\\tmine"') + await host.close() + }) +}) + +describe('real paths', () => { + it('writes to the resolved path of a link', async () => { + const { backend, files } = memoryBackend() + files.set('/workspace/real.txt', new TextEncoder().encode('real')) + // `alias.txt` is a link to `real.txt`. + const linked: WorkspaceBackend = { + ...backend, + realpath: async (path) => + path === '/workspace/alias.txt' ? '/workspace/real.txt' : path, + } + const set = toolsOf({ root: '/workspace', backend: linked }) + expect( + await runIn(set, 'write_file', { path: 'alias.txt', content: 'new' }), + ).toBe('Wrote real.txt.') + expect(decoder.decode(files.get('/workspace/real.txt'))).toBe('new') + expect(files.has('/workspace/alias.txt')).toBe(false) + }) + + it('refuses a write when a link appears between the check and the write', async () => { + const { backend, files } = memoryBackend() + let checks = 0 + // The first check sees a file. Then a link to a folder outside takes + // its place, before the write. + const swapping: WorkspaceBackend = { + ...backend, + realpath: async (path) => { + if (path !== '/workspace/swap.txt') return path + checks += 1 + return checks === 1 ? path : '/elsewhere/swap.txt' + }, + } + const set = toolsOf({ root: '/workspace', backend: swapping }) + await expect( + runIn(set, 'write_file', { path: 'swap.txt', content: 'x' }), + ).rejects.toThrow('swap.txt changed while the tool ran') + expect([...files.keys()]).toEqual([]) + }) + + it('refuses every path when the backend has no realpath', async () => { + const { backend } = memoryBackend() + // A backend from JavaScript, where the type cannot ask for realpath. + Reflect.deleteProperty(backend, 'realpath') + const set = toolsOf({ root: '/workspace', backend }) + await expect(runIn(set, 'read_file', { path: 'a.txt' })).rejects.toThrow( + 'The workspace backend has no realpath', + ) + }) + + it('refuses a working folder outside the workspace', async () => { + const set = toolsOf( + { root: '/workspace', backend: memoryBackend().backend }, + { ask: async () => 'y', mode: () => 'default', cwd: () => '../other' }, + ) + await expect(runIn(set, 'read_file', { path: 'a.txt' })).rejects.toThrow( + 'The working folder ../other is outside the workspace.', + ) + }) +}) diff --git a/packages/ai-sandbox/src/journal.ts b/packages/ai-sandbox/src/journal.ts index 4ae858e3fa..9cda90d5d5 100644 --- a/packages/ai-sandbox/src/journal.ts +++ b/packages/ai-sandbox/src/journal.ts @@ -159,7 +159,7 @@ export function exitSentinelLine( } /** Single-quote a shell word, escaping embedded single quotes POSIX-style. */ -function shellQuote(value: string): string { +export function shellQuote(value: string): string { return `'${value.replaceAll("'", `'\\''`)}'` } diff --git a/packages/ai-sandbox/src/workspace-backend.ts b/packages/ai-sandbox/src/workspace-backend.ts index b5e3c51665..6293bb7930 100644 --- a/packages/ai-sandbox/src/workspace-backend.ts +++ b/packages/ai-sandbox/src/workspace-backend.ts @@ -1,8 +1,12 @@ +import { posix } from 'node:path' +import { shellQuote } from './journal' import type { WorkspaceBackend } from '@tanstack/ai-harness/plugins/coding' import type { ExecResult, SandboxHandle } from './contracts' /** The most output a background job keeps, in characters. */ const MAX_OUTPUT = 10 * 1024 * 1024 +/** `realpath` follows at most this many links, like Linux. */ +const MAX_LINKS = 40 /** * A `WorkspaceBackend` over a live sandbox. With it, `workspaceTools()` @@ -17,7 +21,9 @@ const MAX_OUTPUT = 10 * 1024 * 1024 * `capabilities.backgroundProcesses`. * - `stat` gives `mtimeMs: 0`. The sandbox file system does not give the * time of the last change. - * - There is no `realpath`, so the tools do not check where links go. + * - `realpath` resolves links in the sandbox: `lstat` finds each link, and + * `readlink` in the sandbox reads its target. A sandbox without `lstat` + * cannot find links, so `realpath` throws, and the tools refuse each path. * * @param handle - The sandbox, from `provider.create()` or `provider.resume()`. * @@ -61,6 +67,74 @@ export function sandboxWorkspaceBackend(handle: SandboxHandle) { } } + /** The target of the link `path`, as `readlink` in the sandbox gives it. */ + const readlink = async (path: string) => { + // A relative name in its folder: the sandbox maps the folder, not the + // arguments, to its own files. + const name = shellQuote(posix.basename(path)) + const { exitCode, stdout, stderr } = await handle.process.exec( + `readlink -- ${name}`, + { cwd: posix.dirname(path) }, + ) + if (exitCode !== 0) { + throw new Error(`Cannot read the link ${path}: ${stderr.trim()}`) + } + return stdout.replace(/\n$/, '') + } + + /** + * `path` with every link resolved, part by part, or `undefined` when + * nothing is at `path`. Throws for a link to a missing target, and when + * there are more than 40 links, like a loop. + */ + const realpath: WorkspaceBackend['realpath'] = async (path) => { + const { lstat } = fs + if (!lstat) { + throw new Error( + `This sandbox has no lstat, so links in ${path} cannot be checked.`, + ) + } + let parts = posix.resolve('/', path).split('/').filter(Boolean) + // The first parts that came from a link target. A missing one means a + // link to a missing target. A missing part of `path` itself means + // nothing is at `path`. + let fromLink = 0 + let real = '/' + let links = 0 + while (parts.length > 0) { + const [name = '', ...rest] = parts + parts = rest + const isFromLink = fromLink > 0 + fromLink = Math.max(0, fromLink - 1) + if (name === '.') continue + if (name === '..') { + real = posix.dirname(real) + continue + } + const next = posix.join(real, name) + const info = await lstat(next) + if (!info) { + if (!isFromLink) return undefined + throw new Error(`ENOENT: ${path} is a link to ${next}, not there.`) + } + if (info.type !== 'symlink') { + real = next + continue + } + links += 1 + if (links > MAX_LINKS) + throw new Error(`ELOOP: too many links in ${path}.`) + const target = await readlink(next) + // An absolute target starts again at `/`. A relative one is in `real`. + if (target.startsWith('/')) real = '/' + const targetParts = target.split('/').filter(Boolean) + parts = [...targetParts, ...parts] + // The parts after it can come from an earlier link too. + fromLink += targetParts.length + } + return real + } + const spawn: WorkspaceBackend['spawn'] = (command, options) => { // ponytail: keeps the last 10 MiB of a long job, like `hostBackend`. let output = '' @@ -107,6 +181,7 @@ export function sandboxWorkspaceBackend(handle: SandboxHandle) { await fs.remove(path) }, stat, + realpath, readdir: async (path) => { const entries = await fs.list(path) const { lstat } = fs diff --git a/packages/ai-sandbox/tests/workspace-backend.test.ts b/packages/ai-sandbox/tests/workspace-backend.test.ts index ff0d2fc9f2..bb84fbd7bb 100644 --- a/packages/ai-sandbox/tests/workspace-backend.test.ts +++ b/packages/ai-sandbox/tests/workspace-backend.test.ts @@ -1,3 +1,4 @@ +import { posix } from 'node:path' import { describe, expect, it } from 'vitest' import { EventType } from '@tanstack/ai' import { memoryPersistence } from '@tanstack/ai-persistence' @@ -98,7 +99,7 @@ function memoryFs(seed: FsSeed, withLstat = true) { exists: async (path) => files.has(target(path)) || dirs.has(target(path)), ...(withLstat ? { lstat } : {}), } - return { fs, files, dirs } + return { fs, files, dirs, links } } async function* fromChunks(chunks: Array) { @@ -112,12 +113,26 @@ async function* fromChunks(chunks: Array) { * * `spawn` prints ` out` and ` err`. `serve` runs until it * is killed, then ends with 143. `nope` does not start. + * + * `readlink -- ''` prints the target of the link `name` in the folder, + * from `links`, like `readlink` in a sandbox. */ -function fakeProcesses() { +function fakeProcesses(links: ReadonlyMap = new Map()) { const aborted: Array = [] const killed: Array = [] const process: SandboxHandle['process'] = { exec: (command, options = {}) => { + const readlink = /^readlink -- '(.+)'$/.exec(command) + if (readlink) { + const target = links.get( + posix.join(options.cwd ?? '/', readlink[1] ?? ''), + ) + return Promise.resolve( + target === undefined + ? { exitCode: 1, stdout: '', stderr: 'not a link' } + : { exitCode: 0, stdout: `${target}\n`, stderr: '' }, + ) + } if (command.startsWith('sleep')) { options.signal?.addEventListener('abort', () => aborted.push(command)) return new Promise(() => {}) @@ -163,8 +178,8 @@ function fakeSandbox( seed: FsSeed = {}, options: { withLstat?: boolean; caps?: SandboxCapabilities } = {}, ) { - const { fs, files, dirs } = memoryFs(seed, options.withLstat) - const { process, aborted, killed } = fakeProcesses() + const { fs, files, dirs, links } = memoryFs(seed, options.withLstat) + const { process, aborted, killed } = fakeProcesses(links) const handle: SandboxHandle = { ...makeFakeHandle('sandbox-1', 'fake', options.caps), fs, @@ -186,10 +201,9 @@ const SEED: FsSeed = { } describe('sandboxWorkspaceBackend files', () => { - it('uses POSIX paths and quoting, and checks no links', () => { + it('uses POSIX paths and quoting', () => { const { backend } = fakeSandbox() expect(backend.shell).toBe('sh') - expect(backend).not.toHaveProperty('realpath') }) it('reads and writes the bytes of a file in the sandbox', async () => { @@ -271,6 +285,53 @@ describe('sandboxWorkspaceBackend files', () => { }) }) +describe('sandboxWorkspaceBackend realpath', () => { + /** SEED, plus a link out of the workspace, a relative link, and a loop. */ + const LINKS: FsSeed = { + files: { ...SEED.files, '/home/app/.ssh/id_rsa': 'KEY' }, + dirs: [...(SEED.dirs ?? []), '/home', '/home/app', '/home/app/.ssh'], + links: { + ...SEED.links, + '/workspace/docs': '/home/app', + '/workspace/src/up': '../a.txt', + '/workspace/loop-a': 'loop-b', + '/workspace/loop-b': 'loop-a', + }, + } + + it.each([ + ['/workspace/a.txt', '/workspace/a.txt'], + ['/workspace/to-file', '/workspace/a.txt'], + ['/workspace/to-src', '/workspace/src'], + ['/workspace/src/up', '/workspace/a.txt'], + ['/workspace/docs/.ssh/id_rsa', '/home/app/.ssh/id_rsa'], + ])('resolves %s to %s', async (path, real) => { + const { backend } = fakeSandbox(LINKS) + expect(await backend.realpath(path)).toBe(real) + }) + + it('gives undefined when nothing is at the path', async () => { + const { backend } = fakeSandbox(LINKS) + expect(await backend.realpath('/workspace/gone')).toBeUndefined() + expect(await backend.realpath('/workspace/to-src/gone.ts')).toBeUndefined() + }) + + it('throws for a link to a missing target, and for a loop of links', async () => { + const { backend } = fakeSandbox(LINKS) + await expect(backend.realpath('/workspace/broken')).rejects.toThrow( + 'ENOENT', + ) + await expect(backend.realpath('/workspace/loop-a')).rejects.toThrow('ELOOP') + }) + + it('throws when the sandbox has no lstat, so no link goes unchecked', async () => { + const { backend } = fakeSandbox(LINKS, { withLstat: false }) + await expect(backend.realpath('/workspace/a.txt')).rejects.toThrow( + 'This sandbox has no lstat', + ) + }) +}) + describe('sandboxWorkspaceBackend exec', () => { it('runs a command in the sandbox with the folder and the env', async () => { const { backend } = fakeSandbox() @@ -455,4 +516,36 @@ describe('sandboxWorkspaceBackend with workspaceTools', () => { expect.stringContaining('Your workspace is /workspace.'), ) }) + + it('refuses a read through a link that leads out of /workspace', async () => { + const sandbox = fakeSandbox({ + files: { '/home/app/.ssh/id_rsa': 'PRIVATE KEY' }, + dirs: ['/workspace', '/home', '/home/app', '/home/app/.ssh'], + links: { '/workspace/docs': '/home/app' }, + }) + const { adapter, calls } = scriptedModel([ + toolCall('c1', 'read_file', { path: 'docs/.ssh/id_rsa' }), + ]) + const host = createHarnessHost({ persistence: memoryPersistence() }) + const session = await host.open( + defineHarness({ + name: 'test/sandbox-links', + adapter, + plugins: () => [ + workspaceTools({ + root: '/workspace', + backend: sandbox.backend, + web: false, + }), + ], + }), + { threadId: 't' }, + ) + await session.send('read the key') + await host.close() + + const seen = JSON.stringify(calls[1]?.messages) + expect(seen).toContain('outside the workspace') + expect(seen).not.toContain('PRIVATE KEY') + }) }) From 5e10199674516a5e949524e3a83f7d7a6cb5ef99 Mon Sep 17 00:00:00 2001 From: Alem Tuzlak Date: Sat, 17 Oct 2026 17:06:20 +0200 Subject: [PATCH 02/17] wip: unfinished node work, saved before a shutdown (not reviewed) --- .../ai-harness/src/first-party/coding/edit.ts | 79 +++- .../src/first-party/coding/patch.ts | 30 +- .../src/first-party/coding/snapshots.ts | 411 ++++++++++++++---- .../src/first-party/coding/workspace.ts | 43 +- .../ai-harness/tests/coding-locks.test.ts | 204 +++++++++ .../ai-harness/tests/coding-patch.test.ts | 3 +- .../ai-harness/tests/coding-snapshots.test.ts | 23 +- 7 files changed, 687 insertions(+), 106 deletions(-) create mode 100644 packages/ai-harness/tests/coding-locks.test.ts diff --git a/packages/ai-harness/src/first-party/coding/edit.ts b/packages/ai-harness/src/first-party/coding/edit.ts index fdb0da13b9..c73b196d19 100644 --- a/packages/ai-harness/src/first-party/coding/edit.ts +++ b/packages/ai-harness/src/first-party/coding/edit.ts @@ -118,12 +118,73 @@ export function findMatches(content: string, find: string) { } /** - * `write_file` and `edit_file`. After each write, the `afterWrite` hooks run - * for the file. `edit_file` finds the old text with {@link findMatches}. It - * writes the new text with the line endings of the file, and keeps a byte - * order mark. + * What one workspace binding saw of each file: its size and change time + * when a tool read it or wrote it. A write checks that the file is still as + * the binding saw it, so an edit of another session or agent is never lost. + * Call each function under the lock of the path. The paths come from + * `env.reach`. + * + * ponytail: a change that keeps the size and the `mtimeMs` is not seen. + */ +export function fileVersions(env: ToolEnv) { + const seen = new Map() + return { + /** Remember the file at `path` as it is now: a tool read or wrote it. */ + saw: async (path: string) => { + const info = await env.backend.stat(path) + if (info?.type === 'file') { + seen.set(path, { size: info.size, mtimeMs: info.mtimeMs }) + } else { + seen.delete(path) + } + }, + /** Forget `path`: a tool removed it. */ + forget: (path: string) => { + seen.delete(path) + }, + /** Is something at `path`? Call it before the lock. */ + exists: async (path: string) => (await env.backend.stat(path)) !== undefined, + /** + * Throws when the file at `path` changed since a tool read or wrote it. + * For a file that no tool read, it throws only when the tool saw no file + * there before the lock (`missing`), and a file is there now. + */ + check: async (path: string, { missing }: { missing: boolean }) => { + const known = seen.get(path) + const now = await env.backend.stat(path) + const shown = env.shown(path) + if (known) { + const isSame = + now?.type === 'file' && + now.size === known.size && + now.mtimeMs === known.mtimeMs + if (!isSame) { + throw new Error( + `${shown} changed since it was read. Read it again before you edit it.`, + ) + } + return + } + if (missing && now !== undefined) { + throw new Error( + `${shown} was created while the tool waited. Read it before you edit it.`, + ) + } + }, + } +} + +/** The file record of a binding, from {@link fileVersions}. */ +export type FileVersions = ReturnType + +/** + * `write_file` and `edit_file`. Each write first checks with `versions` that + * the file did not change since a tool read it. After each write, the + * `afterWrite` hooks run for the file. `edit_file` finds the old text with + * {@link findMatches}. It writes the new text with the line endings of the + * file, and keeps a byte order mark. */ -export function editTools(env: ToolEnv) { +export function editTools(env: ToolEnv, versions: FileVersions) { return [ toolDefinition({ name: 'write_file', @@ -140,9 +201,13 @@ export function editTools(env: ToolEnv) { 'file', ) const content = stringArg(args, 'content') + const missing = !(await versions.exists(full)) return env.lock(full, async () => { + await versions.check(full, { missing }) await env.backend.writeFile(full, content) const notes = await afterWrite(env, full) + // After the hooks: a formatter can write the file again. + await versions.saw(full) return [`Wrote ${env.shown(full)}.`, ...notes].join('\n') }) }), @@ -166,6 +231,9 @@ export function editTools(env: ToolEnv) { const newText = stringArg(args, 'new') const replaceAll = isRecord(args) && args.replaceAll === true return env.lock(full, async () => { + // The edit reads the file itself, so a file that no tool read can + // change. One that a tool read must be as that read saw it. + await versions.check(full, { missing: false }) const before = await readText(env.backend, full) const matches = findMatches(before, oldText) const count = matches.length @@ -187,6 +255,7 @@ export function editTools(env: ToolEnv) { const keepBom = before.startsWith(BOM) && !after.startsWith(BOM) await env.backend.writeFile(full, keepBom ? BOM + after : after) const notes = await afterWrite(env, full) + await versions.saw(full) return [ `Edited ${env.shown(full)} (${count} change${count > 1 ? 's' : ''}).`, ...notes, diff --git a/packages/ai-harness/src/first-party/coding/patch.ts b/packages/ai-harness/src/first-party/coding/patch.ts index 8caed0d643..58becd2a52 100644 --- a/packages/ai-harness/src/first-party/coding/patch.ts +++ b/packages/ai-harness/src/first-party/coding/patch.ts @@ -2,6 +2,7 @@ import { toolDefinition } from '@tanstack/ai' import { afterWrite, readText, stringArg } from './backend' import { findMatches, toEndingsOf } from './edit' import type { ToolEnv } from './backend' +import type { FileVersions } from './edit' const BOM = '' @@ -306,9 +307,14 @@ interface Step { /** * Apply each step. Every file is read and checked first. The files are * written and removed only when every operation fits. Then the `afterWrite` - * hooks run for the written files. + * hooks run for the written files, and `versions` remembers the files as the + * patch left them. */ -async function applyAll(env: ToolEnv, steps: ReadonlyArray) { +async function applyAll( + env: ToolEnv, + versions: FileVersions, + steps: ReadonlyArray, +) { const { backend } = env // The new state of each file: its text, or `null` to remove it. A later // operation on a file sees it. A removed file is always a file that was @@ -394,6 +400,9 @@ async function applyAll(env: ToolEnv, steps: ReadonlyArray) { for (const full of removed) await backend.remove?.(full) const notes: Array = [] for (const { full } of written) notes.push(...(await afterWrite(env, full))) + // After the hooks: a formatter can write a file again. + for (const { full } of written) await versions.saw(full) + for (const full of removed) versions.forget(full) return ['Applied the patch:', ...changes, ...notes].join('\n') } @@ -412,8 +421,10 @@ async function applyAll(env: ToolEnv, steps: ReadonlyArray) { * refuses them, before it writes anything. * - A file that the patch removed or moved away cannot change again in the * same patch. A file cannot be removed or moved after the patch changed it. + * - A file that changed since a tool read it fails the patch, before it + * writes anything (see `versions`). */ -export function patchTools(env: ToolEnv) { +export function patchTools(env: ToolEnv, versions: FileVersions) { return [ toolDefinition({ name: 'patch', @@ -443,7 +454,18 @@ export function patchTools(env: ToolEnv) { const paths = [ ...new Set(steps.flatMap((step) => [step.full, step.to])), ].sort() - let run = () => applyAll(env, steps) + // The paths with no file before the locks. A file there under the + // locks is new, and the patch did not see it. + const missing = new Set() + for (const path of paths) { + if (!(await versions.exists(path))) missing.add(path) + } + let run = async () => { + for (const path of paths) { + await versions.check(path, { missing: missing.has(path) }) + } + return applyAll(env, versions, steps) + } for (const path of paths) { const inner = run run = () => env.lock(path, inner) diff --git a/packages/ai-harness/src/first-party/coding/snapshots.ts b/packages/ai-harness/src/first-party/coding/snapshots.ts index 6e1abcf1dd..207e9555f0 100644 --- a/packages/ai-harness/src/first-party/coding/snapshots.ts +++ b/packages/ai-harness/src/first-party/coding/snapshots.ts @@ -2,10 +2,13 @@ import { createHash } from 'node:crypto' import { LogRecordsCapability, getLogRecords } from '@tanstack/ai' import { defineCommand } from '../../commands' import { RevertFiles, RevertStanding, definePlugin } from '../../plugins' -import { hostBackend } from './backend' import { isRecord } from '../../utils' +import { WorkspaceHooks } from '../workspace-hooks' +import { hostBackend, pathsOf } from './backend' +import { parsePatch } from './patch' import { GIT, GIT_ENV, quoteArg, shellPlatform } from './search' -import type { ModelMessage } from '@tanstack/ai' +import { createLock } from './workspace' +import type { ModelMessage, ToolCall } from '@tanstack/ai' import type { WorkspaceBackend } from './backend' export interface SnapshotsOptions { @@ -13,8 +16,8 @@ export interface SnapshotsOptions { root: string /** * A folder for the shadow git repositories, an absolute path. Each - * workspace gets its own repository in it. The project's own `.git` is - * never used. + * thread of each tenant gets its own repository in it. The project's own + * `.git` is never used. */ dataDir: string /** Where the files are and git runs. Default: {@link hostBackend}. */ @@ -33,6 +36,13 @@ export interface SnapshotStep { toolCallIds: Array } +/** Whose shadow repository it is: a thread, in a tenant. */ +export interface SnapshotThread { + /** The tenant. `undefined` in the default tenant. */ + tenantId?: string + threadId: string +} + const isStep = (value: unknown): value is SnapshotStep => isRecord(value) && typeof value.from === 'string' && @@ -55,20 +65,100 @@ const isReverted = ( /** The log record type of a {@link SnapshotStep}. */ const STEP_RECORD = 'tanstack/snapshots:step' const TIMEOUT_MS = 60_000 +/** The most steps that the plugin state keeps for a thread. */ +const MAX_STEPS = 500 +/** + * The waits, in ms, before each new try of a git command that found the + * index locked by another git process. + */ +const LOCKED_RETRY_MS = [50, 100, 200, 400] const BUSY = 'Wait until the turn ends, then try again.' const REVERTED = 'A revert stands. Run unrevert, or send a prompt, first.' +/** The tools that change no file. */ +const READ_TOOLS = new Set([ + 'read_file', + 'list_files', + 'grep', + 'webfetch', + 'websearch', +]) + /** - * A shadow git repository for `options.root`. Git commands of one - * repository run one at a time, so sessions on one workspace do not clash - * on its index. + * Does `call` give each file that it changes to the `afterWrite` hook? The + * read tools change no file. `write_file` and `edit_file` give each write. + * `patch` gives its writes, but not a file that it deletes or moves away. + * Other tools, like `bash`, give nothing. */ -function shadowRepo(options: SnapshotsOptions) { +function reportsChanges(call: ToolCall) { + const { name } = call.function + if (READ_TOOLS.has(name) || name === 'write_file' || name === 'edit_file') { + return true + } + if (name !== 'patch') return false + try { + const args: unknown = JSON.parse(call.function.arguments) + const text = isRecord(args) && typeof args.patch === 'string' ? args.patch : '' + const operations = parsePatch(text) + return operations.every( + (operation) => + operation.type === 'add' || + (operation.type === 'update' && operation.movePath === undefined), + ) + } catch { + // A patch that does not parse changes no file. + return true + } +} + +/** + * The files of `steps`, by the tree that they go back to: for each file, + * the tree at the start of the first step that changed it. `steps` are in + * order. + */ +function firstTrees(steps: ReadonlyArray) { + const treeOf = new Map() + for (const step of steps) { + for (const file of step.files) { + if (!treeOf.has(file)) treeOf.set(file, step.from) + } + } + const byTree = new Map>() + for (const [file, tree] of treeOf) { + const files = byTree.get(tree) ?? new Set() + files.add(file) + byTree.set(tree, files) + } + return byTree +} + +/** + * The git commands of each shadow repository, by its folder, run one at a + * time in this process. Two plugins or sessions on one repository then do + * not clash on its index. + * + * ponytail: another process can still lock the index. Then a command tries + * again a few times, and fails with a clear error. + */ +const repoLocks = createLock() + +/** + * A shadow git repository for the workspace `options.root` and `thread`. It + * never stores `.env` files. + */ +function shadowRepo(options: SnapshotsOptions, thread: SnapshotThread) { const backend = options.backend ?? hostBackend const platform = shellPlatform(backend) const quote = (value: string) => quoteArg(value, platform) - const hash = createHash('sha1').update(options.root).digest('hex') - const gitDir = `${options.dataDir}/${hash}` + // Each tenant, thread, and root gets its own repository, so one tenant + // never sees the files of another. + const key = JSON.stringify([ + thread.tenantId ?? null, + thread.threadId, + options.root, + ]) + const hash = createHash('sha1').update(key).digest('hex') + const gitDir = pathsOf(backend).join(options.dataDir, hash) // The env wins over a GIT_DIR or GIT_INDEX_FILE from a git hook, so git // never writes to the project's repository. const env = { @@ -79,17 +169,33 @@ function shadowRepo(options: SnapshotsOptions) { // `core.autocrlf=false` keeps the line endings of every file. const git = `${GIT} -c core.autocrlf=false --git-dir ${quote(gitDir)} --work-tree ${quote(options.root)}` - /** Run a command in the workspace. Throws with git's message when it fails. */ + /** + * Run a command in the workspace, and give its output. When another git + * process holds the index lock, it tries again a few times. Throws with + * git's message when it fails. + */ const exec = async (command: string) => { - const result = await backend.exec(command, { - cwd: options.root, - env, - timeoutMs: TIMEOUT_MS, - }) - if (result.exitCode !== 0) { - throw new Error(`${command} failed: ${result.stderr.trim()}`) + let tries = 0 + for (;;) { + const result = await backend.exec(command, { + cwd: options.root, + env, + timeoutMs: TIMEOUT_MS, + }) + if (result.exitCode === 0) return result.stdout + const stderr = result.stderr.trim() + if (!stderr.includes('index.lock')) { + throw new Error(`${command} failed: ${stderr}`) + } + const wait = LOCKED_RETRY_MS[tries] + if (wait === undefined) { + throw new Error( + `another git process holds the index lock of ${gitDir}: ${stderr}`, + ) + } + tries += 1 + await new Promise((resolve) => setTimeout(resolve, wait)) } - return result.stdout } const run = (args: string) => exec(`${git} ${args}`) @@ -106,25 +212,12 @@ function shadowRepo(options: SnapshotsOptions) { return names.split('\0').filter((name) => name !== '') } - let queue: Promise = Promise.resolve() - const serial = (fn: () => Promise) => { - const result = queue.then(fn) - queue = result.then( - () => undefined, - () => undefined, - ) - return result - } + const serial = (fn: () => Promise) => repoLocks(gitDir, fn) return { - /** Create the repository. `false` when the backend has no git. */ + /** Create the repository, or keep the one that is there. */ init: () => serial(async () => { - const probe = await backend.exec(`${GIT} --version`, { - env: GIT_ENV, - timeoutMs: TIMEOUT_MS, - }) - if (probe.exitCode !== 0) return false // A bare repository: it has no work tree of its own. await exec(`${GIT} init --bare -q ${quote(gitDir)}`) // Keep the bytes of every file: no line ending change, no filter. @@ -132,7 +225,18 @@ function shadowRepo(options: SnapshotsOptions) { `${gitDir}/info/attributes`, '* -text -ident -filter\n', ) - return true + // `.env` files hold secrets. The snapshots never store them. + await backend.writeFile(`${gitDir}/info/exclude`, '.env*\n') + }), + /** Is `id` a tree of this repository? */ + has: (id: string) => + serial(async () => { + const result = await backend.exec(`${git} cat-file -t ${quote(id)}`, { + cwd: options.root, + env, + timeoutMs: TIMEOUT_MS, + }) + return result.exitCode === 0 && result.stdout.trim() === 'tree' }), track: () => serial(track), changed: (from: string, to: string) => serial(() => changed(from, to)), @@ -147,10 +251,10 @@ function shadowRepo(options: SnapshotsOptions) { /** * Make the files that differ between `from` and `tree` match `tree`, and * give their count. Of these files, the ones that are not in `tree` are - * removed. With `only`, only the files in it change. All other files stay - * as they are. + * removed. Only the files in `only` change. All other files stay as they + * are. */ - restore: (tree: string, from: string, only?: ReadonlySet) => + restore: (tree: string, from: string, only: ReadonlySet) => serial(async () => { // Pairs of a status letter and a path. `D`: the path is not in `tree`. const status = ( @@ -162,7 +266,7 @@ function shadowRepo(options: SnapshotsOptions) { const restored: Array = [] for (let i = 1; i < status.length; i += 2) { const path = status[i] - if (!path || (only && !only.has(path))) continue + if (!path || !only.has(path)) continue if (status[i - 1] === 'D') removed.push(path) else restored.push(path) } @@ -191,15 +295,50 @@ function shadowRepo(options: SnapshotsOptions) { if (commands.length > 0) await exec(commands.join(' && ')) return removed.length + restored.length }), + /** Remove the repository, after the git commands that wait for it. */ + purge: () => + serial(async () => { + const remove = + platform === 'win32' + ? `rmdir /s /q ${quote(gitDir)}` + : `rm -rf ${quote(gitDir)}` + // `rmdir` fails for a folder that is not there, so only a folder + // that is still there is an error. + await backend.exec(remove, { timeoutMs: TIMEOUT_MS }) + if (await backend.stat(gitDir)) { + throw new Error(`snapshots: could not remove ${gitDir}.`) + } + }), } } +/** The shadow repository of one thread, from {@link shadowRepo}. */ +type ShadowRepo = ReturnType + const messageOf = (error: unknown) => error instanceof Error ? error.message : String(error) +/** + * Remove the shadow repository of `thread`, that `snapshots(options)` made. + * It waits for the git commands of that repository that run in this + * process. It resolves also when there is no repository. + * + * @example + * ```ts + * await purgeSnapshots({ root, dataDir }, { tenantId: 'acme', threadId: 't-1' }) + * ``` + */ +export function purgeSnapshots( + options: SnapshotsOptions, + thread: SnapshotThread, +) { + return shadowRepo(options, thread).purge() +} + /** * Snapshots of the workspace files at the start and the end of each model - * step, in a shadow git repository in `dataDir`. Adds two commands: + * step, in a shadow git repository in `dataDir`: one repository for each + * thread of each tenant. Adds two commands: * * - `/undo`: put the files that the last turn changed back as they were * before it, and remove that turn from the transcript. Other files stay @@ -209,10 +348,17 @@ const messageOf = (error: unknown) => * `session.revert(messageId)` also puts back the files that the tool calls * after that message changed, and `session.unrevert()` brings them back. * + * A step keeps only the files that this session changed. The workspace + * tools give each file that they write, so a change of another session or + * agent in the same time is not in the step. A step with a tool that does + * not give its changes, like `bash`, keeps every file that changed in it. + * The snapshots never store `.env` files. + * * The changed files of each step go to the session log as a * `tanstack/snapshots:step` record on a durable host, else to the plugin - * state. `diff(from, to?)` gives the changed files and a unified diff - * between two tree ids. Without `to`, it compares with the files now. + * state, which keeps the last 500 steps. `diff(from, to?)` gives the changed + * files and a unified diff between two tree ids of an open session. Without + * `to`, it compares with the files now. * * Needs the git CLI on the backend. Without git, the plugin does nothing. * @@ -226,31 +372,68 @@ const messageOf = (error: unknown) => * ``` */ export function snapshots(options: SnapshotsOptions) { - const repo = shadowRepo(options) - let ready: Promise | undefined - const isReady = () => - (ready ??= repo.init().then( - (hasGit) => { - if (!hasGit) - console.warn('snapshots: git was not found. Snapshots are off.') - return hasGit - }, - (error: unknown) => { - console.warn(`snapshots: ${messageOf(error)}. Snapshots are off.`) - return false - }, - )) + const backend = options.backend ?? hostBackend + const paths = pathsOf(backend) + // The repositories of the open sessions, for `diff`. + const repos = new Set() + let found: Promise | undefined + /** Is git on the backend? It asks once, and warns once when it is not. */ + const hasGit = () => + (found ??= backend + .exec(`${GIT} --version`, { env: GIT_ENV, timeoutMs: TIMEOUT_MS }) + .then( + (probe) => { + const isFound = probe.exitCode === 0 + if (!isFound) { + console.warn('snapshots: git was not found. Snapshots are off.') + } + return isFound + }, + (error: unknown) => { + console.warn(`snapshots: ${messageOf(error)}. Snapshots are off.`) + return false + }, + )) + + /** `full` as a path in the workspace with `/`, or `undefined` outside it. */ + const inRoot = (full: string) => { + const path = paths.relative(options.root, full) + const isOutside = + path === '' || + path === '..' || + path.startsWith(`..${paths.sep}`) || + paths.isAbsolute(path) + return isOutside ? undefined : path.split(paths.sep).join('/') + } return definePlugin({ name: 'tanstack/snapshots', diff: async (from: string, to?: string) => { - if (!(await isReady())) - throw new Error('snapshots: git is not available.') - return repo.diff(from, to) + if (!(await hasGit())) throw new Error('snapshots: git is not available.') + // ponytail: each thread has its own repository, so `diff` asks the + // repositories of the open sessions, one by one, which has `from`. + const open = [...repos] + for (const repo of open) { + if (await repo.has(from)) return repo.diff(from, to) + } + throw new Error(`snapshots: no open session has the snapshot ${from}.`) }, provides: [RevertFiles], setup: async (ctx) => { - if (!(await isReady())) { + const repo = shadowRepo(options, { + tenantId: ctx.session.principal?.tenantId, + threadId: ctx.session.threadId, + }) + const isReady = + (await hasGit()) && + (await repo.init().then( + () => true, + (error: unknown) => { + console.warn(`snapshots: ${messageOf(error)}. Snapshots are off.`) + return false + }, + )) + if (!isReady) { // Without git, a revert changes only the transcript. ctx.provide(RevertFiles, { revert: async () => undefined, @@ -258,6 +441,10 @@ export function snapshots(options: SnapshotsOptions) { }) return } + await ctx.resources.acquire( + () => repos.add(repo), + () => repos.delete(repo), + ) const reverted = ctx.getOptional(RevertStanding) const noSteps: Array = [] const steps = ctx.state(noSteps) @@ -273,13 +460,16 @@ export function snapshots(options: SnapshotsOptions) { const later = all .filter(isStep) .filter((step) => step.toolCallIds.some((id) => calls.has(id))) - const first = later[0] - if (!first) return undefined - // Only the files that the steps changed go back. The edits of the - // user in other files stay. `now` keeps all files for `unrevert`. + if (later.length === 0) return undefined + // Only the files that the steps changed go back, each to its state + // before the first of them. The edits of the user in other files + // stay. `now` keeps all files for `unrevert`. const files = [...new Set(later.flatMap((step) => step.files))] const now = await repo.track() - await repo.restore(first.from, now, new Set(files)) + const groups = firstTrees(later) + for (const [tree, group] of groups) { + await repo.restore(tree, now, group) + } return { now, files } }, unrevert: async (saved) => { @@ -293,12 +483,24 @@ export function snapshots(options: SnapshotsOptions) { }) /** The tree at the start of the step that runs. */ let stepStart: string | undefined - /** The tree at the start of the last turn: what `/undo` restores. */ - let turnStart: string | undefined - /** What the last `/undo` removed: the trees of the turn, and its messages. */ + /** + * The steps of the last turn, with changed files: what `/undo` puts + * back. `undefined` when there is nothing to undo. + */ + let turnSteps: Array | undefined + /** What the last `/undo` removed: the steps, the end, and the messages. */ let redo: - | { start: string; end: string; messages: Array } + | { + end: string + steps: Array + messages: Array + } | undefined + /** + * The files in the workspace that the workspace tools of this session + * wrote in the step that runs. A session runs one turn at a time. + */ + const written = new Set() // A failed snapshot must not fail the turn. It only says so. const guard = async (fn: () => Promise) => { @@ -310,7 +512,32 @@ export function snapshots(options: SnapshotsOptions) { } } + /** + * Put each file of `undone` back to `target`, the tree that `/undo` + * or `/redo` goes to, and give the count of the files that changed. + */ + const restoreTurn = async ( + undone: ReadonlyArray, + target: (start: string) => { tree: string; from: string }, + ) => { + let count = 0 + const groups = firstTrees(undone) + for (const [start, files] of groups) { + const { tree, from } = target(start) + count += await repo.restore(tree, from, files) + } + return count + } + return { + contribute: [ + WorkspaceHooks.item({ + afterWrite: async (full) => { + const path = inRoot(full) + if (path !== undefined) written.add(path) + }, + }), + ], middleware: [ { name: 'tanstack/snapshots', @@ -321,24 +548,41 @@ export function snapshots(options: SnapshotsOptions) { if (iteration > 0 && stepStart !== undefined) return stepStart = await guard(repo.track) if (iteration > 0) return - turnStart = stepStart + written.clear() + turnSteps = stepStart === undefined ? undefined : [] redo = undefined }, onToolPhaseComplete: async (ctx, { toolCalls }) => { const from = stepStart + const mine = new Set(written) + written.clear() if (from === undefined) return + // ponytail: a step with a tool like `bash` keeps every file + // that changed in it, also a change of another agent then. + const keepsAll = !toolCalls.every(reportsChanges) const step = await guard(async () => { const to = await repo.track() - const files = to === from ? [] : await repo.changed(from, to) + const hasNoChange = to === from || (!keepsAll && mine.size === 0) + const changed = hasNoChange ? [] : await repo.changed(from, to) + // Git skips some files, like `.env`, so the step keeps the + // files that git saw change and the tools wrote. + const files = keepsAll + ? changed + : changed.filter((file) => mine.has(file)) const toolCallIds = toolCalls.map((call) => call.id) return { from, to, files, toolCallIds } }) // A failed snapshot leaves no start: the next step takes one. stepStart = step?.to if (!step || step.files.length === 0) return + turnSteps?.push(step) const log = getLogRecords(ctx, { optional: true }) if (log) await log.append([{ type: STEP_RECORD, ...step }]) - else await steps.update((saved) => [...saved, step]) + else { + await steps.update((saved) => + [...saved, step].slice(-MAX_STEPS), + ) + } }, }, ], @@ -361,21 +605,26 @@ export function snapshots(options: SnapshotsOptions) { // end of the last turn. const turnEnd = stepStart if ( - turnStart === undefined || + turnSteps === undefined || turnEnd === undefined || start === -1 ) { return 'Nothing to undo.' } - // Only the files that the turn changed go back. A later edit of - // the user in one of them is lost. Other files stay as they are. - const count = await repo.restore(turnStart, turnEnd) + // Only the files that the steps of the turn changed go back, + // each to its state before the first step that changed it. A + // later edit of the user in one of them is lost. Other files + // stay as they are. + const count = await restoreTurn(turnSteps, (tree) => ({ + tree, + from: turnEnd, + })) redo = { - start: turnStart, end: turnEnd, + steps: turnSteps, messages: messages.slice(start), } - turnStart = undefined + turnSteps = undefined await session.replaceTranscript(messages.slice(0, start)) return `Undid the last turn. Files restored: ${count}.` }, @@ -386,10 +635,10 @@ export function snapshots(options: SnapshotsOptions) { if (session.snapshot().status !== 'idle') return BUSY if (reverted?.()) return REVERTED if (!redo) return 'Nothing to redo.' - const { start, end, messages } = redo - await repo.restore(end, start) + const { end, steps: undone, messages } = redo + await restoreTurn(undone, (from) => ({ tree: end, from })) // The same turn can be undone again. `stepStart` is still `end`. - turnStart = start + turnSteps = undone redo = undefined const transcript = await session.transcript() await session.replaceTranscript([...transcript, ...messages]) diff --git a/packages/ai-harness/src/first-party/coding/workspace.ts b/packages/ai-harness/src/first-party/coding/workspace.ts index d70acc7f06..90c281149f 100644 --- a/packages/ai-harness/src/first-party/coding/workspace.ts +++ b/packages/ai-harness/src/first-party/coding/workspace.ts @@ -9,7 +9,7 @@ import { import { WorkspaceHooks } from '../workspace-hooks' import { hostBackend, optionalString, pathsOf, stringArg } from './backend' import { bashResources, bashTools } from './bash' -import { editTools } from './edit' +import { editTools, fileVersions } from './edit' import { patchPaths, patchTools } from './patch' import { readTools } from './read' import { searchTools } from './search' @@ -122,8 +122,8 @@ function within(dir: string, full: string, sep: string) { return full === dir || full.startsWith(dir.endsWith(sep) ? dir : dir + sep) } -/** A lock per path: work on one path runs one at a time, in order. */ -function createLock() { +/** A lock per key: work on one key runs one at a time, in order. */ +export function createLock() { const queues = new Map>() return (path: string, fn: () => Promise) => { const result = (queues.get(path) ?? Promise.resolve()).then(fn) @@ -139,6 +139,17 @@ function createLock() { } } +/** + * The file locks of this process, by the workspace root and the path. Every + * session and agent on one workspace shares them, so two edits of one file + * run one after the other. + * + * ponytail: a lock is for this process only. Two processes or machines on + * one shared sandbox do not wait for each other. The unchanged check of the + * edit tools then fails the second write, so no edit is lost. + */ +const fileLocks = createLock() + /** * Does `editStyle: 'auto'` give `model` the `patch` tool? GPT models do: * `gpt-` ids and the `o` series, like `o3`. @@ -352,15 +363,14 @@ function bindTools( return inRoot } - const lockPath = createLock() const env: ToolEnv = { backend, root, hooks: session.hooks ?? (() => []), - // ponytail: one lock per workspace binding. Sessions on one folder do - // not wait for each other. + // The key has the root, so the fixed lock order of `patch` (by path) is + // the same for every session on this root. lock: (path, fn) => - lockPath(path, async () => { + fileLocks(`${root}\0${path}`, async () => { // `reach` checked `path` before the lock. A link made since then // would move the work to another file, so check it again here. const now = await resolved(path) @@ -386,10 +396,23 @@ function bindTools( ]) !== 'allow' }, } + // What this binding read and wrote, so a write fails when the file changed + // since then. + const versions = fileVersions(env) + // `read_file` reads under the lock. Its lock also remembers the file, as + // the read sees it. + const readEnv: ToolEnv = { + ...env, + lock: (path, fn) => + env.lock(path, async () => { + await versions.saw(path) + return fn() + }), + } const tools = [ - ...readTools(env), - ...editTools(env), - ...patchTools(env), + ...readTools(readEnv), + ...editTools(env, versions), + ...patchTools(env, versions), ...searchTools(env), ...bashTools(env, { timeoutMs: options.bashTimeoutMs, diff --git a/packages/ai-harness/tests/coding-locks.test.ts b/packages/ai-harness/tests/coding-locks.test.ts new file mode 100644 index 0000000000..695bac9aa4 --- /dev/null +++ b/packages/ai-harness/tests/coding-locks.test.ts @@ -0,0 +1,204 @@ +import { mkdtemp, 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 { hostBackend } from '../src/first-party/coding/backend' +import { createWorkspaceTools } from '../src/first-party/coding/workspace' +import { gate } from './helpers' +import type { WorkspaceBackend } from '../src/first-party/coding/backend' + +let root = '' + +beforeEach(async () => { + root = await mkdtemp(join(tmpdir(), 'harness-locks-')) +}) + +afterEach(async () => { + await rm(root, { recursive: true, force: true }) +}) + +const put = (path: string, text: string) => writeFile(join(root, path), text) +const read = (path: string) => readFile(join(root, path), 'utf8') + +/** The workspace tools of one session on `root`. Runs a tool by its name. */ +function openSession(backend: WorkspaceBackend = hostBackend) { + const { tools } = createWorkspaceTools({ root, backend, web: false }) + return async (name: string, args: object) => { + const tool = tools.find((each) => each.name === name) + if (!tool?.execute) throw new Error(`No tool ${name}.`) + const result: unknown = await tool.execute(args) + return result + } +} + +/** + * The host backend, with a count of the `realpath` and `stat` calls. With + * `held`, each write waits for `release`, and `reached` resolves at the + * first write. + */ +function watchedBackend({ held = false } = {}) { + const reached = gate() + const released = gate() + const calls = { realpath: 0, stat: 0 } + const backend: WorkspaceBackend = { + ...hostBackend, + realpath: (path) => { + calls.realpath += 1 + return hostBackend.realpath(path) + }, + stat: (path) => { + calls.stat += 1 + return hostBackend.stat(path) + }, + writeFile: async (path, data) => { + if (held) { + reached.open() + await released.opened + } + return hostBackend.writeFile(path, data) + }, + } + return { backend, calls, reached: reached.opened, release: released.open } +} + +/** Long enough for a call that does not wait at the lock to go past it. */ +const settle = () => new Promise((resolve) => setTimeout(resolve, 50)) + +/** A promise, and whether it settled. */ +function tracked(promise: Promise) { + const state = { settled: false } + const done = promise.finally(() => { + state.settled = true + }) + return { done, state } +} + +describe('file locks across sessions', () => { + it('makes the second edit wait, then fails it: the file changed since its read', async () => { + await put('a.txt', 'one\n') + const first = watchedBackend({ held: true }) + const second = watchedBackend() + const firstCall = openSession(first.backend) + const secondCall = openSession(second.backend) + await firstCall('read_file', { path: 'a.txt' }) + await secondCall('read_file', { path: 'a.txt' }) + + const firstEdit = firstCall('edit_file', { + path: 'a.txt', + old: 'one', + new: 'first', + }) + await first.reached + const before = second.calls.realpath + const secondEdit = tracked( + secondCall('edit_file', { path: 'a.txt', old: 'one', new: 'second' }), + ) + // `reach` resolves the path once. The lock resolves it again only when + // the first edit is done. + await vi.waitFor(() => expect(second.calls.realpath).toBe(before + 1)) + await settle() + expect(second.calls.realpath).toBe(before + 1) + expect(secondEdit.state.settled).toBe(false) + + first.release() + await firstEdit + await expect(secondEdit.done).rejects.toThrow( + 'a.txt changed since it was read. Read it again before you edit it.', + ) + expect(await read('a.txt')).toBe('first\n') + }) + + it('lets the second session edit after it reads the file again', async () => { + await put('a.txt', 'one two\n') + const firstCall = openSession() + const secondCall = openSession() + await firstCall('read_file', { path: 'a.txt' }) + await secondCall('read_file', { path: 'a.txt' }) + await firstCall('edit_file', { path: 'a.txt', old: 'one', new: 'ONE' }) + await expect( + secondCall('edit_file', { path: 'a.txt', old: 'two', new: 'TWO' }), + ).rejects.toThrow('a.txt changed since it was read.') + + await secondCall('read_file', { path: 'a.txt' }) + await secondCall('edit_file', { path: 'a.txt', old: 'two', new: 'TWO' }) + + expect(await read('a.txt')).toBe('ONE TWO\n') + }) + + it('fails write_file when another session created the file while it waited', async () => { + const first = watchedBackend({ held: true }) + const second = watchedBackend() + const firstWrite = openSession(first.backend)('write_file', { + path: 'new.txt', + content: 'first', + }) + await first.reached + const secondWrite = tracked( + openSession(second.backend)('write_file', { + path: 'new.txt', + content: 'second', + }), + ) + // One `stat` before the lock sees no file. The check under the lock + // waits for the first write. + await vi.waitFor(() => expect(second.calls.stat).toBe(1)) + await settle() + expect(secondWrite.state.settled).toBe(false) + + first.release() + await firstWrite + await expect(secondWrite.done).rejects.toThrow( + 'new.txt was created while the tool waited. Read it before you edit it.', + ) + expect(await read('new.txt')).toBe('first') + }) +}) + +describe('the unchanged check', () => { + it.each([ + ['edit_file', { path: 'a.txt', old: 'one', new: 'two' }], + ['write_file', { path: 'a.txt', content: 'two\n' }], + [ + 'patch', + { + patch: [ + '*** Begin Patch', + '*** Update File: a.txt', + '@@', + '-one', + '+two', + '*** End Patch', + ].join('\n'), + }, + ], + ])('%s fails for a file that changed on disk after the read', async (name, args) => { + await put('a.txt', 'one\n') + const call = openSession() + await call('read_file', { path: 'a.txt' }) + await put('a.txt', 'changed on disk\n') + + await expect(call(name, args)).rejects.toThrow( + 'a.txt changed since it was read. Read it again before you edit it.', + ) + expect(await read('a.txt')).toBe('changed on disk\n') + }) + + it("keeps the session's own writes: a second edit after the first works", async () => { + await put('a.txt', 'one two\n') + const call = openSession() + await call('read_file', { path: 'a.txt' }) + + await call('edit_file', { path: 'a.txt', old: 'one', new: 'ONE' }) + await call('edit_file', { path: 'a.txt', old: 'two', new: 'TWO' }) + + expect(await read('a.txt')).toBe('ONE TWO\n') + }) + + it('lets write_file replace a file that the session never read', async () => { + await put('a.txt', 'one\n') + + await openSession()('write_file', { path: 'a.txt', content: 'two\n' }) + + expect(await read('a.txt')).toBe('two\n') + }) +}) diff --git a/packages/ai-harness/tests/coding-patch.test.ts b/packages/ai-harness/tests/coding-patch.test.ts index 6b5ee0c0aa..d2b97a2585 100644 --- a/packages/ai-harness/tests/coding-patch.test.ts +++ b/packages/ai-harness/tests/coding-patch.test.ts @@ -1,6 +1,7 @@ import { tmpdir } from 'node:os' import { join, relative, resolve, sep } from 'node:path' import { describe, expect, it } from 'vitest' +import { fileVersions } from '../src/first-party/coding/edit' import { parsePatch, patchPaths, @@ -79,7 +80,7 @@ function workspace( reach: async (path) => resolve(root, path), shown, } - const tools: Array = patchTools(env) + const tools: Array = patchTools(env, fileVersions(env)) const [tool] = tools return { written, diff --git a/packages/ai-harness/tests/coding-snapshots.test.ts b/packages/ai-harness/tests/coding-snapshots.test.ts index eca059c719..2ac4a7ec0c 100644 --- a/packages/ai-harness/tests/coding-snapshots.test.ts +++ b/packages/ai-harness/tests/coding-snapshots.test.ts @@ -14,7 +14,11 @@ import { memoryLogStore, memoryPersistence } from '@tanstack/ai-persistence' import { createHarnessHost, defineHarness, logIdOf } from '../src' import { hostBackend } from '../src/first-party/coding/backend' import { quoteArg } from '../src/first-party/coding/search' -import { snapshots } from '../src/first-party/coding/snapshots' +import { + purgeSnapshots, + snapshots, +} from '../src/first-party/coding/snapshots' +import { workspaceTools } from '../src/first-party/coding/workspace' import { after, gate, @@ -23,7 +27,7 @@ import { text, toolCall, } from './helpers' -import type { HarnessPersistence, HarnessSession } from '../src' +import type { HarnessPersistence, HarnessSession, Principal } from '../src' import type { WorkspaceBackend } from '../src/first-party/coding/backend' import type { Reply } from './helpers' @@ -66,13 +70,18 @@ const editTurn = (answer: string): Array => [ /** * Open a session with `snapshots()` and a `change` tool. Each call of the - * tool runs the next function of `changes`. + * tool runs the next function of `changes`. With `workspace`, the session + * also has `workspaceTools()` on `root`. The thread is `THREAD`, unless + * `threadId` names another. */ async function open(options: { replies: Array changes?: Array<() => Promise> persistence?: HarnessPersistence backend?: WorkspaceBackend + workspace?: boolean + threadId?: string + principal?: Principal }) { const changes = options.changes ?? [] const change = toolDefinition({ @@ -92,10 +101,14 @@ async function open(options: { name: 'test/snapshots', adapter, tools: [change], - plugins: () => [plugin], + plugins: () => + options.workspace + ? [workspaceTools({ root, web: false }), plugin] + : [plugin], }), { - threadId: THREAD, + threadId: options.threadId ?? THREAD, + principal: options.principal, }, ) return { host, session, plugin, calls } From 76af6ac06144f58743bb60983e1027d5e8854808 Mon Sep 17 00:00:00 2001 From: Alem Tuzlak Date: Sat, 17 Oct 2026 17:47:22 +0200 Subject: [PATCH 03/17] fix(ai-harness): shared file locks with an unchanged check, and snapshots per thread and per run --- .../ai-harness/src/first-party/coding/edit.ts | 3 +- .../src/first-party/coding/index.ts | 8 +- .../src/first-party/coding/snapshots.ts | 139 ++++++++---- .../src/first-party/coding/workspace.ts | 42 +++- .../ai-harness/tests/coding-locks.test.ts | 25 ++- .../ai-harness/tests/coding-snapshots.test.ts | 212 +++++++++++++++++- 6 files changed, 355 insertions(+), 74 deletions(-) diff --git a/packages/ai-harness/src/first-party/coding/edit.ts b/packages/ai-harness/src/first-party/coding/edit.ts index c73b196d19..cb17ad02c7 100644 --- a/packages/ai-harness/src/first-party/coding/edit.ts +++ b/packages/ai-harness/src/first-party/coding/edit.ts @@ -143,7 +143,8 @@ export function fileVersions(env: ToolEnv) { seen.delete(path) }, /** Is something at `path`? Call it before the lock. */ - exists: async (path: string) => (await env.backend.stat(path)) !== undefined, + exists: async (path: string) => + (await env.backend.stat(path)) !== undefined, /** * Throws when the file at `path` changed since a tool read or wrote it. * For a file that no tool read, it throws only when the tool saw no file diff --git a/packages/ai-harness/src/first-party/coding/index.ts b/packages/ai-harness/src/first-party/coding/index.ts index b3bc11ecd8..faae8310a7 100644 --- a/packages/ai-harness/src/first-party/coding/index.ts +++ b/packages/ai-harness/src/first-party/coding/index.ts @@ -6,7 +6,11 @@ export { hostBackend } from './backend' export type { BoundWorkspace, WorkspaceBackend } from './backend' export { WorkspaceHooks } from '../workspace-hooks' export type { SearchProvider, WebToolsOptions } from './web' -export { snapshots } from './snapshots' -export type { SnapshotStep, SnapshotsOptions } from './snapshots' +export { purgeSnapshots, snapshots } from './snapshots' +export type { + SnapshotStep, + SnapshotThread, + SnapshotsOptions, +} from './snapshots' export { FormatFailed, formatter } from './formatter' export type { Formatter, FormatterOptions, FormatterProject } from './formatter' diff --git a/packages/ai-harness/src/first-party/coding/snapshots.ts b/packages/ai-harness/src/first-party/coding/snapshots.ts index 207e9555f0..1491a0e1a4 100644 --- a/packages/ai-harness/src/first-party/coding/snapshots.ts +++ b/packages/ai-harness/src/first-party/coding/snapshots.ts @@ -7,7 +7,7 @@ import { WorkspaceHooks } from '../workspace-hooks' import { hostBackend, pathsOf } from './backend' import { parsePatch } from './patch' import { GIT, GIT_ENV, quoteArg, shellPlatform } from './search' -import { createLock } from './workspace' +import { createLock, runOf, writingRun } from './workspace' import type { ModelMessage, ToolCall } from '@tanstack/ai' import type { WorkspaceBackend } from './backend' @@ -98,7 +98,8 @@ function reportsChanges(call: ToolCall) { if (name !== 'patch') return false try { const args: unknown = JSON.parse(call.function.arguments) - const text = isRecord(args) && typeof args.patch === 'string' ? args.patch : '' + const text = + isRecord(args) && typeof args.patch === 'string' ? args.patch : '' const operations = parsePatch(text) return operations.every( (operation) => @@ -315,13 +316,37 @@ function shadowRepo(options: SnapshotsOptions, thread: SnapshotThread) { /** The shadow repository of one thread, from {@link shadowRepo}. */ type ShadowRepo = ReturnType +/** What the snapshots know of one run of a session. */ +interface RunSteps { + /** The tree at the start of the step that runs. */ + start: string | undefined + /** + * The files in the workspace that the workspace tools of this run wrote + * in the step that runs. + */ + written: Set + /** The steps of this run with changed files. */ + steps: Array +} + +/** + * The id of the run of a middleware call. `''` outside a harness run: all + * such calls count as one run. + */ +const runIdOf = (ctx: unknown) => runOf(ctx)?.runId ?? '' + const messageOf = (error: unknown) => error instanceof Error ? error.message : String(error) /** * Remove the shadow repository of `thread`, that `snapshots(options)` made. - * It waits for the git commands of that repository that run in this - * process. It resolves also when there is no repository. + * Call it when the thread is deleted, after its sessions close. It waits + * for the git commands of that repository that run in this process. It + * resolves also when there is no repository, and throws when the folder is + * still there after the remove. + * + * - `options`: the options of the `snapshots()` plugin. + * - `thread`: the tenant and the thread of the repository. * * @example * ```ts @@ -348,11 +373,12 @@ export function purgeSnapshots( * `session.revert(messageId)` also puts back the files that the tool calls * after that message changed, and `session.unrevert()` brings them back. * - * A step keeps only the files that this session changed. The workspace - * tools give each file that they write, so a change of another session or - * agent in the same time is not in the step. A step with a tool that does - * not give its changes, like `bash`, keeps every file that changed in it. - * The snapshots never store `.env` files. + * A step keeps only the files that its run changed. The workspace tools + * give each file that they write, and the run that wrote it. So a change of + * another run, session, or agent in the same time is not in the step, and + * `/undo` keeps it. A step with a tool that does not give its changes, like + * `bash`, keeps every file that changed in it. The snapshots never store + * `.env*` files. * * The changed files of each step go to the session log as a * `tanstack/snapshots:step` record on a durable host, else to the plugin @@ -481,13 +507,16 @@ export function snapshots(options: SnapshotsOptions) { ) }, }) - /** The tree at the start of the step that runs. */ - let stepStart: string | undefined /** - * The steps of the last turn, with changed files: what `/undo` puts - * back. `undefined` when there is nothing to undo. + * The runs of this session that run now, by run id. A turn and the + * agents that it starts in the background can run at the same time. */ - let turnSteps: Array | undefined + const active = new Map() + /** + * The run that started last: what `/undo` puts back. `undefined` when + * there is nothing to undo. + */ + let last: RunSteps | undefined /** What the last `/undo` removed: the steps, the end, and the messages. */ let redo: | { @@ -496,11 +525,6 @@ export function snapshots(options: SnapshotsOptions) { messages: Array } | undefined - /** - * The files in the workspace that the workspace tools of this session - * wrote in the step that runs. A session runs one turn at a time. - */ - const written = new Set() // A failed snapshot must not fail the turn. It only says so. const guard = async (fn: () => Promise) => { @@ -532,9 +556,12 @@ export function snapshots(options: SnapshotsOptions) { return { contribute: [ WorkspaceHooks.item({ + // The write goes to the run of the tool call that wrote it. afterWrite: async (full) => { + const run = writingRun() const path = inRoot(full) - if (path !== undefined) written.add(path) + if (!run || path === undefined) return + active.get(run.runId)?.written.add(path) }, }), ], @@ -542,27 +569,38 @@ export function snapshots(options: SnapshotsOptions) { { name: 'tanstack/snapshots', optionalRequires: [LogRecordsCapability], - onIteration: async (_ctx, { iteration }) => { + onIteration: async (ctx, { iteration }) => { + const id = runIdOf(ctx) + const known = active.get(id) // A later step starts with the files the step before it left, // so it reuses that snapshot. The model call changes no file. - if (iteration > 0 && stepStart !== undefined) return - stepStart = await guard(repo.track) - if (iteration > 0) return - written.clear() - turnSteps = stepStart === undefined ? undefined : [] + if (iteration > 0 && known) { + known.start ??= await guard(repo.track) + return + } + const run: RunSteps = { + start: await guard(repo.track), + written: new Set(), + steps: [], + } + active.set(id, run) + last = run.start === undefined ? undefined : run redo = undefined }, onToolPhaseComplete: async (ctx, { toolCalls }) => { - const from = stepStart - const mine = new Set(written) - written.clear() + const run = active.get(runIdOf(ctx)) + if (!run) return + const from = run.start + const mine = new Set(run.written) + run.written.clear() if (from === undefined) return // ponytail: a step with a tool like `bash` keeps every file // that changed in it, also a change of another agent then. const keepsAll = !toolCalls.every(reportsChanges) const step = await guard(async () => { const to = await repo.track() - const hasNoChange = to === from || (!keepsAll && mine.size === 0) + const hasNoChange = + to === from || (!keepsAll && mine.size === 0) const changed = hasNoChange ? [] : await repo.changed(from, to) // Git skips some files, like `.env`, so the step keeps the // files that git saw change and the tools wrote. @@ -573,9 +611,9 @@ export function snapshots(options: SnapshotsOptions) { return { from, to, files, toolCallIds } }) // A failed snapshot leaves no start: the next step takes one. - stepStart = step?.to + run.start = step?.to if (!step || step.files.length === 0) return - turnSteps?.push(step) + run.steps.push(step) const log = getLogRecords(ctx, { optional: true }) if (log) await log.append([{ type: STEP_RECORD, ...step }]) else { @@ -584,6 +622,16 @@ export function snapshots(options: SnapshotsOptions) { ) } }, + // `last` keeps the run for `/undo` after it ends. + onFinish: (ctx) => { + active.delete(runIdOf(ctx)) + }, + onAbort: (ctx) => { + active.delete(runIdOf(ctx)) + }, + onError: (ctx) => { + active.delete(runIdOf(ctx)) + }, }, ], commands: { @@ -601,30 +649,25 @@ export function snapshots(options: SnapshotsOptions) { const start = messages.findLastIndex( (message) => message.role === 'user', ) - // When the session is idle, the step start is the tree at the - // end of the last turn. - const turnEnd = stepStart - if ( - turnSteps === undefined || - turnEnd === undefined || - start === -1 - ) { + // When the run is done, its step start is the tree at its end. + const turnEnd = last?.start + if (!last || turnEnd === undefined || start === -1) { return 'Nothing to undo.' } - // Only the files that the steps of the turn changed go back, + // Only the files that the steps of the run changed go back, // each to its state before the first step that changed it. A - // later edit of the user in one of them is lost. Other files - // stay as they are. - const count = await restoreTurn(turnSteps, (tree) => ({ + // later edit in one of them is lost. Other files, also the + // files of other runs and sessions, stay as they are. + const count = await restoreTurn(last.steps, (tree) => ({ tree, from: turnEnd, })) redo = { end: turnEnd, - steps: turnSteps, + steps: last.steps, messages: messages.slice(start), } - turnSteps = undefined + last = undefined await session.replaceTranscript(messages.slice(0, start)) return `Undid the last turn. Files restored: ${count}.` }, @@ -637,8 +680,8 @@ export function snapshots(options: SnapshotsOptions) { if (!redo) return 'Nothing to redo.' const { end, steps: undone, messages } = redo await restoreTurn(undone, (from) => ({ tree: end, from })) - // The same turn can be undone again. `stepStart` is still `end`. - turnSteps = undone + // The same turn can be undone again. + last = { start: end, written: new Set(), steps: undone } redo = undefined const transcript = await session.transcript() await session.replaceTranscript([...transcript, ...messages]) diff --git a/packages/ai-harness/src/first-party/coding/workspace.ts b/packages/ai-harness/src/first-party/coding/workspace.ts index 90c281149f..75cd886234 100644 --- a/packages/ai-harness/src/first-party/coding/workspace.ts +++ b/packages/ai-harness/src/first-party/coding/workspace.ts @@ -1,3 +1,4 @@ +import { AsyncLocalStorage } from 'node:async_hooks' import { BackgroundJobs, SessionSignal, definePlugin } from '../../plugins' import { isRecord } from '../../utils' import { @@ -140,9 +141,9 @@ export function createLock() { } /** - * The file locks of this process, by the workspace root and the path. Every - * session and agent on one workspace shares them, so two edits of one file - * run one after the other. + * The file locks of this process, by the bound workspace root and the real + * path of the file. Every session and agent on one workspace shares them, so + * two edits of one file run one after the other. * * ponytail: a lock is for this process only. Two processes or machines on * one shared sandbox do not wait for each other. The unchanged check of the @@ -172,13 +173,27 @@ function isRun(value: unknown): value is RunContext { ) } -/** The run of a tool call: `ctx.context.run`, or `undefined` without one. */ -function runOf(context: unknown) { +/** + * The run of a tool call or a middleware call: `ctx.context.run`, or + * `undefined` without one. + */ +export function runOf(context: unknown) { if (!isRecord(context) || !isRecord(context.context)) return undefined const { run } = context.context return isRun(run) ? run : undefined } +/** The run of the workspace tool call that runs now. */ +const callRuns = new AsyncLocalStorage() + +/** + * The run of the workspace tool call that runs now, or `undefined` outside + * a call. A `WorkspaceHooks` hook reads it to know which run wrote a file. + */ +export function writingRun() { + return callRuns.getStore() +} + /** * The tools for one workspace: `backend`, with the boundary `root`. Each * call gets the paths through `reach`, which resolves every link, and works @@ -367,8 +382,9 @@ function bindTools( backend, root, hooks: session.hooks ?? (() => []), - // The key has the root, so the fixed lock order of `patch` (by path) is - // the same for every session on this root. + // `path` is the real path from `reach`. The key has the root, so the + // fixed lock order of `patch` (by path) is the same for every session on + // this root. lock: (path, fn) => fileLocks(`${root}\0${path}`, async () => { // `reach` checked `path` before the lock. A link made since then @@ -573,8 +589,18 @@ export function createWorkspaceTools( }, bash: { commands: bashResources }, } + // Each call runs with its run in `callRuns`, for `writingRun`. + const tools = source.tools.map((tool): AnyTool => { + const { execute } = tool + if (!execute) return tool + return { + ...tool, + execute: (args: unknown, context?: ToolExecutionContext) => + callRuns.run(runOf(context), () => execute(args, context)), + } + }) return { - tools: source.tools, + tools, resources, prompt: source.prompt, bind: source.bind, diff --git a/packages/ai-harness/tests/coding-locks.test.ts b/packages/ai-harness/tests/coding-locks.test.ts index 695bac9aa4..76a0d9136f 100644 --- a/packages/ai-harness/tests/coding-locks.test.ts +++ b/packages/ai-harness/tests/coding-locks.test.ts @@ -171,17 +171,20 @@ describe('the unchanged check', () => { ].join('\n'), }, ], - ])('%s fails for a file that changed on disk after the read', async (name, args) => { - await put('a.txt', 'one\n') - const call = openSession() - await call('read_file', { path: 'a.txt' }) - await put('a.txt', 'changed on disk\n') - - await expect(call(name, args)).rejects.toThrow( - 'a.txt changed since it was read. Read it again before you edit it.', - ) - expect(await read('a.txt')).toBe('changed on disk\n') - }) + ])( + '%s fails for a file that changed on disk after the read', + async (name, args) => { + await put('a.txt', 'one\n') + const call = openSession() + await call('read_file', { path: 'a.txt' }) + await put('a.txt', 'changed on disk\n') + + await expect(call(name, args)).rejects.toThrow( + 'a.txt changed since it was read. Read it again before you edit it.', + ) + expect(await read('a.txt')).toBe('changed on disk\n') + }, + ) it("keeps the session's own writes: a second edit after the first works", async () => { await put('a.txt', 'one two\n') diff --git a/packages/ai-harness/tests/coding-snapshots.test.ts b/packages/ai-harness/tests/coding-snapshots.test.ts index 2ac4a7ec0c..f81cca1b4d 100644 --- a/packages/ai-harness/tests/coding-snapshots.test.ts +++ b/packages/ai-harness/tests/coding-snapshots.test.ts @@ -14,10 +14,7 @@ import { memoryLogStore, memoryPersistence } from '@tanstack/ai-persistence' import { createHarnessHost, defineHarness, logIdOf } from '../src' import { hostBackend } from '../src/first-party/coding/backend' import { quoteArg } from '../src/first-party/coding/search' -import { - purgeSnapshots, - snapshots, -} from '../src/first-party/coding/snapshots' +import { purgeSnapshots, snapshots } from '../src/first-party/coding/snapshots' import { workspaceTools } from '../src/first-party/coding/workspace' import { after, @@ -156,6 +153,54 @@ function fakeBackend( return backend } +/** + * A backend with a fake git. Each save of the files gives a new tree id, and + * each diff gives `a.txt`. The first `locked` saves fail, like git when + * another git process holds the index lock. + */ +function fakeGit({ locked = 0 } = {}) { + let trees = 0 + let failures = locked + return fakeBackend(async (command) => { + if (!/write-tree/.test(command)) { + const isDiff = / diff --name-only /.test(command) + return { exitCode: 0, stdout: isDiff ? 'a.txt\0' : '', stderr: '' } + } + if (failures > 0) { + failures -= 1 + return { + exitCode: 128, + stdout: '', + stderr: "fatal: Unable to create '/data/index.lock': File exists.", + } + } + trees += 1 + return { + exitCode: 0, + stdout: `${trees.toString(16).padStart(40, '0')}\n`, + stderr: '', + } + }) +} + +/** The one shadow repository in the data folder. */ +async function onlyRepo() { + const [repo, ...others] = await readdir(dataDir) + if (repo === undefined || others.length > 0) { + throw new Error('The data folder must have one repository.') + } + return join(dataDir, repo) +} + +/** The paths of the files in the tree `tree` of the repository `repo`. */ +async function filesOfTree(repo: string, tree: string) { + const quote = (value: string) => quoteArg(value, process.platform) + const result = await hostBackend.exec( + `git --git-dir ${quote(repo)} ls-tree -r --name-only ${tree}`, + ) + return result.stdout.split('\n').filter((line) => line !== '') +} + describe.skipIf(!hasGit)('snapshots', () => { // Each git call starts a shell and a process, which is slow on Windows. it('starts few git commands for each step', async () => { @@ -413,6 +458,99 @@ describe.skipIf(!hasGit)('snapshots', () => { }) }) +describe.skipIf(!hasGit)('snapshots of many runs', () => { + it("/undo of run A keeps the change that run B made during A's step", async () => { + await put('a.txt', 'one') + // B runs on the same thread and root, so both use one shadow repository. + const runB = await open({ + workspace: true, + replies: [ + () => toolCall('write_file', { path: 'b.txt', content: 'B' }, 'b-1'), + () => text('B done.'), + ], + }) + const runA = await open({ + workspace: true, + replies: [ + // B writes its file while the step of A runs. + async function* () { + await runB.session.send('Write b.txt') + yield* toolCall('write_file', { path: 'a.txt', content: 'A' }, 'a-1') + }, + () => text('A done.'), + ], + }) + + await runA.session.send('Write a.txt') + expect(await read('b.txt')).toBe('B') + expect(savedSteps(runA.session)).toMatchObject([{ files: ['a.txt'] }]) + expect(await runA.session.command({ name: 'undo' })).toBe( + 'Undid the last turn. Files restored: 1.', + ) + + expect(await read('a.txt')).toBe('one') + expect(await read('b.txt')).toBe('B') + await runA.host.close() + await runB.host.close() + }) + + it('gives two tenants on one root and one thread id two repositories', async () => { + const first = await open({ + replies: [() => text('Hi.')], + principal: { id: 'user', tenantId: 'acme' }, + }) + const second = await open({ + replies: [() => text('Hi.')], + principal: { id: 'user', tenantId: 'globex' }, + }) + + expect(await readdir(dataDir)).toHaveLength(2) + await first.host.close() + await second.host.close() + }) + + it('never stores a .env file', async () => { + await put('a.txt', 'one') + await put('.env', 'SECRET=1') + await put('.env.local', 'SECRET=2') + const { host, session } = await open({ + replies: editTurn('Done.'), + changes: [ + async () => { + await put('a.txt', 'two') + await put('.env', 'SECRET=3') + }, + ], + }) + + await session.send('Edit the files') + + const [step] = savedSteps(session) + expect(step.files).toEqual(['a.txt']) + const repo = await onlyRepo() + expect(await filesOfTree(repo, step.from)).toEqual(['a.txt']) + expect(await filesOfTree(repo, step.to)).toEqual(['a.txt']) + await host.close() + }) + + it('purgeSnapshots removes the repository of that thread only', async () => { + const first = await open({ replies: [() => text('Hi.')] }) + const second = await open({ replies: [() => text('Hi.')], threadId: 't2' }) + await first.host.close() + await second.host.close() + expect(await readdir(dataDir)).toHaveLength(2) + + await purgeSnapshots({ root, dataDir }, { threadId: THREAD }) + expect(await readdir(dataDir)).toHaveLength(1) + await purgeSnapshots({ root, dataDir }, { threadId: 't2' }) + expect(await readdir(dataDir)).toEqual([]) + // A thread with no repository: nothing to remove. + await expect( + purgeSnapshots({ root, dataDir }, { threadId: THREAD }), + ).resolves.toBeUndefined() + }) +}) + /** Three turns. Each one changes the files once, with its own tool call. */ const threeTurns = () => ({ replies: (['First.', 'Second.', 'Third.'] as const).flatMap( @@ -752,6 +890,72 @@ describe('snapshots git safety', () => { }) }) +describe('snapshots with a fake git', () => { + it('keeps the last 500 steps in the plugin state', async () => { + // 11 turns of 46 steps: 506 steps. A turn has at most 50 model calls. + let callId = 0 + const replies = Array.from( + { length: 11 }, + (): Array => [ + ...Array.from({ length: 46 }, (): Reply => { + const id = `call-${callId++}` + return () => toolCall('change', {}, id) + }), + () => text('Done.'), + ], + ).flat() + const { host, session } = await open({ replies, backend: fakeGit() }) + + for (let turn = 0; turn < 11; turn += 1) await session.send('Edit') + + const steps = savedSteps(session) + expect(steps).toHaveLength(500) + expect(steps[0]).toMatchObject({ toolCallIds: ['call-6'] }) + expect(steps[499]).toMatchObject({ toolCallIds: ['call-505'] }) + await host.close() + }, 60_000) + + it('tries again when another git process holds the index lock', async () => { + const warn = vi.spyOn(console, 'warn').mockImplementation(() => {}) + const { host, session } = await open({ + replies: editTurn('Done.'), + backend: fakeGit({ locked: 2 }), + }) + + await session.send('Edit a.txt') + + expect(savedSteps(session)).toMatchObject([{ files: ['a.txt'] }]) + // Only the snapshot warnings: the host can warn about a handle of an + // earlier test that the garbage collector found. + const snapshotWarnings = warn.mock.calls.filter(([message]) => + String(message).startsWith('snapshots:'), + ) + expect(snapshotWarnings).toEqual([]) + await host.close() + }) + + it('reports an index lock that stays, and the turn still runs', async () => { + const warn = vi.spyOn(console, 'warn').mockImplementation(() => {}) + const { host, session } = await open({ + replies: [() => text('Hi.')], + backend: fakeGit({ locked: Infinity }), + }) + + await session.send('Hello') + + expect(await session.transcript()).toMatchObject([ + { role: 'user', content: 'Hello' }, + { role: 'assistant', content: 'Hi.' }, + ]) + expect(warn).toHaveBeenCalledWith( + expect.stringMatching( + /^snapshots: another git process holds the index lock of .+: fatal: Unable to create/, + ), + ) + await host.close() + }) +}) + describe('snapshots without git', () => { it('adds no commands, warns once, and the turn still runs', async () => { const warn = vi.spyOn(console, 'warn').mockImplementation(() => {}) From dd05481806af2f23533941aba316cb67e866b75b Mon Sep 17 00:00:00 2001 From: Alem Tuzlak Date: Sat, 17 Oct 2026 17:54:42 +0200 Subject: [PATCH 04/17] fix(ai-harness, ai-sandbox, ai-sandbox-local-process): bounded background jobs with a ring buffer, a kill that waits for the spawn, workspace close, and an env allowlist in the local sandbox --- .../src/first-party/coding/backend.ts | 193 +++++++++++--- .../ai-harness/src/first-party/coding/bash.ts | 105 ++++++-- .../src/first-party/coding/ring-buffer.ts | 65 +++++ .../src/first-party/coding/workspace.ts | 17 ++ .../ai-harness/tests/coding-backend.test.ts | 34 +++ packages/ai-harness/tests/coding-bash.test.ts | 7 +- packages/ai-harness/tests/coding-jobs.test.ts | 242 ++++++++++++++++++ .../ai-harness/tests/workspace-tools.test.ts | 25 ++ .../ai-sandbox-local-process/src/handle.ts | 48 +++- .../ai-sandbox-local-process/src/provider.ts | 5 +- .../tests/local-process.test.ts | 22 ++ packages/ai-sandbox/src/workspace-backend.ts | 127 +++++++-- .../tests/workspace-backend.test.ts | 65 ++++- 13 files changed, 873 insertions(+), 82 deletions(-) create mode 100644 packages/ai-harness/src/first-party/coding/ring-buffer.ts create mode 100644 packages/ai-harness/tests/coding-jobs.test.ts diff --git a/packages/ai-harness/src/first-party/coding/backend.ts b/packages/ai-harness/src/first-party/coding/backend.ts index f3759cbb5d..3265b802ba 100644 --- a/packages/ai-harness/src/first-party/coding/backend.ts +++ b/packages/ai-harness/src/first-party/coding/backend.ts @@ -2,6 +2,7 @@ import * as childProcess from 'node:child_process' import * as fs from 'node:fs/promises' import * as nodePath from 'node:path' import { isRecord } from '../../utils' +import { createOutputTail } from './ring-buffer' import type { WorkspaceHooks } from '../workspace-hooks' /** @@ -83,7 +84,10 @@ export interface WorkspaceBackend { ) => Promise<{ exitCode: number; stdout: string; stderr: string }> /** * Start a shell command in the background. `output()` is stdout and - * stderr so far. Optional: without it, background commands give an error. + * stderr so far: the newest 1 MiB, with a first line that tells how many + * older bytes were dropped. `kill()` stops the command and every command + * it started, also right after the start. Optional: without it, + * background commands give an error. */ spawn?: ( command: string, @@ -93,6 +97,20 @@ export interface WorkspaceBackend { kill: () => void output: () => string } + /** + * The tools are done with this backend. Before they call it, they stop + * the jobs of the session. `reason` tells why: + * + * - `'closed'`: the session closed. Stop the jobs that this backend + * started, and destroy what it made for itself, like a sandbox that it + * created. + * - `'moved'`: the session continues on another machine. Keep what that + * machine can resume, like a sandbox. + * + * Never destroy a sandbox that the backend did not create. Optional: a + * backend with nothing to release has none. + */ + close?: (input: { reason: 'closed' | 'moved' }) => Promise } /** @@ -159,7 +177,9 @@ export function pathsOf(backend: WorkspaceBackend) { /** * The variables of this process that a command gets. The others, like - * provider keys and `DATABASE_URL`, stay in the server. + * provider keys and `DATABASE_URL`, stay in the server. The + * `ai-sandbox-local-process` package has a copy of these lists and of + * {@link commandEnv} in `handle.ts`. Change both. */ const KEPT_ENV = ['PATH', 'HOME', 'LANG', 'LC_ALL', 'TMPDIR', 'TEMP', 'TMP'] /** On Windows, a command also gets these, so programs can start. */ @@ -212,10 +232,121 @@ const pageDecoder = new TextDecoder() /** `exec` stops a command that writes more than 10 MiB of output. */ const MAX_EXEC_OUTPUT = 10 * 1024 * 1024 +/** `spawn` keeps the newest 1 MiB of the output of a background command. */ +const MAX_JOB_OUTPUT = 1024 * 1024 + +/** Resolves after `ms` milliseconds, or with the first of `events`. */ +const firstOf = (ms: number, events: ReadonlyArray>) => + Promise.race([ + new Promise((done) => setTimeout(done, ms).unref()), + ...events, + ]) + +/** + * Run a program of Windows by its full path, and resolve with its stdout. + * Node looks for a bare name in the current folder first, and that folder + * can be a cloned repo. A failed program resolves too. + */ +function runWindowsProgram(path: string, args: Array) { + const system = nodePath.join( + process.env.SystemRoot ?? 'C:\\Windows', + 'System32', + ) + return new Promise((done) => { + childProcess.execFile( + nodePath.join(system, path), + args, + { windowsHide: true }, + (_error, stdout) => done(String(stdout)), + ) + }) +} + +/** Kill `pid` and the processes it started, with `taskkill`. */ +const taskkill = (pid: number) => + runWindowsProgram('taskkill.exe', ['/pid', String(pid), '/T', '/F']) + +/** + * The processes that `pid` started between `from` and `to` (times in + * milliseconds). The time range skips the processes of another process + * that gets the same pid later. + */ +async function windowsChildren(pid: number, from: number, to: number) { + // One statement on each line. + const script = [ + `$from = [DateTimeOffset]::FromUnixTimeMilliseconds(${from}).LocalDateTime`, + `$to = [DateTimeOffset]::FromUnixTimeMilliseconds(${to}).LocalDateTime`, + `Get-CimInstance Win32_Process -Filter 'ParentProcessId=${pid}' | Where-Object { $_.CreationDate -ge $from -and $_.CreationDate -le $to } | ForEach-Object { $_.ProcessId }`, + ].join('\n') + const stdout = await runWindowsProgram( + 'WindowsPowerShell\\v1.0\\powershell.exe', + ['-NoProfile', '-NonInteractive', '-Command', script], + ) + return stdout + .split(/\s+/) + .filter((line) => /^\d+$/.test(line)) + .map(Number) +} + +/** True when a process with `pid` runs now. Signal 0 only checks. */ +function isAlive(pid: number) { + try { + process.kill(pid, 0) + return true + } catch { + return false + } +} + +/** + * Kill `child`, a `cmd.exe`, and every command it started. A killed + * `cmd.exe` alone leaves its commands running. `taskkill /T` finds the + * commands when it starts, so it can miss a command that the shell starts + * at the same moment. That command keeps the output pipes open, and the + * shell never gives `close`. So: + * + * 1. `taskkill /T` again while the shell runs. + * 2. When the shell ended but the pipes stay open, kill the commands that + * it started, found by their parent pid. + * 3. When the pipes are still open after that, close them, so `close` comes. + */ +async function killWindowsShell(shell: ShellProcess, pid: number) { + // `exit` comes at once after the end of the shell. So before it, the pid + // is still the shell, and a new `taskkill` cannot hit another process. + for (let tries = 0; tries < 5 && shell.exitedAt === undefined; tries++) { + await taskkill(pid) + await firstOf(500, [shell.exited]) + } + const isClosed = await firstOf(1000, [shell.closed.then(() => true)]) + const { exitedAt } = shell + if (isClosed === true || exitedAt === undefined) return + // Another process with the same pid would run now. Its commands are not ours. + if (!isAlive(pid)) { + const started = await windowsChildren( + pid, + shell.startedAt - 100, + exitedAt + 100, + ) + await Promise.all(started.map(taskkill)) + } + await firstOf(2000, [shell.closed]) + shell.child.stdout?.destroy() + shell.child.stderr?.destroy() +} + +/** A shell from {@link startShell}, and when it started and ended. */ +interface ShellProcess { + child: childProcess.ChildProcessWithoutNullStreams + startedAt: number + /** Set in the `exit` event, so it is there at once. */ + exitedAt?: number + exited: Promise + closed: Promise +} + /** * Start `command` in the system shell. `kill` stops the shell and every - * command it started: on Windows, a killed `cmd.exe` leaves its commands - * running. + * command it started, also when it comes right after the start. */ function startShell( command: string, @@ -230,23 +361,30 @@ function startShell( shell: true, detached: !isWindows, }) + // The listeners are added now, so `kill` does not miss an event that + // came before it. + const shell: ShellProcess = { + child, + startedAt: Date.now(), + exited: new Promise((done) => + child.once('exit', () => { + shell.exitedAt = Date.now() + done() + }), + ), + closed: new Promise((done) => child.once('close', () => done())), + } + let isKilled = false const kill = () => { - if (child.pid === undefined) return + // The shell did not start. `error` ends the command. + const { pid } = child + if (isKilled || pid === undefined) return + isKilled = true if (isWindows) { - // The full path: Node looks for a bare name in the current folder - // first, and that folder can be a cloned repo. - const system = nodePath.join( - process.env.SystemRoot ?? 'C:\\Windows', - 'System32', - ) - childProcess.execFile( - nodePath.join(system, 'taskkill.exe'), - ['/pid', String(child.pid), '/T', '/F'], - () => undefined, - ) + void killWindowsShell(shell, pid) return } - const group = -child.pid + const group = -pid try { process.kill(group, 'SIGTERM') } catch { @@ -426,29 +564,22 @@ export const hostBackend = { }), spawn: (command, options = {}) => { const { child, kill } = startShell(command, options) - // ponytail: keeps the last 10 MiB of a long job, like `exec` keeps the first. - let output = '' - const collect = (chunk: string) => { - output = keepTail(output + chunk) - } - // The stream decodes UTF-8, so a character split over two chunks stays whole. - child.stdout.setEncoding('utf8').on('data', collect) - child.stderr.setEncoding('utf8').on('data', collect) + // Bytes, not text: the text is made at `output()`, so a character split + // over two chunks stays whole. + const output = createOutputTail(MAX_JOB_OUTPUT) + child.stdout.on('data', output.push) + child.stderr.on('data', output.push) const exited = new Promise<{ exitCode: number }>((done) => { child.on('close', (code) => done({ exitCode: code ?? 1 })) child.on('error', (error) => { - output += String(error) + output.push(String(error)) done({ exitCode: 1 }) }) }) - return { wait: () => exited, kill, output: () => output } + return { wait: () => exited, kill, output: output.text } }, } satisfies WorkspaceBackend -/** Keep the last `MAX_EXEC_OUTPUT` characters. Cuts at twice that, so it copies rarely. */ -const keepTail = (text: string) => - text.length > 2 * MAX_EXEC_OUTPUT ? text.slice(-MAX_EXEC_OUTPUT) : text - const MAX_OUTPUT = 20_000 /** Cut text after 20,000 characters, and say how many were cut. */ diff --git a/packages/ai-harness/src/first-party/coding/bash.ts b/packages/ai-harness/src/first-party/coding/bash.ts index 00709e4015..9f972d847b 100644 --- a/packages/ai-harness/src/first-party/coding/bash.ts +++ b/packages/ai-harness/src/first-party/coding/bash.ts @@ -104,20 +104,52 @@ export function bashResources(input: unknown) { return splitCommand(stringArg(input, 'command')) } -/** The number of background jobs started, by the signal that stops them. */ -const jobCounts = new WeakMap() +/** At most this many background jobs run at a time in a session. */ +const DEFAULT_MAX_JOBS = 8 + +/** The jobs of one session: the ones that run, and how many started. */ +interface Jobs { + started: number + running: Set<{ kill: () => void }> +} + +/** + * The jobs of each session, by its signal. A session signal outlives a + * reload, so the jobs from before a reload are still found, killed, and + * counted. Each signal gets one abort listener, not one for each setup. + */ +const sessionJobs = new WeakMap() + +/** The jobs of the session of `signal`. Without it, new jobs of their own. */ +function jobsOf(signal: AbortSignal | undefined) { + const known = signal && sessionJobs.get(signal) + if (known) return known + const jobs: Jobs = { started: 0, running: new Set() } + if (!signal) return jobs + sessionJobs.set(signal, jobs) + signal.addEventListener( + 'abort', + () => { + for (const job of jobs.running) job.kill() + }, + { once: true }, + ) + return jobs +} /** * `bash`: run a shell command in the workspace folder. Every command gets * `AGENT=1` in its environment. The model gets the exit code and the last - * part of the output. + * part of the output. The abort signal of the run stops the command. * * With `background: true`, the call returns at once with a job id, and * `note` tells the model when the job ends. Background jobs need - * `env.backend.spawn`. They are killed when `signal` aborts. + * `env.backend.spawn`. At most `maxJobs` run at a time in a session. They + * are killed when `signal` aborts. * * A command that runs in the foreground supports `detach` of the tool - * context: the host can move it to the background while it runs. + * context: the host can move it to the background while it runs. Then the + * end of the run does not stop it, and `signal` kills it like a job. */ export function bashTools( env: ToolEnv, @@ -136,6 +168,11 @@ export function bashTools( signal?: AbortSignal /** Told when a background job starts and ends. A durable session logs it. */ jobs?: { started: (jobId: string) => void; ended: (jobId: string) => void } + /** + * At most this many background jobs run at a time in the session of + * `signal`. One more gives a tool error. Default: 8. + */ + maxJobs?: number } = {}, ) { const { @@ -143,17 +180,13 @@ export function bashTools( note, spillDir, signal, + maxJobs = DEFAULT_MAX_JOBS, } = options // A restart cannot reach these processes. `options.jobs` lets a durable - // session note the jobs that a crash stopped. - const jobs = new Set<{ kill: () => void }>() - // The count belongs to the signal. A session signal outlives a reload, so a - // job after a reload does not reuse the id of a job that still runs. - const count = (signal && jobCounts.get(signal)) || { started: 0 } - if (signal) jobCounts.set(signal, count) - signal?.addEventListener('abort', () => { - for (const job of jobs) job.kill() - }) + // session note the jobs that a crash stopped. The ids come from the + // session, so a job after a reload does not reuse the id of a job that + // still runs. + const jobs = jobsOf(signal) /** Save the full output in `dir`. Gives back the note for the model. */ const spill = async (dir: string, output: string) => { @@ -183,14 +216,22 @@ export function bashTools( if (!env.backend.spawn) { throw new Error('This workspace cannot run commands in the background.') } + if (signal?.aborted) throw new Error('The session is closed.') + if (jobs.running.size >= maxJobs) { + throw new Error( + `${jobs.running.size} background jobs run already, and the limit is ${maxJobs}. Wait until a job ends, then start this one.`, + ) + } const job = env.backend.spawn(command, { cwd, env: AGENT_ENV }) - count.started += 1 - const jobId = `bash-${count.started}` - jobs.add(job) + jobs.started += 1 + const jobId = `bash-${jobs.started}` + jobs.running.add(job) options.jobs?.started(jobId) const tell = async () => { - const { exitCode } = await job.wait() - jobs.delete(job) + // A job that fails to wait gives its place back too. + const { exitCode } = await job + .wait() + .finally(() => jobs.running.delete(job)) // A closed session keeps the job as running in the log, so the next // host tells the model that it stopped. if (signal?.aborted) return @@ -239,6 +280,14 @@ export function bashTools( const requested = isRecord(args) ? args.timeoutMs : undefined // A timeout of 0 would mean no timeout, so it gets the default. const isTimeout = typeof requested === 'number' && requested > 0 + // The run stops the command while it runs in the foreground. + const stop = new AbortController() + const runSignal = context?.abortSignal + const onRunAbort = () => stop.abort() + runSignal?.addEventListener('abort', onRunAbort, { once: true }) + if (runSignal?.aborted) stop.abort() + const job = { kill: () => stop.abort() } + let isDone = false const work = env.backend .exec(command, { cwd, @@ -246,6 +295,12 @@ export function bashTools( timeoutMs: isTimeout ? Math.min(requested, MAX_TIMEOUT_MS) : defaultTimeoutMs, + signal: stop.signal, + }) + .finally(() => { + isDone = true + runSignal?.removeEventListener('abort', onRunAbort) + jobs.running.delete(job) }) .then(({ exitCode, stdout, stderr }) => report(exitCode, `${stdout}${stderr ? `\nstderr:\n${stderr}` : ''}`), @@ -253,7 +308,17 @@ export function bashTools( // The host can move the command to the background. The timeout still // applies to it. const detach = context?.detach - return detach ? Promise.race([work, detach(work)]) : work + if (!detach) return work + const moved = detach(work).then((text) => { + // The run can end now, and the command goes on as a job of the + // session. The session signal kills it, not the run. + runSignal?.removeEventListener('abort', onRunAbort) + if (!isDone) jobs.running.add(job) + // The session closed already, so its listener does not run again. + if (signal?.aborted) job.kill() + return text + }) + return Promise.race([work, moved]) }), ] } diff --git a/packages/ai-harness/src/first-party/coding/ring-buffer.ts b/packages/ai-harness/src/first-party/coding/ring-buffer.ts new file mode 100644 index 0000000000..815c432037 --- /dev/null +++ b/packages/ai-harness/src/first-party/coding/ring-buffer.ts @@ -0,0 +1,65 @@ +const encoder = new TextEncoder() +const decoder = new TextDecoder() + +/** True for a byte in the middle of a UTF-8 character (`10xxxxxx`). */ +const isContinuation = (byte: number) => (byte & 0xc0) === 0x80 + +/** + * The newest output of a command, at most `maxBytes` bytes. When more comes, + * the oldest bytes are dropped. `text()` gives the output as UTF-8 text, and + * starts with a line that tells how many bytes were dropped. + * + * The `ai-sandbox` package has a copy of this function in + * `workspace-backend.ts`. Change both. + * + * @example + * ```ts + * const output = createOutputTail(1024 * 1024) + * child.stdout.on('data', output.push) + * output.text() // '[12 earlier bytes were dropped.]\n...' + * ``` + */ +export function createOutputTail(maxBytes: number) { + // The chunks from `start` on are the output. The ones before are dropped. + const chunks: Array = [] + let start = 0 + let size = 0 + let dropped = 0 + + const push = (chunk: Uint8Array | string) => { + const bytes = typeof chunk === 'string' ? encoder.encode(chunk) : chunk + // A copy of the end only, so a big chunk is not kept whole. + const kept = bytes.byteLength > maxBytes ? bytes.slice(-maxBytes) : bytes + dropped += bytes.byteLength - kept.byteLength + chunks.push(kept) + size += kept.byteLength + while (size > maxBytes) { + const first = chunks[start] + if (first === undefined) break + const over = size - maxBytes + const cut = Math.min(over, first.byteLength) + if (cut === first.byteLength) start += 1 + else chunks[start] = first.subarray(cut) + size -= cut + dropped += cut + } + // Remove the dropped chunks from the array now and then, not each time. + if (start > 1024 && start * 2 > chunks.length) { + chunks.splice(0, start) + start = 0 + } + } + + const text = () => { + let bytes = Buffer.concat(chunks.slice(start)) + if (dropped === 0) return decoder.decode(bytes) + // The cut can be in a character. Start at the next whole one. + let from = 0 + while (from < 3 && from < bytes.length && isContinuation(bytes[from] ?? 0)) + from += 1 + bytes = bytes.subarray(from) + return `[${dropped + from} earlier bytes were dropped.]\n${decoder.decode(bytes)}` + } + + return { push, text } +} diff --git a/packages/ai-harness/src/first-party/coding/workspace.ts b/packages/ai-harness/src/first-party/coding/workspace.ts index 75cd886234..7bb09601a8 100644 --- a/packages/ai-harness/src/first-party/coding/workspace.ts +++ b/packages/ai-harness/src/first-party/coding/workspace.ts @@ -79,6 +79,8 @@ export type WorkspaceToolsOptions = WorkspaceSource & { * full output in this folder. Relative to `root`, or absolute. */ spillDir?: string + /** How many background jobs can run at a time in a session. Default 8. */ + maxJobs?: number } /** What the tools need from the session for a path outside the workspace. */ @@ -436,6 +438,7 @@ function bindTools( spillDir: options.spillDir, signal: session.signal, jobs: session.jobs, + maxJobs: options.maxJobs, }), ...(options.web === false ? [] : webTools(options.web)), ] @@ -487,6 +490,19 @@ function threadTools( const bindings = new Map>() const bound = new Map() const keyOf = (run: RunContext) => `${run.logId}\0${run.threadId}` + // The workspaces that `bind` gave. They belong to the tools, so the tools + // close them when the session ends. A fixed `backend` option belongs to + // the caller and is never closed here. + const workspaces = new Set() + session.signal?.addEventListener( + 'abort', + () => { + for (const workspace of workspaces) { + void workspace.close?.({ reason: 'closed' }).catch(() => undefined) + } + }, + { once: true }, + ) /** The tools of the workspace of `run`. */ const bind = (run: RunContext | undefined) => { @@ -499,6 +515,7 @@ function threadTools( const known = bindings.get(key) if (known) return known const binding = binder.bind({ run }).then((workspace) => { + workspaces.add(workspace) const tools = bindTools( { backend: workspace, root: workspace.root }, options, diff --git a/packages/ai-harness/tests/coding-backend.test.ts b/packages/ai-harness/tests/coding-backend.test.ts index a8669a686b..81359fdb75 100644 --- a/packages/ai-harness/tests/coding-backend.test.ts +++ b/packages/ai-harness/tests/coding-backend.test.ts @@ -113,6 +113,40 @@ describe('hostBackend', () => { expect((await job.wait()).exitCode).not.toBe(0) }) + it('keeps the newest 1 MiB of the output of a background command', async () => { + // 1 MiB of x, then `end`: 3 bytes too many. + const job = hostBackend.spawn( + `node --no-warnings -e "process.stdout.write('x'.repeat(1048576) + 'end')"`, + ) + await job.wait() + const marker = '[3 earlier bytes were dropped.]\n' + const output = job.output() + expect(output.length).toBe(marker.length + 1024 * 1024) + expect(output.startsWith(marker)).toBe(true) + expect(output.endsWith('xxend')).toBe(true) + }) + + // `start /b` runs node, and `cmd.exe` ends at once. Node keeps the output + // pipes open, so the job does not end. `taskkill /T` cannot find node from + // the shell that ended. + it.runIf(process.platform === 'win32')( + 'kills a command that its shell left running', + async () => { + const job = hostBackend.spawn( + `start /b node --no-warnings -e "console.log(process.pid); setTimeout(function () {}, 60000)"`, + ) + let pid = Number.NaN + for (let tries = 0; tries < 100 && Number.isNaN(pid); tries++) { + await new Promise((done) => setTimeout(done, 100)) + pid = Number.parseInt(job.output(), 10) + } + job.kill() + await job.wait() + expect(() => process.kill(pid, 0)).toThrow() + }, + 20_000, + ) + // cmd.exe looks in the current folder before the PATH. A cloned repo can // put an `rg.cmd` there, and grep runs `rg` with no question. it.runIf(process.platform === 'win32')( diff --git a/packages/ai-harness/tests/coding-bash.test.ts b/packages/ai-harness/tests/coding-bash.test.ts index aeeeb655e9..9f028a0e94 100644 --- a/packages/ai-harness/tests/coding-bash.test.ts +++ b/packages/ai-harness/tests/coding-bash.test.ts @@ -160,7 +160,12 @@ describe('bash', () => { expect(execs).toEqual([ { command: 'make', - options: { cwd: root, env: { AGENT: '1' }, timeoutMs: 120_000 }, + options: { + cwd: root, + env: { AGENT: '1' }, + timeoutMs: 120_000, + signal: expect.any(AbortSignal), + }, }, ]) }) diff --git a/packages/ai-harness/tests/coding-jobs.test.ts b/packages/ai-harness/tests/coding-jobs.test.ts new file mode 100644 index 0000000000..87076eb0da --- /dev/null +++ b/packages/ai-harness/tests/coding-jobs.test.ts @@ -0,0 +1,242 @@ +import { getEventListeners } from 'node:events' +import { tmpdir } from 'node:os' +import { resolve } from 'node:path' +import { describe, expect, it } from 'vitest' +import { bashTools } from '../src/first-party/coding/bash' +import { createOutputTail } from '../src/first-party/coding/ring-buffer' +import type { ToolExecutionContext } from '@tanstack/ai' +import type { + ToolEnv, + WorkspaceBackend, +} from '../src/first-party/coding/backend' + +type BashOptions = Parameters[1] +type ExecOptions = Parameters[1] + +const root = resolve(tmpdir(), 'harness-jobs') + +/** A job that runs until `kill`, which ends it with 143. */ +function runningJob() { + let end!: (exitCode: number) => void + const exited = new Promise<{ exitCode: number }>((done) => { + end = (exitCode) => done({ exitCode }) + }) + const job = { + killed: false, + wait: () => exited, + kill: () => { + job.killed = true + end(143) + }, + output: () => '', + } + return { job, end } +} + +/** + * A workspace with `backend`. Its `spawn` starts a {@link runningJob} and + * keeps it in `jobs`. Its `exec` runs until its signal aborts, then gives + * exit code 1, and keeps each call in `execs`. + */ +function workspace(backend: Partial = {}) { + const jobs: Array> = [] + const execs: Array<{ options: ExecOptions; aborted: boolean }> = [] + const env: ToolEnv = { + backend: { + readFile: async (path) => { + throw new Error(`No file ${path}`) + }, + writeFile: async () => undefined, + stat: async () => undefined, + realpath: async (path) => path, + readdir: async () => [], + exec: (_command, options) => { + const call = { options, aborted: false } + execs.push(call) + return new Promise((done) => { + const onAbort = () => { + call.aborted = true + done({ exitCode: 1, stdout: '', stderr: '' }) + } + // Like a real backend: a signal that aborted already stops it too. + if (options?.signal?.aborted) onAbort() + options?.signal?.addEventListener('abort', onAbort) + }) + }, + spawn: () => { + const started = runningJob() + jobs.push(started) + return started.job + }, + ...backend, + }, + root, + hooks: () => [], + lock: (_path, fn) => fn(), + reach: async (path) => resolve(root, path), + shown: (full) => full, + } + return { env, jobs, execs } +} + +/** What a call of the tool gets from the run: its signal and `detach`. */ +type CallContext = Pick + +/** The `bash` tool of one plugin setup. */ +function bashTool(env: ToolEnv, options?: BashOptions) { + const [tool] = bashTools(env, options) + if (!tool?.execute) throw new Error('No bash tool') + const { execute } = tool + return (args: object, context: CallContext = {}) => + execute(args, { emitCustomEvent: () => undefined, ...context }) +} + +/** Start `count` background jobs with `bash`. */ +async function startJobs(bash: ReturnType, count: number) { + for (let i = 0; i < count; i++) { + await bash({ command: `job ${i}`, background: true }) + } +} + +/** Wait until the promises that are ready now have run. */ +const settle = () => new Promise((done) => setTimeout(done, 0)) + +describe('createOutputTail', () => { + it('gives all output under the limit, with no note', () => { + const output = createOutputTail(10) + output.push('ab') + output.push(new TextEncoder().encode('cd')) + expect(output.text()).toBe('abcd') + }) + + it('keeps the newest bytes over the limit, and tells how many were dropped', () => { + const output = createOutputTail(4) + output.push('abc') + output.push('def') + expect(output.text()).toBe('[2 earlier bytes were dropped.]\ncdef') + }) + + it('keeps the end of one chunk that is bigger than the limit', () => { + const output = createOutputTail(3) + output.push('abcdef') + expect(output.text()).toBe('[3 earlier bytes were dropped.]\ndef') + }) + + it('starts at a whole character when the cut is in one', () => { + const output = createOutputTail(4) + // `é` is 2 bytes. The cut drops its first byte, so it drops all of it. + output.push('é') + output.push('abc') + expect(output.text()).toBe('[2 earlier bytes were dropped.]\nabc') + }) +}) + +describe('bash job limit', () => { + it('refuses a 9th background job, and names the limit', async () => { + const { env } = workspace() + const bash = bashTool(env, { signal: new AbortController().signal }) + await startJobs(bash, 8) + await expect( + bash({ command: 'one more', background: true }), + ).rejects.toThrow( + '8 background jobs run already, and the limit is 8. Wait until a job ends, then start this one.', + ) + }) + + it('starts a job again when one ended', async () => { + const { env, jobs } = workspace() + const bash = bashTool(env, { maxJobs: 2 }) + await startJobs(bash, 2) + jobs[0]?.end(0) + await settle() + expect(await bash({ command: 'next', background: true })).toEqual({ + jobId: 'bash-3', + status: 'started', + }) + }) + + it('counts the jobs of every setup of one session', async () => { + const { env } = workspace() + const signal = new AbortController().signal + await startJobs(bashTool(env, { signal, maxJobs: 2 }), 2) + // A reload sets the plugin up again, with the same session signal. + const reloaded = bashTool(env, { signal, maxJobs: 2 }) + await expect( + reloaded({ command: 'one more', background: true }), + ).rejects.toThrow('the limit is 2') + }) +}) + +describe('bash session close', () => { + it('adds one abort listener for many setups, and kills the jobs of all of them', async () => { + const { env, jobs } = workspace() + const controller = new AbortController() + const { signal } = controller + const setups = Array.from({ length: 5 }, () => bashTool(env, { signal })) + for (const bash of setups) await startJobs(bash, 1) + expect(getEventListeners(signal, 'abort')).toHaveLength(1) + controller.abort() + expect(jobs.map(({ job }) => job.killed)).toEqual([ + true, + true, + true, + true, + true, + ]) + }) + + it('refuses a background job after the session closed', async () => { + const { env } = workspace() + const controller = new AbortController() + const bash = bashTool(env, { signal: controller.signal }) + controller.abort() + await expect( + bash({ command: 'npm run dev', background: true }), + ).rejects.toThrow('The session is closed.') + }) +}) + +describe('bash in the foreground', () => { + it('stops the command when the run aborts', async () => { + const { env, execs } = workspace() + const run = new AbortController() + const result = bashTool(env)( + { command: 'npm test' }, + { abortSignal: run.signal }, + ) + run.abort() + expect(await result).toBe('exit code: 1\n') + expect(execs[0]?.aborted).toBe(true) + }) + + it('keeps a moved command running after the run, and kills it when the session closes', async () => { + const { env, execs } = workspace() + const run = new AbortController() + const session = new AbortController() + const result = await bashTool(env, { signal: session.signal })( + { command: 'npm run dev' }, + { + abortSignal: run.signal, + detach: async () => 'Moved to the background.', + }, + ) + expect(result).toBe('Moved to the background.') + run.abort() + expect(execs[0]?.aborted).toBe(false) + session.abort() + expect(execs[0]?.aborted).toBe(true) + }) + + it('counts a moved command as a job of the session', async () => { + const { env } = workspace() + const signal = new AbortController().signal + const bash = bashTool(env, { signal, maxJobs: 1 }) + await bash( + { command: 'npm run dev' }, + { detach: async () => 'Moved to the background.' }, + ) + await expect( + bash({ command: 'one more', background: true }), + ).rejects.toThrow('the limit is 1') + }) +}) diff --git a/packages/ai-harness/tests/workspace-tools.test.ts b/packages/ai-harness/tests/workspace-tools.test.ts index 15650de074..eca8ff8064 100644 --- a/packages/ai-harness/tests/workspace-tools.test.ts +++ b/packages/ai-harness/tests/workspace-tools.test.ts @@ -1126,6 +1126,31 @@ describe('a workspace for each thread', () => { ) } + it('closes the workspace that bind gave when the session closes', async () => { + const { workspace } = threadWorkspaces() + const closed: Array = [] + const closing: WorkspaceBinder = { + bind: async (input) => ({ + ...(await workspace.bind(input)), + close: async ({ reason }) => { + closed.push(reason) + }, + }), + } + const { adapter } = mockAdapter([ + () => toolCall('read_file', { path: 'mine.txt' }, 'c0'), + () => text('done'), + ]) + const host = createHarnessHost({ persistence: memoryPersistence() }) + const session = await openThread(host, closing, adapter, 'a') + await session.send('read') + expect(closed).toEqual([]) + + await host.close() + + expect(closed).toEqual(['closed']) + }) + it('keeps a thread in its root: no path through .. or a link reaches another root', async () => { const { workspace } = threadWorkspaces() const { adapter, calls } = mockAdapter([ diff --git a/packages/ai-sandbox-local-process/src/handle.ts b/packages/ai-sandbox-local-process/src/handle.ts index cb066394be..7657dedbe0 100644 --- a/packages/ai-sandbox-local-process/src/handle.ts +++ b/packages/ai-sandbox-local-process/src/handle.ts @@ -111,6 +111,45 @@ function prependShellPath(env: NodeJS.ProcessEnv): NodeJS.ProcessEnv { return env } +/** + * The variables of this process that a command gets. The others, like + * provider keys and `DATABASE_URL`, stay in the server. + */ +const KEPT_ENV = ['PATH', 'HOME', 'LANG', 'LC_ALL', 'TMPDIR', 'TEMP', 'TMP'] +/** On Windows, a command also gets these, so programs can start. */ +const KEPT_WINDOWS_ENV = [ + 'SystemRoot', + 'ComSpec', + 'PATHEXT', + 'WINDIR', + 'USERPROFILE', +] + +/** + * The environment of a command: the variables of this process in a short + * list (`PATH`, `HOME`, `LANG`, `LC_ALL`, `TMPDIR`, `TEMP`, `TMP`, and on + * Windows `SystemRoot`, `ComSpec`, `PATHEXT`, `WINDIR`, `USERPROFILE`), plus + * `env`. No other variable of this process goes to the command. On Windows + * it sets `NoDefaultCurrentDirectoryInExePath`, so a program name is found + * on the PATH only, and not in the current folder. + * + * A copy of `commandEnv` in `@tanstack/ai-harness` + * (`src/first-party/coding/backend.ts`). This package does not depend on + * that one. Change both. + */ +function commandEnv(env: Record, platform: NodeJS.Platform) { + const isWindows = platform === 'win32' + const names = isWindows ? [...KEPT_ENV, ...KEPT_WINDOWS_ENV] : KEPT_ENV + const kept: Record = {} + for (const name of names) { + // On Windows, `process.env` finds a name in any letter case. + const value = process.env[name] + if (value !== undefined) kept[name] = value + } + if (!isWindows) return { ...kept, ...env } + return { ...kept, ...env, NoDefaultCurrentDirectoryInExePath: '1' } +} + export const LOCAL_PROCESS_CAPS: SandboxCapabilities = { fs: true, exec: true, @@ -581,7 +620,7 @@ export interface LocalProcessHandleOptions { removeOnDestroy: boolean /** Create a fork by copying this sandbox's dir to a new root. */ forkFactory: (sourceRoot: string) => Promise - /** Env vars to delete from the inherited `process.env` before spawning. */ + /** Env vars to delete from the env of each command (see `commandEnv`). */ scrubEnv?: Array /** * Sink for non-fatal teardown diagnostics — currently a `killTree` that could @@ -842,7 +881,12 @@ export class LocalProcessHandle implements SandboxHandle { } private mergedEnv(extra?: Record): NodeJS.ProcessEnv { - const env: NodeJS.ProcessEnv = { ...process.env, ...this.envVars, ...extra } + // Only a few variables of the server, plus the env that was given. A + // provider key or `DATABASE_URL` of the server does not reach a command. + const env: NodeJS.ProcessEnv = commandEnv( + { ...this.envVars, ...extra }, + process.platform, + ) // Drop scrubbed vars so a host CLI falls back to its own stored auth // (e.g. remove ANTHROPIC_API_KEY → Claude Code uses the logged-in // subscription instead of billing the API). Delete (not blank) so the var diff --git a/packages/ai-sandbox-local-process/src/provider.ts b/packages/ai-sandbox-local-process/src/provider.ts index bd3f9a919b..8cccf2687f 100644 --- a/packages/ai-sandbox-local-process/src/provider.ts +++ b/packages/ai-sandbox-local-process/src/provider.ts @@ -30,7 +30,10 @@ export interface LocalProcessSandboxConfig { /** Remove the backing dir on destroy. Defaults: true (generated), false (fixed `dir`). */ removeOnDestroy?: boolean /** - * Env vars to remove from the inherited `process.env` before spawning. Use to + * Env vars to remove from the env of each command. A command gets only a + * few variables of this process (`PATH`, `HOME`, `LANG`, `LC_ALL`, `TMPDIR`, + * `TEMP`, `TMP`, and on Windows `SystemRoot`, `ComSpec`, `PATHEXT`, `WINDIR`, + * `USERPROFILE`), plus the env that `env.set` and the call give. Use it to * let a host CLI fall back to its own stored auth — e.g. scrub * `ANTHROPIC_API_KEY` so Claude Code uses your logged-in subscription instead * of billing the API. diff --git a/packages/ai-sandbox-local-process/tests/local-process.test.ts b/packages/ai-sandbox-local-process/tests/local-process.test.ts index 38503a81e2..c516645104 100644 --- a/packages/ai-sandbox-local-process/tests/local-process.test.ts +++ b/packages/ai-sandbox-local-process/tests/local-process.test.ts @@ -121,6 +121,28 @@ describe('local-process process', () => { await sbx.destroy() }) + it('does not give a variable of the server to a command', async () => { + const sbx = await fresh() + process.env.LP_SERVER_SECRET = 'hunter2' + try { + const r = await sbx.process.exec('echo "${LP_SERVER_SECRET:-unset}"') + expect(r.stdout.trim()).toBe('unset') + } finally { + delete process.env.LP_SERVER_SECRET + await sbx.destroy() + } + }) + + it('gives a command the env of env.set and of the call', async () => { + const sbx = await fresh() + await sbx.env.set({ LP_SET: 'from-set' }) + const r = await sbx.process.exec('echo "$LP_SET $LP_CALL"', { + env: { LP_CALL: 'from-call' }, + }) + expect(r.stdout.trim()).toBe('from-set from-call') + await sbx.destroy() + }) + it('spawn streams stdout and resolves wait()', async () => { const sbx = await fresh() const proc = await sbx.process.spawn('echo streamed') diff --git a/packages/ai-sandbox/src/workspace-backend.ts b/packages/ai-sandbox/src/workspace-backend.ts index 6293bb7930..9520365d06 100644 --- a/packages/ai-sandbox/src/workspace-backend.ts +++ b/packages/ai-sandbox/src/workspace-backend.ts @@ -3,15 +3,83 @@ import { shellQuote } from './journal' import type { WorkspaceBackend } from '@tanstack/ai-harness/plugins/coding' import type { ExecResult, SandboxHandle } from './contracts' -/** The most output a background job keeps, in characters. */ -const MAX_OUTPUT = 10 * 1024 * 1024 +/** A background job keeps the newest 1 MiB of its output. */ +const MAX_JOB_OUTPUT = 1024 * 1024 /** `realpath` follows at most this many links, like Linux. */ const MAX_LINKS = 40 +const encoder = new TextEncoder() +const decoder = new TextDecoder() + +/** True for a byte in the middle of a UTF-8 character (`10xxxxxx`). */ +const isContinuation = (byte: number) => (byte & 0xc0) === 0x80 + +/** + * The newest output of a command, at most `maxBytes` bytes. When more comes, + * the oldest bytes are dropped. `text()` gives the output as UTF-8 text, and + * starts with a line that tells how many bytes were dropped. + * + * A copy of `createOutputTail` in `@tanstack/ai-harness` + * (`src/first-party/coding/ring-buffer.ts`). That package does not export + * it. Change both. + */ +function createOutputTail(maxBytes: number) { + // The chunks from `start` on are the output. The ones before are dropped. + const chunks: Array = [] + let start = 0 + let size = 0 + let dropped = 0 + + const push = (chunk: Uint8Array | string) => { + const bytes = typeof chunk === 'string' ? encoder.encode(chunk) : chunk + // A copy of the end only, so a big chunk is not kept whole. + const kept = bytes.byteLength > maxBytes ? bytes.slice(-maxBytes) : bytes + dropped += bytes.byteLength - kept.byteLength + chunks.push(kept) + size += kept.byteLength + while (size > maxBytes) { + const first = chunks[start] + if (first === undefined) break + const over = size - maxBytes + const cut = Math.min(over, first.byteLength) + if (cut === first.byteLength) start += 1 + else chunks[start] = first.subarray(cut) + size -= cut + dropped += cut + } + // Remove the dropped chunks from the array now and then, not each time. + if (start > 1024 && start * 2 > chunks.length) { + chunks.splice(0, start) + start = 0 + } + } + + const text = () => { + let bytes = Buffer.concat(chunks.slice(start)) + if (dropped === 0) return decoder.decode(bytes) + // The cut can be in a character. Start at the next whole one. + let from = 0 + while (from < 3 && from < bytes.length && isContinuation(bytes[from] ?? 0)) + from += 1 + bytes = bytes.subarray(from) + return `[${dropped + from} earlier bytes were dropped.]\n${decoder.decode(bytes)}` + } + + return { push, text } +} + /** * A `WorkspaceBackend` over a live sandbox. With it, `workspaceTools()` * reads and writes files and runs commands in the sandbox, not on this - * machine. You start the sandbox before, and you stop it after. + * machine. + * + * Who stops the sandbox: + * + * - By default, you. `close()` of the backend stops the jobs that it + * started, and leaves the sandbox running. + * - With `owned: true`, the backend. `close({ reason: 'closed' })` also + * destroys the sandbox. Use it for a sandbox that you create in + * `workspace.bind`, one for each thread. * * - Paths are POSIX paths in the sandbox, like `/workspace/src/a.ts`. * - `exec` gives exit code 124 when `timeoutMs` stops a command, and 1 when @@ -25,19 +93,35 @@ const MAX_LINKS = 40 * `readlink` in the sandbox reads its target. A sandbox without `lstat` * cannot find links, so `realpath` throws, and the tools refuse each path. * + * - `spawn` keeps the newest 1 MiB of the output of a job. + * - `close({ reason: 'moved' })` keeps the jobs and the sandbox, so the + * next machine can resume them. + * * @param handle - The sandbox, from `provider.create()` or `provider.resume()`. + * @param options - `owned: true` when the backend owns the sandbox. Then + * `close` destroys it. Default: `false`. * * @example * ```ts - * const handle = await localProcessSandbox().create({}) + * const provider = localProcessSandbox() * workspaceTools({ - * root: '/workspace', - * backend: sandboxWorkspaceBackend(handle), + * workspace: { + * bind: async () => ({ + * ...sandboxWorkspaceBackend(await provider.create({}), { owned: true }), + * root: '/workspace', + * }), + * }, * }) * ``` */ -export function sandboxWorkspaceBackend(handle: SandboxHandle) { +export function sandboxWorkspaceBackend( + handle: SandboxHandle, + { owned = false }: { owned?: boolean } = {}, +) { const { fs } = handle + // The jobs that this backend started and that still run. + const running = new Set<{ kill: () => void }>() + let destroyed: Promise | undefined // The read of the sandbox follows links, so this follows links too. A // path that `readBytes` cannot read is a folder. @@ -136,13 +220,9 @@ export function sandboxWorkspaceBackend(handle: SandboxHandle) { } const spawn: WorkspaceBackend['spawn'] = (command, options) => { - // ponytail: keeps the last 10 MiB of a long job, like `hostBackend`. - let output = '' + const output = createOutputTail(MAX_JOB_OUTPUT) const collect = async (stream: AsyncIterable) => { - for await (const chunk of stream) { - output += chunk - if (output.length > 2 * MAX_OUTPUT) output = output.slice(-MAX_OUTPUT) - } + for await (const chunk of stream) output.push(chunk) } const started = handle.process.spawn(command, options) // Wait for the end of the output too, so `output()` is complete. @@ -156,17 +236,29 @@ export function sandboxWorkspaceBackend(handle: SandboxHandle) { return { exitCode } }) .catch((error: unknown) => { - output += String(error) + output.push(String(error)) return { exitCode: 1 } }) - return { + const job = { wait: () => exited, // Before the process starts, this kills it when it starts. kill: () => { - started.then((job) => job.kill()).catch(() => undefined) + started.then((each) => each.kill()).catch(() => undefined) }, - output: () => output, + output: output.text, } + running.add(job) + void exited.then(() => running.delete(job)) + return job + } + + const close: NonNullable = async ({ reason }) => { + // The next machine resumes the sandbox and its jobs. + if (reason === 'moved') return + for (const job of running) job.kill() + if (!owned) return + destroyed ??= handle.destroy() + await destroyed } return { @@ -231,5 +323,6 @@ export function sandboxWorkspaceBackend(handle: SandboxHandle) { }) }, ...(handle.capabilities.backgroundProcesses ? { spawn } : {}), + close, } satisfies WorkspaceBackend } diff --git a/packages/ai-sandbox/tests/workspace-backend.test.ts b/packages/ai-sandbox/tests/workspace-backend.test.ts index bb84fbd7bb..adcf1fe38f 100644 --- a/packages/ai-sandbox/tests/workspace-backend.test.ts +++ b/packages/ai-sandbox/tests/workspace-backend.test.ts @@ -16,6 +16,8 @@ import type { const encoder = new TextEncoder() const decoder = new TextDecoder() +/** The output that a background job keeps: 1 MiB. */ +const MIB = 1024 * 1024 interface FsSeed { files?: Record @@ -112,7 +114,8 @@ async function* fromChunks(chunks: Array) { * that cannot stop it, and only notes the abort in `aborted`. * * `spawn` prints ` out` and ` err`. `serve` runs until it - * is killed, then ends with 143. `nope` does not start. + * is killed, then ends with 143. `nope` does not start. `flood` prints 1 MiB + * of `x` before `flood out`. * * `readlink -- ''` prints the target of the link `name` in the folder, * from `links`, like `readlink` in a sandbox. @@ -151,6 +154,7 @@ function fakeProcesses(links: ReadonlyMap = new Map()) { ? new Promise((resolve) => (end = resolve)) : Promise.resolve() async function* stdout() { + if (command === 'flood') yield 'x'.repeat(MIB) yield `${command} out\n` await ended } @@ -173,20 +177,24 @@ function fakeProcesses(links: ReadonlyMap = new Map()) { return { process, aborted, killed } } -/** A sandbox handle with `memoryFs` and `fakeProcesses`. */ +/** + * A sandbox handle with `memoryFs` and `fakeProcesses`, and its backend. + * `sandbox.destroyed` tells whether the sandbox was destroyed. + */ function fakeSandbox( seed: FsSeed = {}, - options: { withLstat?: boolean; caps?: SandboxCapabilities } = {}, + options: { + withLstat?: boolean + caps?: SandboxCapabilities + owned?: boolean + } = {}, ) { const { fs, files, dirs, links } = memoryFs(seed, options.withLstat) const { process, aborted, killed } = fakeProcesses(links) - const handle: SandboxHandle = { - ...makeFakeHandle('sandbox-1', 'fake', options.caps), - fs, - process, - } - const backend = sandboxWorkspaceBackend(handle) - return { backend, files, dirs, aborted, killed } + const sandbox = makeFakeHandle('sandbox-1', 'fake', options.caps) + const handle: SandboxHandle = { ...sandbox, fs, process } + const backend = sandboxWorkspaceBackend(handle, { owned: options.owned }) + return { backend, files, dirs, aborted, killed, sandbox } } const SEED: FsSeed = { @@ -398,6 +406,43 @@ describe('sandboxWorkspaceBackend spawn', () => { ) expect(backend.spawn).toBeUndefined() }) + + it('keeps the newest 1 MiB of the output, and tells how much was dropped', async () => { + const { backend } = fakeSandbox() + const job = backend.spawn?.('flood') + await job?.wait() + // 1 MiB of x, then `flood out\n` and `flood err\n`: 20 bytes too many. + const marker = '[20 earlier bytes were dropped.]\n' + const output = job?.output() ?? '' + expect(output.startsWith(marker)).toBe(true) + expect(output.length).toBe(marker.length + MIB) + expect(output.endsWith('xflood err\nflood out\n')).toBe(true) + }) +}) + +describe('sandboxWorkspaceBackend close', () => { + it('kills its jobs, and leaves a sandbox that it did not create', async () => { + const { backend, killed, sandbox } = fakeSandbox() + const job = backend.spawn?.('serve') + await backend.close({ reason: 'closed' }) + expect(await job?.wait()).toEqual({ exitCode: 143 }) + expect(killed).toEqual(['serve']) + expect(sandbox.destroyed).toBe(false) + }) + + it('destroys a sandbox that it owns', async () => { + const { backend, sandbox } = fakeSandbox({}, { owned: true }) + await backend.close({ reason: 'closed' }) + expect(sandbox.destroyed).toBe(true) + }) + + it('keeps the jobs and an owned sandbox when the session moves', async () => { + const { backend, killed, sandbox } = fakeSandbox({}, { owned: true }) + backend.spawn?.('serve') + await backend.close({ reason: 'moved' }) + expect(killed).toEqual([]) + expect(sandbox.destroyed).toBe(false) + }) }) const now = () => Date.now() From 8a5f44fe1fa13972ea47ac1d00926326c977aa8e Mon Sep 17 00:00:00 2001 From: Alem Tuzlak Date: Sat, 17 Oct 2026 18:17:43 +0200 Subject: [PATCH 05/17] perf(ai-harness, ai-sandbox): size caps for webfetch and read_file, range reads in the sandbox, one grep listing, and coding events for jobs, snapshots, and the formatter --- .../src/first-party/coding/backend.ts | 68 +++++-- .../ai-harness/src/first-party/coding/bash.ts | 81 +++++++- .../src/first-party/coding/events.ts | 41 +++++ .../src/first-party/coding/formatter.ts | 44 +++-- .../src/first-party/coding/index.ts | 7 + .../ai-harness/src/first-party/coding/read.ts | 26 ++- .../src/first-party/coding/search.ts | 57 ++++-- .../src/first-party/coding/snapshots.ts | 56 +++--- .../ai-harness/src/first-party/coding/web.ts | 23 ++- .../src/first-party/coding/workspace.ts | 6 + .../ai-harness/tests/coding-backend.test.ts | 44 ++++- .../ai-harness/tests/coding-formatter.test.ts | 88 ++++++++- packages/ai-harness/tests/coding-jobs.test.ts | 108 ++++++++++- packages/ai-harness/tests/coding-read.test.ts | 41 ++++- .../ai-harness/tests/coding-search.test.ts | 38 +++- .../ai-harness/tests/coding-snapshots.test.ts | 46 ++--- packages/ai-harness/tests/coding-web.test.ts | 26 ++- .../ai-harness/tests/workspace-tools.test.ts | 45 ++++- packages/ai-sandbox/src/workspace-backend.ts | 160 +++++++++++++--- .../tests/workspace-backend.test.ts | 174 +++++++++++++++++- 20 files changed, 1009 insertions(+), 170 deletions(-) create mode 100644 packages/ai-harness/src/first-party/coding/events.ts diff --git a/packages/ai-harness/src/first-party/coding/backend.ts b/packages/ai-harness/src/first-party/coding/backend.ts index 3265b802ba..ec009bd5c4 100644 --- a/packages/ai-harness/src/first-party/coding/backend.ts +++ b/packages/ai-harness/src/first-party/coding/backend.ts @@ -24,12 +24,16 @@ export interface WorkspaceBackend { readFile: (path: string) => Promise /** * One page of a file, without the whole file in memory: the first 8 KiB - * (`head`), the size in bytes, the lines from line `offset` (from 1) as - * UTF-8 text, and the number of lines. It gives at most `limit` lines, and - * fewer when the lines are long: no more lines after 64 KiB. A final - * newline ends the last line. It does not start a new one. Throws when - * there is no file. Optional: without it, `read_file` reads the whole - * file. + * (`head`), the size in bytes, and the lines from line `offset` (from 1) + * as UTF-8 text. It gives at most `limit` lines, and fewer when the lines + * are long: no more lines after 64 KiB. A final newline ends the last + * line. It does not start a new one. Throws when there is no file. + * + * It stops reading after the page. `total`, the number of lines, is there + * only when the read got to the end of the file. `undefined` means that + * the file has more lines after the page. + * + * Optional: without it, `read_file` reads the whole file. */ readLines?: ( path: string, @@ -38,7 +42,7 @@ export interface WorkspaceBackend { head: Uint8Array size: number lines: Array - total: number + total?: number }> /** Create or replace a file. Missing parent folders are created. */ writeFile: (path: string, data: Uint8Array | string) => Promise @@ -86,12 +90,18 @@ export interface WorkspaceBackend { * Start a shell command in the background. `output()` is stdout and * stderr so far: the newest 1 MiB, with a first line that tells how many * older bytes were dropped. `kill()` stops the command and every command - * it started, also right after the start. Optional: without it, - * background commands give an error. + * it started, also right after the start. `onOutput` gets the new text + * of stdout and stderr as it comes. A backend that cannot stream the + * output does not call it. Optional: without `spawn`, background commands + * give an error. */ spawn?: ( command: string, - options?: { cwd?: string; env?: Record }, + options?: { + cwd?: string + env?: Record + onOutput?: (text: string) => void + }, ) => { wait: () => Promise<{ exitCode: number }> kill: () => void @@ -426,9 +436,12 @@ export const hostBackend = { let end = offset + limit const page: Array = [] let pageBytes = 0 + // True when a byte of line `end` was found: the file goes on after + // the page, and the read stops there. + let hasMore = false // ponytail: reads up to the size at the start. Bytes added later wait // for the next read. - for (let position = 0; position < size;) { + for (let position = 0; position < size && !hasMore;) { const { bytesRead } = await handle.read( chunk, 0, @@ -446,6 +459,10 @@ export const hostBackend = { let pageTo = -1 // Count the lines with `indexOf`. Only the page is decoded. for (let from = 0; from < bytesRead;) { + if (line >= end) { + hasMore = true + break + } const newline = data.indexOf(10, from) const to = newline === -1 ? bytesRead : newline + 1 if (line >= offset && line < end) { @@ -461,11 +478,17 @@ export const hostBackend = { if (pageFrom !== -1) page.push(Buffer.from(data.subarray(pageFrom, pageTo))) } - // An empty file has one empty line, like `''.split('\n')`. - const total = line - 1 + (size === 0 || last !== 10 ? 1 : 0) const text = pageDecoder.decode(Buffer.concat(page)) - const lines = offset > total ? [] : text.split('\n') - if (text.endsWith('\n')) lines.pop() + const split = text.split('\n') + if (text.endsWith('\n')) split.pop() + // An empty file has one empty line, like `''.split('\n')`. + const total = hasMore + ? undefined + : line - 1 + (size === 0 || last !== 10 ? 1 : 0) + // No lines for an offset after the end, or a page of no lines, like + // with a limit of 0. + const isEmpty = total === undefined ? text === '' : offset > total + const lines = isEmpty ? [] : split return { head: head ?? Buffer.alloc(0), size, lines, total } } finally { await handle.close() @@ -567,8 +590,19 @@ export const hostBackend = { // Bytes, not text: the text is made at `output()`, so a character split // over two chunks stays whole. const output = createOutputTail(MAX_JOB_OUTPUT) - child.stdout.on('data', output.push) - child.stderr.on('data', output.push) + const { onOutput } = options + /** Keep the chunks of `stream`, and give their text to `onOutput`. */ + const collect = (stream: NodeJS.ReadableStream) => { + // One decoder for each stream keeps a character that is split over + // two chunks whole. + const decoder = new TextDecoder() + stream.on('data', (chunk: Buffer) => { + output.push(chunk) + if (onOutput) onOutput(decoder.decode(chunk, { stream: true })) + }) + } + collect(child.stdout) + collect(child.stderr) const exited = new Promise<{ exitCode: number }>((done) => { child.on('close', (code) => done({ exitCode: code ?? 1 })) child.on('error', (error) => { diff --git a/packages/ai-harness/src/first-party/coding/bash.ts b/packages/ai-harness/src/first-party/coding/bash.ts index 9f972d847b..9fe429f156 100644 --- a/packages/ai-harness/src/first-party/coding/bash.ts +++ b/packages/ai-harness/src/first-party/coding/bash.ts @@ -2,6 +2,9 @@ import { toolDefinition } from '@tanstack/ai' import { isRecord } from '../../utils' import { boundText } from '../bound-output' import { pathsOf, stringArg } from './backend' +import { JobEnded, JobOutput, JobStarted } from './events' +import { createOutputTail } from './ring-buffer' +import type { PluginEvent } from '../../extensions' import type { ToolEnv } from './backend' /** The longest timeout that one call can ask for: 10 minutes. */ @@ -106,6 +109,50 @@ export function bashResources(input: unknown) { /** At most this many background jobs run at a time in a session. */ const DEFAULT_MAX_JOBS = 8 +/** A job sends at most one output event in this time. */ +const OUTPUT_EVERY_MS = 250 +/** One output event has at most the newest 64 KiB. */ +const MAX_EVENT_OUTPUT = 64 * 1024 + +/** Sends a plugin event, like `ctx.emit` of a plugin. */ +type Emit = (event: PluginEvent, value: T) => void + +/** + * The output events of job `jobId`. `push` keeps new text. The text goes + * out at most once each {@link OUTPUT_EVERY_MS}. `flush` sends what waits + * now, and `stop` drops it. + */ +function outputEvents(jobId: string, emit: Emit) { + let waiting = createOutputTail(MAX_EVENT_OUTPUT) + let hasText = false + let timer: ReturnType | undefined + let lastSent = 0 + const flush = () => { + clearTimeout(timer) + timer = undefined + if (!hasText) return + const output = waiting.text() + waiting = createOutputTail(MAX_EVENT_OUTPUT) + hasText = false + lastSent = Date.now() + emit(JobOutput, { jobId, output }) + } + const push = (text: string) => { + if (text === '') return + waiting.push(text) + hasText = true + if (timer) return + const wait = Math.max(0, lastSent + OUTPUT_EVERY_MS - Date.now()) + timer = setTimeout(flush, wait) + // The timer must not keep the process alive. + timer.unref?.() + } + const stop = () => { + clearTimeout(timer) + timer = undefined + } + return { push, flush, stop } +} /** The jobs of one session: the ones that run, and how many started. */ interface Jobs { @@ -145,7 +192,9 @@ function jobsOf(signal: AbortSignal | undefined) { * With `background: true`, the call returns at once with a job id, and * `note` tells the model when the job ends. Background jobs need * `env.backend.spawn`. At most `maxJobs` run at a time in a session. They - * are killed when `signal` aborts. + * are killed when `signal` aborts. With `emit`, each job sends events: + * `tanstack/workspace-tools:job-started`, `tanstack/workspace-tools:job-output` (at most one each 250 ms), and + * `tanstack/workspace-tools:job-ended`. * * A command that runs in the foreground supports `detach` of the tool * context: the host can move it to the background while it runs. Then the @@ -168,6 +217,11 @@ export function bashTools( signal?: AbortSignal /** Told when a background job starts and ends. A durable session logs it. */ jobs?: { started: (jobId: string) => void; ended: (jobId: string) => void } + /** + * Sends the events of the background jobs: {@link JobStarted}, + * {@link JobOutput}, and {@link JobEnded}. Without it, there are none. + */ + emit?: Emit /** * At most this many background jobs run at a time in the session of * `signal`. One more gives a tool error. Default: 8. @@ -222,19 +276,36 @@ export function bashTools( `${jobs.running.size} background jobs run already, and the limit is ${maxJobs}. Wait until a job ends, then start this one.`, ) } - const job = env.backend.spawn(command, { cwd, env: AGENT_ENV }) + const { emit } = options + // The output can come before `spawn` returns, so the id is made first. + const jobId = `bash-${jobs.started + 1}` + const events = emit && outputEvents(jobId, emit) + const job = env.backend.spawn(command, { + cwd, + env: AGENT_ENV, + ...(events ? { onOutput: events.push } : {}), + }) jobs.started += 1 - const jobId = `bash-${jobs.started}` jobs.running.add(job) options.jobs?.started(jobId) + emit?.(JobStarted, { jobId, command }) const tell = async () => { // A job that fails to wait gives its place back too. const { exitCode } = await job .wait() + .catch((error: unknown) => { + events?.stop() + throw error + }) .finally(() => jobs.running.delete(job)) // A closed session keeps the job as running in the log, so the next - // host tells the model that it stopped. - if (signal?.aborted) return + // host tells the model that it stopped. No one gets its events now. + if (signal?.aborted) { + events?.stop() + return + } + events?.flush() + emit?.(JobEnded, { jobId, exitCode }) try { if (note === undefined) return const result = await report(exitCode, job.output()) diff --git a/packages/ai-harness/src/first-party/coding/events.ts b/packages/ai-harness/src/first-party/coding/events.ts new file mode 100644 index 0000000000..7dcf3ad2fb --- /dev/null +++ b/packages/ai-harness/src/first-party/coding/events.ts @@ -0,0 +1,41 @@ +import { createPluginEvent } from '../../extensions' + +/** A background `bash` job started. */ +export const JobStarted = createPluginEvent<{ + /** The id that the model gets, like `bash-1`. */ + jobId: string + command: string +}>('tanstack/workspace-tools:job-started') + +/** + * New output of a background `bash` job: stdout and stderr, in the order + * they came. At most one event every 250 ms for each job, with the output + * since the event before it. When a job writes more than 64 KiB in that + * time, the event has the newest 64 KiB, and starts with a line that tells + * how many bytes were dropped. + */ +export const JobOutput = createPluginEvent<{ jobId: string; output: string }>( + 'tanstack/workspace-tools:job-output', +) + +/** A background `bash` job ended. It comes after its last output event. */ +export const JobEnded = createPluginEvent<{ jobId: string; exitCode: number }>( + 'tanstack/workspace-tools:job-ended', +) + +/** + * The `snapshots()` plugin could not take or save a snapshot. The turn goes + * on. When git is missing, or the shadow repository cannot be made, the + * event comes once for each session, and snapshots are off in it. + */ +export const SnapshotFailed = createPluginEvent<{ message: string }>( + 'tanstack/snapshots:failed', +) + +/** The `formatter()` plugin changed a file after a tool wrote it. */ +export const FormatterRewrote = createPluginEvent<{ + /** The absolute path of the file. */ + path: string + /** The name of the formatter, like `prettier`. */ + formatter: string +}>('tanstack/formatter:rewrote') diff --git a/packages/ai-harness/src/first-party/coding/formatter.ts b/packages/ai-harness/src/first-party/coding/formatter.ts index ad46a78431..fab404e7d2 100644 --- a/packages/ai-harness/src/first-party/coding/formatter.ts +++ b/packages/ai-harness/src/first-party/coding/formatter.ts @@ -4,6 +4,7 @@ import { isRecord } from '../../utils' import { PermissionPrompt, PermissionRules } from '../permissions' import { WorkspaceHooks } from '../workspace-hooks' import { hostBackend, pathsOf, readText } from './backend' +import { FormatterRewrote } from './events' import { quoteArg, shellPlatform } from './search' import type { WorkspaceBackend } from './backend' @@ -138,15 +139,22 @@ const messageOf = (error: unknown) => /** * The `afterWrite` hook of {@link formatter}. It runs the first formatter * that the project uses for the file type. The project is checked at the - * first write, once. A failure goes to `onFailure` and never throws. - * `allow` says if the formatter `name` can run on `path`. When it says no, - * the hook resolves to a note for the tool result. + * first write, once. + * + * - `onFailure` gets a failure. The hook never throws. + * - `onRewrite` gets the file and the formatter when a run changed the file. + * - `allow` says if the formatter `name` can run on `path`. When it says no, + * the hook resolves to a note for the tool result. */ export function formatOnWrite( options: FormatterOptions, - onFailure: (failure: FormatFailure) => void, - allow?: (name: string, path: string) => Promise, + report: { + onFailure: (failure: FormatFailure) => void + onRewrite?: (rewrite: { path: string; formatter: string }) => void + allow?: (name: string, path: string) => Promise + }, ) { + const { onFailure, onRewrite, allow } = report const backend = options.backend ?? hostBackend const { extname, join, resolve } = pathsOf(backend) const root = resolve(options.root) @@ -188,8 +196,19 @@ export function formatOnWrite( return `The formatter ${match.name} did not run: it is not allowed.` } const command = match.command(quoteArg(path, platform)) + // The bytes before the run, to tell if the formatter changed the file. + // A file that cannot be read gives no rewrite event, and no failure. + const read = () => backend.readFile(path).catch(() => undefined) + const before = onRewrite && (await read()) const result = await backend.exec(command, { cwd: root, timeoutMs }) - if (result.exitCode === 0) return undefined + if (result.exitCode === 0) { + if (!onRewrite || !before) return undefined + const after = await read() + const isChanged = + after !== undefined && !Buffer.from(before).equals(after) + if (isChanged) onRewrite({ path, formatter: match.name }) + return undefined + } const reason = result.exitCode === 124 ? `it took longer than ${timeoutMs / 1000} seconds` @@ -222,7 +241,8 @@ export function formatOnWrite( * Each run is a command through `backend.exec` in `root`, with a time limit * (`timeoutMs`, default 20 seconds). A formatter that fails, times out, or * is not installed does not fail the tool call: the file stays as written, - * and the plugin sends a {@link FormatFailed} event. + * and the plugin sends a {@link FormatFailed} event. A run that changes the + * file sends a `tanstack/formatter:rewrote` event (`FormatterRewrote`). * * A formatter runs code from the project, like its config files. So with * `permissions()`, each run is a command named `formatter:`, for @@ -250,11 +270,11 @@ export function formatter(options: FormatterOptions) { setup: (ctx) => ({ contribute: [ WorkspaceHooks.item({ - afterWrite: formatOnWrite( - options, - (failure) => ctx.emit(FormatFailed, failure), + afterWrite: formatOnWrite(options, { + onFailure: (failure) => ctx.emit(FormatFailed, failure), + onRewrite: (rewrite) => ctx.emit(FormatterRewrote, rewrite), // Read at each write, so permissions() can come after this plugin. - async (name, path) => { + allow: async (name, path) => { const ask = ctx.getOptional(PermissionPrompt) return ask ? ask( @@ -263,7 +283,7 @@ export function formatter(options: FormatterOptions) { ) : true }, - ), + }), }), PermissionRules.item({ tool: 'formatter:*', diff --git a/packages/ai-harness/src/first-party/coding/index.ts b/packages/ai-harness/src/first-party/coding/index.ts index faae8310a7..36305a8b75 100644 --- a/packages/ai-harness/src/first-party/coding/index.ts +++ b/packages/ai-harness/src/first-party/coding/index.ts @@ -13,4 +13,11 @@ export type { SnapshotsOptions, } from './snapshots' export { FormatFailed, formatter } from './formatter' +export { + FormatterRewrote, + JobEnded, + JobOutput, + JobStarted, + SnapshotFailed, +} from './events' export type { Formatter, FormatterOptions, FormatterProject } from './formatter' diff --git a/packages/ai-harness/src/first-party/coding/read.ts b/packages/ai-harness/src/first-party/coding/read.ts index 8bc9e392fa..d8749376c9 100644 --- a/packages/ai-harness/src/first-party/coding/read.ts +++ b/packages/ai-harness/src/first-party/coding/read.ts @@ -13,9 +13,11 @@ const MAX_PAGE_BYTES = 50 * 1024 const MAX_LINE_CHARS = 2000 // A NUL byte in this many first bytes marks a binary file. const BINARY_CHECK_BYTES = 8 * 1024 -// ponytail: media goes to the model as inline base64, and it stays in the -// transcript. Store it as harness media if big files make turns too large. -const MAX_MEDIA_BYTES = 5 * 1024 * 1024 +// Media goes to the model as inline base64, and it stays in the transcript, +// so every model call after it sends it again. 256 KiB keeps that small. +// ponytail: a plugin cannot put a file in the media store of the session +// yet. With that, big media could go there as a reference. +const MAX_MEDIA_BYTES = 256 * 1024 const MAX_SUGGESTIONS = 3 // The first bytes of the files that go to the model as media parts. The @@ -108,7 +110,7 @@ function linesOf(text: string, offset: number, limit: number) { * `limit` and the file goes on, a note says how to read the next page. */ function page( - { lines: window, total }: { lines: Array; total: number }, + { lines: window, total }: { lines: Array; total?: number }, options: { path: string; offset: number; limit: number | undefined }, ) { const { path, offset, limit } = options @@ -122,12 +124,16 @@ function page( shown.push(row) } if (shown.length === 0) { - return `[No lines from line ${offset}. ${path} has ${total} lines.]` + const count = total === undefined ? '' : ` ${path} has ${total} lines.` + return `[No lines from line ${offset}.${count}]` } const last = offset + shown.length - 1 - const isCut = shown.length < (limit ?? Infinity) && last < total + // Without `total`, the file has more lines after the page. + const goesOn = total === undefined || last < total + const isCut = shown.length < (limit ?? Infinity) && goesOn if (!isCut) return shown.join('\n') - return `${shown.join('\n')}\n\n[Showing lines ${offset}-${last} of ${total}. Read more with offset ${last + 1}.]` + const of = total === undefined ? '' : ` of ${total}` + return `${shown.join('\n')}\n\n[Showing lines ${offset}-${last}${of}. Read more with offset ${last + 1}.]` } /** @@ -141,7 +147,7 @@ export function readTools(env: ToolEnv) { toolDefinition({ name: 'read_file', description: - 'Read a file in the workspace. Text comes back as lines numbered from 1, at most 2000 lines or 50 KiB a page. Use offset and limit to read other lines. Images (PNG, JPEG, GIF, WebP) and PDFs up to 5 MB come back as media.', + 'Read a file in the workspace. Text comes back as lines numbered from 1, at most 2000 lines or 50 KiB a page. Use offset and limit to read other lines. Images (PNG, JPEG, GIF, WebP) and PDFs up to 256 KB come back as media.', inputSchema: { type: 'object', properties: { @@ -184,9 +190,9 @@ export function readTools(env: ToolEnv) { const media = MEDIA.find((item) => item.signature.test(start)) if (media) { if (size > MAX_MEDIA_BYTES) { - const mb = (size / 1024 / 1024).toFixed(1) + const kb = Math.ceil(size / 1024) throw new Error( - `${path} is ${mb} MB. read_file sends images and PDFs up to 5 MB.`, + `${path} is ${kb} KB. read_file sends images and PDFs up to 256 KB. Make a smaller copy, for example with bash, and read that.`, ) } const bytes = whole ?? (await backend.readFile(full)) diff --git a/packages/ai-harness/src/first-party/coding/search.ts b/packages/ai-harness/src/first-party/coding/search.ts index f9c52f3c54..56873d9278 100644 --- a/packages/ai-harness/src/first-party/coding/search.ts +++ b/packages/ai-harness/src/first-party/coding/search.ts @@ -361,32 +361,53 @@ export function searchTools(env: ToolEnv) { /** * The hits for `pattern` under `base`, in the files that `isWanted` keeps. - * `listed` is the result of `filesUnder(base)`, when the caller has it. + * With `withFiles`, also every file under `base`, like `filesUnder(base)`. + * The files are listed once: the JS engines search the list, and rg lists + * and searches in one command. */ async function grepUnder( base: string, pattern: string, isWanted: (path: string) => boolean, - listed?: Awaited>, + withFiles: boolean, ) { // A pattern that is not a valid JS regular expression throws here, // whichever engine runs. const regex = new RegExp(pattern) const command = await findRgOnce() if (command === undefined) { - const { files, partial } = listed ?? (await filesUnder(base)) + const { files, partial } = await filesUnder(base) const hits = await grepFiles(backend, base, files.filter(isWanted), regex) - return { hits, partial } + return { hits, partial, files } + } + const search = `${command} ${RG_GREP} ${rgCommon} -e ${quoteArg(pattern, platform)} .` + const hitsOf = (out: string) => + rgHits(out) + .filter((hit) => isWanted(hit.path)) + .sort(byPath) + if (!withFiles) { + const hits = hitsOf(await runRg(backend, search, base)) + return { hits, partial: false, files: undefined } } + // One shell runs the list and the search, with a line between them that + // no file name has. The exit code is the one of the search. + const marker = `-- ${crypto.randomUUID()} --` + const then = platform === 'win32' ? '&' : ';' const out = await runRg( backend, - `${command} ${RG_GREP} ${rgCommon} -e ${quoteArg(pattern, platform)} .`, + `${command} --files ${rgCommon} . ${then} echo ${marker} ${then} ${search}`, base, ) - const hits = rgHits(out) - .filter((hit) => isWanted(hit.path)) - .sort(byPath) - return { hits, partial: false } + const at = out.indexOf(marker) + if (at === -1) throw new Error('The search failed: rg gave no file list.') + const listing = out.slice(0, at) + // `cmd.exe` keeps the space before `&` in the line of `echo`. + const found = out.slice(at + marker.length).replace(/^[ \t]*\r?\n/, '') + return { + hits: hitsOf(found), + partial: false, + files: tidy(listing.split('\0')), + } } return [ @@ -452,15 +473,19 @@ export function searchTools(env: ToolEnv) { // grep must not show the lines of a file that read_file could not // read without a question. const isProtected = await env.protectedIn?.(base) - // The note counts the protected files in the list, not in the hits, - // so it does not tell whether a protected file has the pattern. - const listed = isProtected ? await filesUnder(base) : undefined - const skipped = - listed?.files.filter((path) => inGlob(path) && isProtected?.(path)) - .length ?? 0 const isWanted = (path: string) => inGlob(path) && !(isProtected?.(path) ?? false) - const { hits, partial } = await grepUnder(base, pattern, isWanted, listed) + const { + hits, + partial, + files: listed, + } = await grepUnder(base, pattern, isWanted, isProtected !== undefined) + // The note counts the protected files in the list, not in the hits, + // so it does not tell whether a protected file has the pattern. + const skipped = isProtected + ? (listed ?? []).filter((path) => inGlob(path) && isProtected(path)) + .length + : 0 const shown = hits.slice(0, MAX_MATCHES).map((hit) => { const text = hit.text.slice(0, MAX_LINE).trim() return `${env.shown(resolve(base, hit.path))}:${hit.line}: ${text}` diff --git a/packages/ai-harness/src/first-party/coding/snapshots.ts b/packages/ai-harness/src/first-party/coding/snapshots.ts index 1491a0e1a4..ebe2a1f531 100644 --- a/packages/ai-harness/src/first-party/coding/snapshots.ts +++ b/packages/ai-harness/src/first-party/coding/snapshots.ts @@ -5,6 +5,7 @@ import { RevertFiles, RevertStanding, definePlugin } from '../../plugins' import { isRecord } from '../../utils' import { WorkspaceHooks } from '../workspace-hooks' import { hostBackend, pathsOf } from './backend' +import { SnapshotFailed } from './events' import { parsePatch } from './patch' import { GIT, GIT_ENV, quoteArg, shellPlatform } from './search' import { createLock, runOf, writingRun } from './workspace' @@ -387,6 +388,8 @@ export function purgeSnapshots( * `to`, it compares with the files now. * * Needs the git CLI on the backend. Without git, the plugin does nothing. + * A snapshot that fails does not fail the turn: the plugin sends a + * `tanstack/snapshots:failed` event (`SnapshotFailed`) and goes on. * * @example * ```ts @@ -402,23 +405,21 @@ export function snapshots(options: SnapshotsOptions) { const paths = pathsOf(backend) // The repositories of the open sessions, for `diff`. const repos = new Set() - let found: Promise | undefined - /** Is git on the backend? It asks once, and warns once when it is not. */ - const hasGit = () => + let found: Promise | undefined + /** + * Why git cannot run on the backend, or `undefined` when it can. It asks + * once. + */ + const gitMissing = () => (found ??= backend .exec(`${GIT} --version`, { env: GIT_ENV, timeoutMs: TIMEOUT_MS }) .then( - (probe) => { - const isFound = probe.exitCode === 0 - if (!isFound) { - console.warn('snapshots: git was not found. Snapshots are off.') - } - return isFound - }, - (error: unknown) => { - console.warn(`snapshots: ${messageOf(error)}. Snapshots are off.`) - return false - }, + (probe) => + probe.exitCode === 0 + ? undefined + : 'snapshots: git was not found. Snapshots are off.', + (error: unknown) => + `snapshots: ${messageOf(error)}. Snapshots are off.`, )) /** `full` as a path in the workspace with `/`, or `undefined` outside it. */ @@ -435,7 +436,9 @@ export function snapshots(options: SnapshotsOptions) { return definePlugin({ name: 'tanstack/snapshots', diff: async (from: string, to?: string) => { - if (!(await hasGit())) throw new Error('snapshots: git is not available.') + if ((await gitMissing()) !== undefined) { + throw new Error('snapshots: git is not available.') + } // ponytail: each thread has its own repository, so `diff` asks the // repositories of the open sessions, one by one, which has `from`. const open = [...repos] @@ -450,16 +453,16 @@ export function snapshots(options: SnapshotsOptions) { tenantId: ctx.session.principal?.tenantId, threadId: ctx.session.threadId, }) - const isReady = - (await hasGit()) && + // Why snapshots are off in this session, or `undefined`. + const off = + (await gitMissing()) ?? (await repo.init().then( - () => true, - (error: unknown) => { - console.warn(`snapshots: ${messageOf(error)}. Snapshots are off.`) - return false - }, + () => undefined, + (error: unknown) => + `snapshots: ${messageOf(error)}. Snapshots are off.`, )) - if (!isReady) { + if (off !== undefined) { + ctx.emit(SnapshotFailed, { message: off }) // Without git, a revert changes only the transcript. ctx.provide(RevertFiles, { revert: async () => undefined, @@ -526,12 +529,15 @@ export function snapshots(options: SnapshotsOptions) { } | undefined - // A failed snapshot must not fail the turn. It only says so. + // A failed snapshot must not fail the turn. It only says so, with a + // `tanstack/snapshots:failed` event. const guard = async (fn: () => Promise) => { try { return await fn() } catch (error) { - console.warn(`snapshots: ${messageOf(error)}`) + ctx.emit(SnapshotFailed, { + message: `snapshots: ${messageOf(error)}`, + }) return undefined } } diff --git a/packages/ai-harness/src/first-party/coding/web.ts b/packages/ai-harness/src/first-party/coding/web.ts index cb45d92cc4..a895232c42 100644 --- a/packages/ai-harness/src/first-party/coding/web.ts +++ b/packages/ai-harness/src/first-party/coding/web.ts @@ -4,6 +4,7 @@ import * as https from 'node:https' import { BlockList, isIP } from 'node:net' import { toolDefinition } from '@tanstack/ai' import { isRecord } from '../../utils' +import { boundText } from '../bound-output' import { optionalString, stringArg } from './backend' import type { IncomingMessage } from 'node:http' import type { LookupFunction } from 'node:net' @@ -12,6 +13,11 @@ const DEFAULT_TIMEOUT_MS = 30_000 const MAX_TIMEOUT_MS = 120_000 /** `webfetch` stops reading a body after this many bytes. */ const MAX_BYTES = 5 * 1024 * 1024 +/** + * The model gets at most this many bytes of a page: its start. The result + * stays in the transcript, so each model call after it sends it again. + */ +const MAX_RESULT_BYTES = 256 * 1024 const MAX_REDIRECTS = 5 /** A NUL byte in this many first bytes marks a binary body. */ const BINARY_CHECK_BYTES = 8 * 1024 @@ -353,7 +359,13 @@ async function readPage( // content type if other charsets matter. const text = new TextDecoder().decode(bytes) const page = isHtml ? await convert(text, format) : text - return isCut ? `${page}\n\n[Cut at 5 MB. The rest was not read.]` : page + const kept = boundText(page, { + maxLines: Number.MAX_SAFE_INTEGER, + maxBytes: MAX_RESULT_BYTES, + }) + if (!kept.truncated && !isCut) return page + const notRead = isCut ? ' The page is over 5 MB. The rest was not read.' : '' + return `${kept.text}\n[webfetch gives at most the first 256 KB of a page.${notRead} To see more, save the page to a file with bash, then use grep or read_file with an offset.]` } /** @@ -361,9 +373,10 @@ async function readPage( * * - `webfetch` reads a URL. HTML comes back as Markdown when the optional * `turndown` package is installed, else as text. Text and JSON come back - * as they are. Other content types are refused. It reads at most 5 MB, - * stops after 30 seconds (the model can ask for up to 120), and follows - * at most 5 redirects. + * as they are. Other content types are refused. The model gets the + * first 256 KB of the page, with a note when there is more. It reads at + * most 5 MB, stops after 30 seconds (the model can ask for up to 120), + * and follows at most 5 redirects. * - `websearch` searches with `options.search`. Without a provider, there * is no `websearch` tool. * @@ -386,7 +399,7 @@ export function webTools(options: WebToolsOptions = {}) { const webfetch = toolDefinition({ name: 'webfetch', description: - 'Read a web page by its http or https URL. HTML comes back as Markdown by default, or set format to text or html. Text and JSON come back as they are. At most 5 MB is read.', + 'Read a web page by its http or https URL. HTML comes back as Markdown by default, or set format to text or html. Text and JSON come back as they are. You get at most the first 256 KB of the page.', inputSchema: { type: 'object', properties: { diff --git a/packages/ai-harness/src/first-party/coding/workspace.ts b/packages/ai-harness/src/first-party/coding/workspace.ts index 7bb09601a8..24aed4f10f 100644 --- a/packages/ai-harness/src/first-party/coding/workspace.ts +++ b/packages/ai-harness/src/first-party/coding/workspace.ts @@ -16,6 +16,7 @@ import { readTools } from './read' import { searchTools } from './search' import { webTools } from './web' import type { AnyTool, ToolExecutionContext } from '@tanstack/ai' +import type { PluginEvent } from '../../extensions' import type { RunContext } from '../../types' import type { PermissionDecision, ToolResources } from '../permissions' import type { BoundWorkspace, ToolEnv, WorkspaceBackend } from './backend' @@ -112,6 +113,8 @@ interface SessionAccess { note?: (text: string) => Promise signal?: AbortSignal jobs?: { started: (jobId: string) => void; ended: (jobId: string) => void } + /** Sends the events of the background `bash` jobs. */ + emit?: (event: PluginEvent, value: T) => void } /** A Windows network path (`\\server\share`). Not `\\?\` or `\\.\`. */ @@ -438,6 +441,7 @@ function bindTools( spillDir: options.spillDir, signal: session.signal, jobs: session.jobs, + emit: session.emit, maxJobs: options.maxJobs, }), ...(options.web === false ? [] : webTools(options.web)), @@ -724,6 +728,8 @@ export function workspaceTools(options: WorkspaceToolsOptions) { // keeps them. Without a session, the plugin cleanup kills them. signal: ctx.getOptional(SessionSignal) ?? ctx.resources.signal, jobs: ctx.getOptional(BackgroundJobs), + // The events of the background jobs go to the clients. + emit: ctx.emit, }) return { tools, diff --git a/packages/ai-harness/tests/coding-backend.test.ts b/packages/ai-harness/tests/coding-backend.test.ts index 81359fdb75..22e8b26401 100644 --- a/packages/ai-harness/tests/coding-backend.test.ts +++ b/packages/ai-harness/tests/coding-backend.test.ts @@ -207,13 +207,17 @@ describe('hostBackend', () => { }) describe('hostBackend.readLines', () => { - /** The page and the line count of a whole-file read, as `read_file` makes them. */ + /** + * The page of a whole-file read, as `read_file` makes it, and the line + * count when the page gets to the end of the file. + */ const expected = (text: string, offset: number, limit: number) => { const all = text.replace(/^/, '').split('\n') if (text.endsWith('\n')) all.pop() + const isLast = offset - 1 + limit >= all.length return { lines: all.slice(offset - 1, offset - 1 + limit), - total: all.length, + total: isLast ? all.length : undefined, } } // 1 MiB is the read size: `é` is split over the edge of the first read. @@ -239,12 +243,44 @@ describe('hostBackend.readLines', () => { expect(read.head).toEqual(Buffer.from(text).subarray(0, 8 * 1024)) }) - it('stops after 64 KiB of lines, and still counts every line', async () => { + it('stops after 64 KiB of lines, and does not count the lines after it', async () => { const file = join(dir, 'long.txt') await writeFile(file, `${'y'.repeat(40 * 1024)}\n`.repeat(4)) const read = await hostBackend.readLines(file, { offset: 1, limit: 10 }) expect(read.lines).toHaveLength(2) - expect(read.total).toBe(4) + expect(read.total).toBeUndefined() + }) + + /** A file of 300,000 lines of 10 bytes: 3 MB, more than one read. */ + async function manyLines() { + const file = join(dir, 'many.txt') + const lines = Array.from( + { length: 300_000 }, + (_, index) => `line${String(index + 1).padStart(6, '0')}`, + ) + await writeFile(file, `${lines.join('\n')}\n`) + return file + } + + it('stops after the first page of a big file, so the line count is not known', async () => { + const read = await hostBackend.readLines(await manyLines(), { + offset: 1, + limit: 3, + }) + expect(read).toMatchObject({ + size: 3_300_000, + lines: ['line000001', 'line000002', 'line000003'], + }) + expect(read.total).toBeUndefined() + }) + + it('counts the lines when the page gets to the end of a big file', async () => { + const read = await hostBackend.readLines(await manyLines(), { + offset: 299_999, + limit: 5, + }) + expect(read.lines).toEqual(['line299999', 'line300000']) + expect(read.total).toBe(300_000) }) }) diff --git a/packages/ai-harness/tests/coding-formatter.test.ts b/packages/ai-harness/tests/coding-formatter.test.ts index 0688ef0c4f..471941517f 100644 --- a/packages/ai-harness/tests/coding-formatter.test.ts +++ b/packages/ai-harness/tests/coding-formatter.test.ts @@ -10,6 +10,7 @@ import { formatOnWrite, formatter, } from '../src/first-party/coding/formatter' +import { FormatterRewrote } from '../src/first-party/coding/events' import { mockAdapter, text, toolCall } from './helpers' import type { HarnessSession } from '../src' import type { WorkspaceBackend } from '../src/first-party/coding/backend' @@ -84,8 +85,9 @@ async function formatIn( ) { const { backend, runs } = project(files, exec) const failures: Array = [] - const afterWrite = formatOnWrite({ root, backend, ...options }, (failure) => - failures.push(failure), + const afterWrite = formatOnWrite( + { root, backend, ...options }, + { onFailure: (failure) => failures.push(failure) }, ) await afterWrite(join(root, path)) return { runs, commands: runs.map((run) => run.command), failures } @@ -201,7 +203,7 @@ describe('formatOnWrite', () => { const { backend, runs } = project(prettier, ok, 'sh') const afterWrite = formatOnWrite( { root: '/workspace', backend }, - () => undefined, + { onFailure: () => undefined }, ) await afterWrite('/workspace/src/a.ts') expect(runs).toEqual([ @@ -227,8 +229,9 @@ describe('formatOnWrite', () => { return readdir(dir) } const failures: Array = [] - const afterWrite = formatOnWrite({ root, backend }, (failure) => - failures.push(failure), + const afterWrite = formatOnWrite( + { root, backend }, + { onFailure: (failure) => failures.push(failure) }, ) await afterWrite(join(root, 'a.ts')) fail = false @@ -239,7 +242,10 @@ describe('formatOnWrite', () => { it('checks the config files once, at the first write', async () => { const { backend, runs, put } = project({}) - const afterWrite = formatOnWrite({ root, backend }, () => undefined) + const afterWrite = formatOnWrite( + { root, backend }, + { onFailure: () => undefined }, + ) await afterWrite(join(root, 'a.ts')) put('.prettierrc', '{}') await afterWrite(join(root, 'a.ts')) @@ -273,6 +279,34 @@ describe('formatOnWrite', () => { expect(failures).toEqual([{ path: join(root, 'a.ts'), message }]) }) + it.each([ + [ + 'changed', + 'let a = 1\n', + [{ path: join(root, 'a.ts'), formatter: 'prettier' }], + ], + ['kept', 'let a=1', []], + ])( + 'reports a rewrite only when the formatter changed the file (%s)', + async (_, formatted, expected) => { + // The fake prettier writes `formatted` to the file. + const files = project({ ...prettier, 'a.ts': 'let a=1' }, () => { + files.put('a.ts', formatted) + return { exitCode: 0, stdout: '', stderr: '' } + }) + const rewrites: Array<{ path: string; formatter: string }> = [] + const afterWrite = formatOnWrite( + { root, backend: files.backend }, + { + onFailure: () => undefined, + onRewrite: (rewrite) => rewrites.push(rewrite), + }, + ) + await afterWrite(join(root, 'a.ts')) + expect(rewrites).toEqual(expected) + }, + ) + it('reports a backend that throws, and does not throw', async () => { const { failures } = await formatIn(prettier, 'a.ts', {}, () => { throw new Error('spawn sh ENOENT') @@ -336,6 +370,48 @@ describe('formatter', () => { }) }) +describe('formatter events', () => { + it('sends tanstack/formatter:rewrote when the formatter changed the written file', async () => { + // The fake prettier puts spaces around `=`. + const files = project(prettier, () => { + files.put('a.ts', files.read('a.ts').replace('=', ' = ')) + return { exitCode: 0, stdout: '', stderr: '' } + }) + const rewrites: Array = [] + const listener = definePlugin({ + name: 'test/listener', + setup: (ctx) => { + ctx.on(FormatterRewrote, (rewrite) => rewrites.push(rewrite)) + }, + }) + const { adapter } = mockAdapter([ + () => toolCall('write_file', { path: 'a.ts', content: 'let a=1' }, 'c1'), + () => text('done'), + ]) + const host = createHarnessHost({ persistence: memoryPersistence() }) + const session = await host.open( + defineHarness({ + name: 'test/formatter-events', + adapter, + plugins: () => [ + workspaceTools({ root, backend: files.backend }), + formatter({ root, backend: files.backend }), + listener, + ], + }), + { threadId: 't' }, + ) + + await session.send('work') + await host.close() + + expect(files.read('a.ts')).toBe('let a = 1') + expect(rewrites).toEqual([ + { path: join(root, 'a.ts'), formatter: 'prettier' }, + ]) + }) +}) + describe('formatter with permissions()', () => { /** * A session with the workspace tools, the formatter, and `permissions()` diff --git a/packages/ai-harness/tests/coding-jobs.test.ts b/packages/ai-harness/tests/coding-jobs.test.ts index 87076eb0da..864fa96899 100644 --- a/packages/ai-harness/tests/coding-jobs.test.ts +++ b/packages/ai-harness/tests/coding-jobs.test.ts @@ -1,10 +1,11 @@ import { getEventListeners } from 'node:events' import { tmpdir } from 'node:os' import { resolve } from 'node:path' -import { describe, expect, it } from 'vitest' +import { afterEach, describe, expect, it, vi } from 'vitest' import { bashTools } from '../src/first-party/coding/bash' import { createOutputTail } from '../src/first-party/coding/ring-buffer' import type { ToolExecutionContext } from '@tanstack/ai' +import type { PluginEvent } from '../src/extensions' import type { ToolEnv, WorkspaceBackend, @@ -240,3 +241,108 @@ describe('bash in the foreground', () => { ).rejects.toThrow('the limit is 1') }) }) + +describe('bash job events', () => { + afterEach(() => { + vi.useRealTimers() + }) + + /** + * A workspace whose one background job gives its output to `write`, and + * ends with `end`. `events` has each event that the tools send. + */ + function streaming() { + vi.useFakeTimers() + const started = runningJob() + let onOutput: (text: string) => void = () => undefined + const { env } = workspace({ + spawn: (_command, options) => { + onOutput = options?.onOutput ?? onOutput + return started.job + }, + }) + const events: Array<{ name: string; value: unknown }> = [] + const emit = (event: PluginEvent, value: T) => { + events.push({ name: event.name, value }) + } + const session = new AbortController() + const bash = bashTool(env, { emit, signal: session.signal }) + return { + start: () => bash({ command: 'npm run dev', background: true }), + write: (text: string) => onOutput(text), + end: started.end, + close: () => session.abort(), + events, + /** The output events, as their text. */ + outputs: () => + events + .filter((event) => event.name === 'tanstack/workspace-tools:job-output') + .map((event) => event.value), + } + } + + it('sends started, then the output, then ended', async () => { + const job = streaming() + await job.start() + job.write('ready\n') + job.end(0) + await vi.advanceTimersByTimeAsync(0) + expect(job.events).toEqual([ + { + name: 'tanstack/workspace-tools:job-started', + value: { jobId: 'bash-1', command: 'npm run dev' }, + }, + { + name: 'tanstack/workspace-tools:job-output', + value: { jobId: 'bash-1', output: 'ready\n' }, + }, + { name: 'tanstack/workspace-tools:job-ended', value: { jobId: 'bash-1', exitCode: 0 } }, + ]) + }) + + it('sends at most one output event each 250 ms, with the new output', async () => { + const job = streaming() + await job.start() + job.write('a') + await vi.advanceTimersByTimeAsync(0) + job.write('b') + job.write('c') + await vi.advanceTimersByTimeAsync(249) + expect(job.outputs()).toEqual([{ jobId: 'bash-1', output: 'a' }]) + await vi.advanceTimersByTimeAsync(1) + expect(job.outputs()).toEqual([ + { jobId: 'bash-1', output: 'a' }, + { jobId: 'bash-1', output: 'bc' }, + ]) + }) + + it('keeps the newest 64 KiB in one output event', async () => { + const job = streaming() + await job.start() + job.write('x'.repeat(70 * 1024)) + await vi.advanceTimersByTimeAsync(0) + expect(job.outputs()).toEqual([ + { + jobId: 'bash-1', + output: `[6144 earlier bytes were dropped.]\n${'x'.repeat(64 * 1024)}`, + }, + ]) + }) + + it('sends no output or ended event after the session closed', async () => { + const job = streaming() + await job.start() + job.write('early') + await vi.advanceTimersByTimeAsync(0) + // This output waits for the next 250 ms. + job.write('late') + // Closing the session kills the job, so it ends with 143. + job.close() + await vi.advanceTimersByTimeAsync(1000) + expect(job.events.map((event) => event.name)).toEqual([ + 'tanstack/workspace-tools:job-started', + 'tanstack/workspace-tools:job-output', + ]) + expect(job.outputs()).toEqual([{ jobId: 'bash-1', output: 'early' }]) + }) +}) diff --git a/packages/ai-harness/tests/coding-read.test.ts b/packages/ai-harness/tests/coding-read.test.ts index e34e898e30..bbeb5d5eb4 100644 --- a/packages/ai-harness/tests/coding-read.test.ts +++ b/packages/ai-harness/tests/coding-read.test.ts @@ -158,6 +158,24 @@ describe('read_file pages', () => { ) }) + it('says where the next page starts when the backend did not count the lines', async () => { + // `readLines` stopped after the page, so the number of lines is not known. + const paged: WorkspaceBackend = { + ...backend, + readLines: async () => ({ + head: ascii('a\nb\n'), + size: 1000, + lines: ['a', 'b'], + total: undefined, + }), + } + const readPaged = reader({ workspace: { root, backend: paged } }) + await put('paged.txt', 'a\nb\n') + expect(await readPaged('paged.txt')).toBe( + '1\ta\n2\tb\n\n[Showing lines 1-2. Read more with offset 3.]', + ) + }) + it('says how many lines there are when offset is past the end', async () => { expect(await read('many.txt', { offset: 3000 })).toBe( '[No lines from line 3000. many.txt has 2500 lines.]', @@ -220,15 +238,32 @@ describe('read_file media and binary files', () => { expect(await read(file)).toEqual([part]) }) - it('refuses an image over 5 MB', async () => { - const big = new Uint8Array(6 * 1024 * 1024) + it('refuses an image over 256 KB, and says how to read it', async () => { + const big = new Uint8Array(300 * 1024) big.set(PNG) await put('big.png', big) await expect(read('big.png')).rejects.toThrow( - 'big.png is 6.0 MB. read_file sends images and PDFs up to 5 MB.', + 'big.png is 300 KB. read_file sends images and PDFs up to 256 KB. Make a smaller copy, for example with bash, and read that.', ) }) + it('sends an image of 256 KB', async () => { + const image = new Uint8Array(256 * 1024) + image.set(PNG) + await put('edge.png', image) + const result = await read('edge.png') + expect(result).toEqual([ + { + type: 'image', + source: { + type: 'data', + value: Buffer.from(image).toString('base64'), + mimeType: 'image/png', + }, + }, + ]) + }) + it.each([ { file: 'data.bin', diff --git a/packages/ai-harness/tests/coding-search.test.ts b/packages/ai-harness/tests/coding-search.test.ts index e3706b6ba7..6a38d8d426 100644 --- a/packages/ai-harness/tests/coding-search.test.ts +++ b/packages/ai-harness/tests/coding-search.test.ts @@ -286,15 +286,19 @@ describe('list_files and grep engines', () => { ) }) - it('rg: grep drops the lines of a protected file, and lists the files to count them', async () => { + it('rg: grep drops the lines of a protected file, and lists the files to count them in the same command', async () => { + // One shell runs the list, an `echo` of a marker line, and the search, + // like `sh` does with `;`. + const listAndSearch = + /^rg --files .* \. ; echo (-- [\w-]+ --) ; rg --line-number .* \.$/ const fake = fakeBackend( { '.env': 'two=1\n', 'a.ts': 'two\n' }, (command) => { if (command === 'rg --version') return ok('ripgrep 14.1.0') - if (command.startsWith('rg --files ')) return ok('./.env\0./a.ts\0') - return command.startsWith('rg --line-number ') - ? ok('./.env\x001:two=1\n./a.ts\x001:two\n') - : undefined + const marker = listAndSearch.exec(command)?.[1] + if (marker === undefined) return undefined + const hits = './.env\x001:two=1\n./a.ts\x001:two\n' + return ok(`./.env\0./a.ts\0${marker}\n${hits}`) }, ) // Rules that ask about every .env file, like permissions() does. @@ -310,10 +314,30 @@ describe('list_files and grep engines', () => { ) expect(fake.commands).toEqual([ 'rg --version', - expect.stringMatching(/^rg --files /), - expect.stringMatching(/^rg --line-number /), + expect.stringMatching(listAndSearch), ]) }) + + it('git: grep with protected files lists the files once', async () => { + const fake = fakeBackend( + { '.env': 'two=1\n', 'a.ts': 'two\n' }, + (command) => (command === GIT ? ok('.env\0a.ts\0') : undefined), + ) + const readdir = vi.spyOn(fake.backend, 'readdir') + const { grep } = searchOn(fake.backend, root, { + ask: async () => 'n', + mode: () => 'default', + rules: () => (_tool, paths) => + paths.some((path) => path.endsWith('.env')) ? 'ask' : 'allow', + }) + + expect(await grep({ pattern: 'two' })).toBe( + 'a.ts:1: two\n[Skipped 1 file that the permission rules protect.]', + ) + expect(fake.commands.filter((command) => command === GIT)).toHaveLength(1) + // The link check of the git list reads the root folder once. + expect(readdir.mock.calls).toEqual([[root]]) + }) }) describe('quoting for the shell of the backend', () => { diff --git a/packages/ai-harness/tests/coding-snapshots.test.ts b/packages/ai-harness/tests/coding-snapshots.test.ts index f81cca1b4d..14c8a3733e 100644 --- a/packages/ai-harness/tests/coding-snapshots.test.ts +++ b/packages/ai-harness/tests/coding-snapshots.test.ts @@ -11,7 +11,8 @@ import { dirname, join } from 'node:path' import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest' import { toolDefinition } from '@tanstack/ai' import { memoryLogStore, memoryPersistence } from '@tanstack/ai-persistence' -import { createHarnessHost, defineHarness, logIdOf } from '../src' +import { createHarnessHost, defineHarness, definePlugin, logIdOf } from '../src' +import { SnapshotFailed } from '../src/first-party/coding/events' import { hostBackend } from '../src/first-party/coding/backend' import { quoteArg } from '../src/first-party/coding/search' import { purgeSnapshots, snapshots } from '../src/first-party/coding/snapshots' @@ -89,6 +90,15 @@ async function open(options: { return 'changed' }) const plugin = snapshots({ root, dataDir, backend: options.backend }) + // The messages of the `tanstack/snapshots:failed` events. The listener comes + // first, so it also gets an event that `snapshots()` sends in its setup. + const failures: Array = [] + const listener = definePlugin({ + name: 'test/snapshot-failures', + setup: (ctx) => { + ctx.on(SnapshotFailed, ({ message }) => failures.push(message)) + }, + }) const { adapter, calls } = mockAdapter(options.replies) const host = createHarnessHost({ persistence: options.persistence ?? memoryPersistence(), @@ -100,15 +110,15 @@ async function open(options: { tools: [change], plugins: () => options.workspace - ? [workspaceTools({ root, web: false }), plugin] - : [plugin], + ? [listener, workspaceTools({ root, web: false }), plugin] + : [listener, plugin], }), { threadId: options.threadId ?? THREAD, principal: options.principal, }, ) - return { host, session, plugin, calls } + return { host, session, plugin, calls, failures } } /** The steps that the plugin state keeps. */ @@ -916,8 +926,7 @@ describe('snapshots with a fake git', () => { }, 60_000) it('tries again when another git process holds the index lock', async () => { - const warn = vi.spyOn(console, 'warn').mockImplementation(() => {}) - const { host, session } = await open({ + const { host, session, failures } = await open({ replies: editTurn('Done.'), backend: fakeGit({ locked: 2 }), }) @@ -925,18 +934,12 @@ describe('snapshots with a fake git', () => { await session.send('Edit a.txt') expect(savedSteps(session)).toMatchObject([{ files: ['a.txt'] }]) - // Only the snapshot warnings: the host can warn about a handle of an - // earlier test that the garbage collector found. - const snapshotWarnings = warn.mock.calls.filter(([message]) => - String(message).startsWith('snapshots:'), - ) - expect(snapshotWarnings).toEqual([]) + expect(failures).toEqual([]) await host.close() }) - it('reports an index lock that stays, and the turn still runs', async () => { - const warn = vi.spyOn(console, 'warn').mockImplementation(() => {}) - const { host, session } = await open({ + it('reports an index lock that stays with an event, and the turn still runs', async () => { + const { host, session, failures } = await open({ replies: [() => text('Hi.')], backend: fakeGit({ locked: Infinity }), }) @@ -947,20 +950,19 @@ describe('snapshots with a fake git', () => { { role: 'user', content: 'Hello' }, { role: 'assistant', content: 'Hi.' }, ]) - expect(warn).toHaveBeenCalledWith( + expect(failures).toEqual([ expect.stringMatching( /^snapshots: another git process holds the index lock of .+: fatal: Unable to create/, ), - ) + ]) await host.close() }) }) describe('snapshots without git', () => { - it('adds no commands, warns once, and the turn still runs', async () => { - const warn = vi.spyOn(console, 'warn').mockImplementation(() => {}) + it('adds no commands, sends one failed event, and the turn still runs', async () => { // Every command fails, like a machine with no git. - const { host, session, plugin } = await open({ + const { host, session, plugin, failures } = await open({ replies: [() => text('Hi.')], backend: fakeBackend(async () => ({ exitCode: 127, @@ -976,8 +978,8 @@ describe('snapshots without git', () => { { role: 'user', content: 'Hello' }, { role: 'assistant', content: 'Hi.' }, ]) - expect(warn.mock.calls).toEqual([ - ['snapshots: git was not found. Snapshots are off.'], + expect(failures).toEqual([ + 'snapshots: git was not found. Snapshots are off.', ]) await expect(plugin.diff('abc')).rejects.toThrow( 'snapshots: git is not available.', diff --git a/packages/ai-harness/tests/coding-web.test.ts b/packages/ai-harness/tests/coding-web.test.ts index 547b343ee5..acfbe489d7 100644 --- a/packages/ai-harness/tests/coding-web.test.ts +++ b/packages/ai-harness/tests/coding-web.test.ts @@ -40,6 +40,8 @@ function resolves(host: string, ...answers: Array>) { const PUBLIC_IP = '93.184.216.34' const MB5 = 5 * 1024 * 1024 +/** The most of a page that the model gets. */ +const KB256 = 256 * 1024 const REFUSED = 'webfetch does not read private or local hosts' type Route = (init: RequestInit | undefined) => Response | Promise @@ -142,14 +144,30 @@ describe('webfetch formats', () => { }) describe('webfetch limits', () => { - it('stops reading at 5 MB and says the page was cut', async () => { - const body = 'a'.repeat(MB5 + 100) - const result = await webfetch( + /** Fetch a text page of `body`. */ + const fetchText = (body: string) => + webfetch( { url: 'https://example.com/big.txt' }, { 'https://example.com/big.txt': page(body, 'text/plain') }, ) + const MORE = + 'To see more, save the page to a file with bash, then use grep or read_file with an offset.]' + + it('gives the first 256 KB of a 1 MB page, and says how much was cut', async () => { + const result = await fetchText('a'.repeat(1024 * 1024)) + expect(result).toBe( + `${'a'.repeat(KB256)}\n\n[Output cut. Showing the first 1 of 1 lines, 262144 of 1048576 bytes.]\n[webfetch gives at most the first 256 KB of a page. ${MORE}`, + ) + }) + + it('gives a page of 256 KB as it is', async () => { + expect(await fetchText('a'.repeat(KB256))).toBe('a'.repeat(KB256)) + }) + + it('stops reading at 5 MB and says the page was cut', async () => { + const result = await fetchText('a'.repeat(MB5 + 100)) expect(result).toBe( - `${'a'.repeat(MB5)}\n\n[Cut at 5 MB. The rest was not read.]`, + `${'a'.repeat(KB256)}\n\n[Output cut. Showing the first 1 of 1 lines, 262144 of 5242880 bytes.]\n[webfetch gives at most the first 256 KB of a page. The page is over 5 MB. The rest was not read. ${MORE}`, ) }) diff --git a/packages/ai-harness/tests/workspace-tools.test.ts b/packages/ai-harness/tests/workspace-tools.test.ts index eca8ff8064..e5e6a1680c 100644 --- a/packages/ai-harness/tests/workspace-tools.test.ts +++ b/packages/ai-harness/tests/workspace-tools.test.ts @@ -3,9 +3,14 @@ import { tmpdir } from 'node:os' import * as nodePath from 'node:path' import { afterAll, beforeAll, describe, expect, it, vi } from 'vitest' import { memoryPersistence } from '@tanstack/ai-persistence' -import { createHarnessHost, defineHarness } from '../src' +import { createHarnessHost, defineHarness, definePlugin } from '../src' import { permissions } from '../src/first-party/permissions' import { hostBackend } from '../src/first-party/coding/backend' +import { + JobEnded, + JobOutput, + JobStarted, +} from '../src/first-party/coding/events' import { createWorkspaceTools, workspaceTools, @@ -834,6 +839,44 @@ describe('background bash in a session', () => { await host.close() }) + it('sends the start, the output, and the end of a job as plugin events', async () => { + const { backend } = memoryBackend({ + spawn: (_command, options) => { + options?.onOutput?.('built\n') + return { + wait: async () => ({ exitCode: 0 }), + kill: () => undefined, + output: () => 'built\n', + } + }, + }) + const events: Array = [] + const listener = definePlugin({ + name: 'test/job-events', + setup: (ctx) => { + ctx.on(JobStarted, (value) => events.push(['started', value])) + ctx.on(JobOutput, (value) => events.push(['output', value])) + ctx.on(JobEnded, (value) => events.push(['ended', value])) + }, + }) + const { adapter } = mockAdapter([ + () => toolCall('bash', { command: 'make', background: true }, 'c1'), + () => text('started'), + ]) + const { host, session } = await openSession(adapter, [ + workspaceTools({ root: '/workspace', backend }), + listener, + ]) + await session.send('build') + await vi.waitFor(() => expect(events).toHaveLength(3)) + expect(events).toEqual([ + ['started', { jobId: 'bash-1', command: 'make' }], + ['output', { jobId: 'bash-1', output: 'built\n' }], + ['ended', { jobId: 'bash-1', exitCode: 0 }], + ]) + await host.close() + }) + it('kills a job that still runs when the session closes', async () => { let killed = false const { backend } = memoryBackend({ diff --git a/packages/ai-sandbox/src/workspace-backend.ts b/packages/ai-sandbox/src/workspace-backend.ts index 9520365d06..4d0c8f1800 100644 --- a/packages/ai-sandbox/src/workspace-backend.ts +++ b/packages/ai-sandbox/src/workspace-backend.ts @@ -1,12 +1,18 @@ import { posix } from 'node:path' import { shellQuote } from './journal' import type { WorkspaceBackend } from '@tanstack/ai-harness/plugins/coding' -import type { ExecResult, SandboxHandle } from './contracts' +import type { ExecResult, SandboxFs, SandboxHandle } from './contracts' /** A background job keeps the newest 1 MiB of its output. */ const MAX_JOB_OUTPUT = 1024 * 1024 /** `realpath` follows at most this many links, like Linux. */ const MAX_LINKS = 40 +/** The first bytes that `readLines` gives, for the file type check. */ +const HEAD_BYTES = 8 * 1024 +/** `readLines` gives no more lines after this many bytes. */ +const MAX_PAGE_BYTES = 64 * 1024 +/** The line in the output of `readLines` between the head and the lines. */ +const LINES_MARK = '@lines' const encoder = new TextEncoder() const decoder = new TextDecoder() @@ -68,6 +74,34 @@ function createOutputTail(maxBytes: number) { return { push, text } } +/** + * The page in the output of the `readLines` command, for a page from line + * `first`. The output is the size, the first bytes as hex from `od`, a line + * {@link LINES_MARK}, the lines of the page, and a last line: `more`, or + * `total` and the number of lines. + */ +function parsePage(stdout: string, first: number) { + const rows = stdout.split('\n') + // The output ends with a newline, so the last row is empty. + if (rows.at(-1) === '') rows.pop() + const mark = rows.indexOf(LINES_MARK) + const status = rows.at(-1) ?? '' + const isPage = + mark > 0 && mark < rows.length - 1 && /^(more|total \d+)$/.test(status) + if (!isPage) { + throw new Error('readLines: the output of the sandbox is not a page.') + } + const size = Number((rows[0] ?? '').trim()) + const hex = rows.slice(1, mark).join(' ').split(/\s+/).filter(Boolean) + const head = Uint8Array.from(hex, (byte) => Number.parseInt(byte, 16)) + const lines = rows.slice(mark + 1, -1) + if (status === 'more') return { head, size, lines, total: undefined } + // An empty file has one empty line, like `''.split('\n')`. + if (size === 0) + return { head, size, lines: first === 1 ? [''] : [], total: 1 } + return { head, size, lines, total: Number(status.slice('total '.length)) } +} + /** * A `WorkspaceBackend` over a live sandbox. With it, `workspaceTools()` * reads and writes files and runs commands in the sandbox, not on this @@ -88,12 +122,15 @@ function createOutputTail(maxBytes: number) { * - `spawn` is there only when the sandbox has * `capabilities.backgroundProcesses`. * - `stat` gives `mtimeMs: 0`. The sandbox file system does not give the - * time of the last change. + * time of the last change. With `lstat`, `stat` never reads the file. + * - `readLines` reads one page of a file in the sandbox with `wc`, `od`, + * and `awk`, so only the page comes out of the sandbox. * - `realpath` resolves links in the sandbox: `lstat` finds each link, and * `readlink` in the sandbox reads its target. A sandbox without `lstat` * cannot find links, so `realpath` throws, and the tools refuse each path. * - * - `spawn` keeps the newest 1 MiB of the output of a job. + * - `spawn` keeps the newest 1 MiB of the output of a job, and gives each + * chunk to `onOutput` as it comes. * - `close({ reason: 'moved' })` keeps the jobs and the sandbox, so the * next machine can resume them. * @@ -123,10 +160,10 @@ export function sandboxWorkspaceBackend( const running = new Set<{ kill: () => void }>() let destroyed: Promise | undefined - // The read of the sandbox follows links, so this follows links too. A - // path that `readBytes` cannot read is a folder. - // ponytail: reads the whole file for its size. Only links, and a sandbox - // without `lstat`, come here. + // Without `lstat`, the read of the sandbox follows links, so this + // follows links too. A path that `readBytes` cannot read is a folder. + // ponytail: reads the whole file for its size. Only a sandbox without + // `lstat` comes here. const follow: WorkspaceBackend['stat'] = async (path) => { if (!(await fs.exists(path))) return undefined const bytes = await fs.readBytes(path).catch(() => undefined) @@ -135,8 +172,9 @@ export function sandboxWorkspaceBackend( } const stat: WorkspaceBackend['stat'] = async (path) => { - if (!fs.lstat) return follow(path) - const info = await fs.lstat(path) + const { lstat } = fs + if (!lstat) return follow(path) + const info = await lstat(path) if (!info) return undefined switch (info.type) { case 'file': @@ -146,8 +184,13 @@ export function sandboxWorkspaceBackend( // Like `fs.stat` on the host, a fifo or a socket is an empty file. case 'other': return { type: 'file', size: 0, mtimeMs: 0 } - case 'symlink': - return follow(path) + case 'symlink': { + // The target, by its real path, so the file is not read. A link to + // a missing target is nothing, like `fs.stat` on the host. + const resolved = await resolveLinks(lstat, path) + if (resolved.missing) return undefined + return stat(resolved.real) + } } } @@ -167,17 +210,14 @@ export function sandboxWorkspaceBackend( } /** - * `path` with every link resolved, part by part, or `undefined` when - * nothing is at `path`. Throws for a link to a missing target, and when - * there are more than 40 links, like a loop. + * `path` with every link resolved, part by part. `missing` is `'path'` + * when nothing is at `path`, and `'target'` for a link to a missing + * target. Throws when there are more than 40 links, like a loop. */ - const realpath: WorkspaceBackend['realpath'] = async (path) => { - const { lstat } = fs - if (!lstat) { - throw new Error( - `This sandbox has no lstat, so links in ${path} cannot be checked.`, - ) - } + const resolveLinks = async ( + lstat: NonNullable, + path: string, + ) => { let parts = posix.resolve('/', path).split('/').filter(Boolean) // The first parts that came from a link target. A missing one means a // link to a missing target. A missing part of `path` itself means @@ -198,8 +238,8 @@ export function sandboxWorkspaceBackend( const next = posix.join(real, name) const info = await lstat(next) if (!info) { - if (!isFromLink) return undefined - throw new Error(`ENOENT: ${path} is a link to ${next}, not there.`) + const missing = isFromLink ? ('target' as const) : ('path' as const) + return { real: undefined, missing } } if (info.type !== 'symlink') { real = next @@ -216,15 +256,82 @@ export function sandboxWorkspaceBackend( // The parts after it can come from an earlier link too. fromLink += targetParts.length } - return real + return { real, missing: undefined } + } + + /** + * `path` with every link resolved, or `undefined` when nothing is at + * `path`. Throws for a link to a missing target, and when there are more + * than 40 links, like a loop. + */ + const realpath: WorkspaceBackend['realpath'] = async (path) => { + const { lstat } = fs + if (!lstat) { + throw new Error( + `This sandbox has no lstat, so links in ${path} cannot be checked.`, + ) + } + const resolved = await resolveLinks(lstat, path) + switch (resolved.missing) { + case undefined: + return resolved.real + case 'path': + return undefined + case 'target': + throw new Error( + `ENOENT: ${path} is a link to a target that is not there.`, + ) + } + } + + /** + * One page of the file `path`, read in the sandbox. One command gives the + * size (`wc`), the first 8 KiB as hex (`od`), and the lines of the page + * (`awk`). So only the page comes out of the sandbox, not the whole file. + * `awk` stops at the first line after the page, so `total` is there only + * when the page gets to the end of the file. + */ + const readLines: NonNullable = async ( + path, + { offset, limit }, + ) => { + const first = Math.max(1, Math.floor(offset)) + const end = first + Math.max(0, Math.floor(limit)) + // A relative name in its folder, like `readlink`. `./` keeps a name + // that starts with `-` from being read as an option. + const name = shellQuote(`./${posix.basename(path)}`) + // `awk` prints the page, then one last line: `more` when the file goes + // on after the page, else `total` and the number of lines. + const awk = `NR >= e { more = 1; exit } NR >= o { print; n += length($0) + 1; if (n > max) e = NR + 1 } END { if (more) print "more"; else print "total " NR }` + const command = [ + `wc -c < ${name}`, + `od -An -v -tx1 -N ${HEAD_BYTES} ${name}`, + `echo ${LINES_MARK}`, + `awk -v o=${first} -v e=${end} -v max=${MAX_PAGE_BYTES} ${shellQuote(awk)} ${name}`, + ].join(' && ') + const { exitCode, stdout, stderr } = await handle.process.exec(command, { + cwd: posix.dirname(path), + }) + if (exitCode !== 0) { + throw new Error(`Cannot read ${path}: ${stderr.trim()}`) + } + return parsePage(stdout, first) } const spawn: WorkspaceBackend['spawn'] = (command, options) => { const output = createOutputTail(MAX_JOB_OUTPUT) + const onOutput = options?.onOutput const collect = async (stream: AsyncIterable) => { - for await (const chunk of stream) output.push(chunk) + for await (const chunk of stream) { + output.push(chunk) + onOutput?.(chunk) + } } - const started = handle.process.spawn(command, options) + // `onOutput` is for this backend. The sandbox gets the folder and the env. + const started = handle.process.spawn(command, { + cwd: options?.cwd, + env: options?.env, + }) // Wait for the end of the output too, so `output()` is complete. const exited = started .then(async (job) => { @@ -272,6 +379,7 @@ export function sandboxWorkspaceBackend( if (!isFile) throw new Error(`Cannot remove ${path}: no such file.`) await fs.remove(path) }, + readLines, stat, realpath, readdir: async (path) => { diff --git a/packages/ai-sandbox/tests/workspace-backend.test.ts b/packages/ai-sandbox/tests/workspace-backend.test.ts index adcf1fe38f..2652c931bc 100644 --- a/packages/ai-sandbox/tests/workspace-backend.test.ts +++ b/packages/ai-sandbox/tests/workspace-backend.test.ts @@ -1,5 +1,5 @@ import { posix } from 'node:path' -import { describe, expect, it } from 'vitest' +import { describe, expect, it, vi } from 'vitest' import { EventType } from '@tanstack/ai' import { memoryPersistence } from '@tanstack/ai-persistence' import { createHarnessHost, defineHarness } from '@tanstack/ai-harness' @@ -104,6 +104,41 @@ function memoryFs(seed: FsSeed, withLstat = true) { return { fs, files, dirs, links } } +/** The `awk` program of the page command of `readLines`. */ +const PAGE_AWK = `'NR >= e { more = 1; exit } NR >= o { print; n += length($0) + 1; if (n > max) e = NR + 1 } END { if (more) print "more"; else print "total " NR }'` + +/** The page command of `readLines` for `name`, from line `first` to `end`. */ +const pageCommand = (name: string, first: number, end: number) => + `wc -c < '${name}' && od -An -v -tx1 -N 8192 '${name}' && echo @lines && awk -v o=${first} -v e=${end} -v max=65536 ${PAGE_AWK} '${name}'` + +/** The page command with any name and range: the name, the first line, the end. */ +const PAGE_COMMAND = + /^wc -c < '(.+)' && od -An -v -tx1 -N 8192 '\1' && echo @lines && awk -v o=(\d+) -v e=(\d+) -v max=65536 '.+' '\1'$/ + +/** + * What `wc -c`, `od -An -v -tx1 -N 8192`, `echo @lines`, and the `awk` + * program print for a file of `data`, for the lines from `first` to before + * `end`. It does not stop after 64 KiB. + */ +function pageOutput(data: Uint8Array, first: number, end: number) { + const text = decoder.decode(data) + const lines = text === '' ? [] : text.split('\n') + if (text.endsWith('\n')) lines.pop() + const hex: Array = [] + const head = data.subarray(0, 8192) + for (let at = 0; at < head.length; at += 16) { + const row = [...head.subarray(at, at + 16)] + hex.push( + row.map((byte) => ` ${byte.toString(16).padStart(2, '0')}`).join(''), + ) + } + const page = lines.slice(first - 1, end - 1) + const status = lines.length >= end ? 'more' : `total ${lines.length}` + return [`${data.byteLength}`, ...hex, '@lines', ...page, status, ''].join( + '\n', + ) +} + async function* fromChunks(chunks: Array) { for (const chunk of chunks) yield chunk } @@ -118,13 +153,40 @@ async function* fromChunks(chunks: Array) { * of `x` before `flood out`. * * `readlink -- ''` prints the target of the link `name` in the folder, - * from `links`, like `readlink` in a sandbox. + * from `links`, like `readlink` in a sandbox. The page command of + * `readLines` prints what `wc`, `od`, and `awk` print for a file in `files` + * (see {@link pageOutput}). `commands` lists each command and its folder. */ -function fakeProcesses(links: ReadonlyMap = new Map()) { +function fakeProcesses( + links: ReadonlyMap = new Map(), + files: ReadonlyMap = new Map(), +) { const aborted: Array = [] const killed: Array = [] + const commands: Array<{ command: string; cwd?: string }> = [] const process: SandboxHandle['process'] = { exec: (command, options = {}) => { + commands.push({ command, cwd: options.cwd }) + const page = PAGE_COMMAND.exec(command) + if (page) { + const [, quoted = '', first = '', end = ''] = page + // The shell reads `'\''` in single quotes as one `'`. + const name = quoted.replaceAll(`'\\''`, `'`) + const data = files.get(posix.join(options.cwd ?? '/', name)) + return Promise.resolve( + data === undefined + ? { + exitCode: 1, + stdout: '', + stderr: `sh: ${name}: No such file or directory`, + } + : { + exitCode: 0, + stdout: pageOutput(data, Number(first), Number(end)), + stderr: '', + }, + ) + } const readlink = /^readlink -- '(.+)'$/.exec(command) if (readlink) { const target = links.get( @@ -174,7 +236,7 @@ function fakeProcesses(links: ReadonlyMap = new Map()) { } }, } - return { process, aborted, killed } + return { process, aborted, killed, commands } } /** @@ -190,11 +252,11 @@ function fakeSandbox( } = {}, ) { const { fs, files, dirs, links } = memoryFs(seed, options.withLstat) - const { process, aborted, killed } = fakeProcesses(links) + const { process, aborted, killed, commands } = fakeProcesses(links, files) const sandbox = makeFakeHandle('sandbox-1', 'fake', options.caps) const handle: SandboxHandle = { ...sandbox, fs, process } const backend = sandboxWorkspaceBackend(handle, { owned: options.owned }) - return { backend, files, dirs, aborted, killed, sandbox } + return { backend, fs, files, dirs, aborted, killed, commands, sandbox } } const SEED: FsSeed = { @@ -242,6 +304,17 @@ describe('sandboxWorkspaceBackend files', () => { }, ) + it('stat of a link does not read the file, when the sandbox has lstat', async () => { + const { backend, fs } = fakeSandbox(SEED) + const readBytes = vi.spyOn(fs, 'readBytes') + expect(await backend.stat('/workspace/to-file')).toEqual({ + type: 'file', + size: 5, + mtimeMs: 0, + }) + expect(readBytes).not.toHaveBeenCalled() + }) + it.each([true, false])( 'stat gives undefined for a missing path and a broken link (lstat: %s)', async (withLstat) => { @@ -340,6 +413,84 @@ describe('sandboxWorkspaceBackend realpath', () => { }) }) +describe('sandboxWorkspaceBackend readLines', () => { + /** A sandbox with `/workspace/notes.txt`, five lines. */ + function notes() { + const sandbox = fakeSandbox({ + files: { '/workspace/notes.txt': 'one\ntwo\nthree\nfour\nfive\n' }, + dirs: ['/workspace'], + }) + const readBytes = vi.spyOn(sandbox.fs, 'readBytes') + return { ...sandbox, readBytes } + } + + it('reads one page in the sandbox, so the file does not come out whole', async () => { + const { backend, commands, readBytes } = notes() + const page = await backend.readLines?.('/workspace/notes.txt', { + offset: 2, + limit: 2, + }) + expect(commands).toEqual([ + { command: pageCommand('./notes.txt', 2, 4), cwd: '/workspace' }, + ]) + expect(page).toEqual({ + head: encoder.encode('one\ntwo\nthree\nfour\nfive\n'), + size: 24, + lines: ['two', 'three'], + total: undefined, + }) + expect(readBytes).not.toHaveBeenCalled() + }) + + it('counts the lines when the page gets to the end of the file', async () => { + const { backend } = notes() + const page = await backend.readLines?.('/workspace/notes.txt', { + offset: 4, + limit: 10, + }) + expect(page).toMatchObject({ lines: ['four', 'five'], total: 5 }) + }) + + it('gives one empty line for an empty file, like the host', async () => { + const { backend } = fakeSandbox({ + files: { '/workspace/empty.txt': '' }, + dirs: ['/workspace'], + }) + const page = await backend.readLines?.('/workspace/empty.txt', { + offset: 1, + limit: 10, + }) + expect(page).toEqual({ + head: new Uint8Array(), + size: 0, + lines: [''], + total: 1, + }) + }) + + it('throws with the error of the sandbox for a missing file', async () => { + const { backend } = notes() + await expect( + backend.readLines?.('/workspace/gone.txt', { offset: 1, limit: 10 }), + ).rejects.toThrow( + 'Cannot read /workspace/gone.txt: sh: ./gone.txt: No such file or directory', + ) + }) + + it('quotes a file name with a quote, and keeps a leading - from being an option', async () => { + const { backend, commands } = fakeSandbox({ + files: { "/workspace/-it's.txt": 'a\n' }, + dirs: ['/workspace'], + }) + const page = await backend.readLines?.("/workspace/-it's.txt", { + offset: 1, + limit: 1, + }) + expect(commands[0]?.command).toContain(`wc -c < './-it'\\''s.txt'`) + expect(page).toMatchObject({ lines: ['a'], total: 1 }) + }) +}) + describe('sandboxWorkspaceBackend exec', () => { it('runs a command in the sandbox with the folder and the env', async () => { const { backend } = fakeSandbox() @@ -383,6 +534,17 @@ describe('sandboxWorkspaceBackend spawn', () => { expect(job?.output()).toBe('build out\nbuild err\n') }) + it('gives each chunk of the output to onOutput as it comes', async () => { + const { backend } = fakeSandbox() + const chunks: Array = [] + const job = backend.spawn?.('build', { + onOutput: (text) => chunks.push(text), + }) + await job?.wait() + // stdout and stderr come at the same time, so their order is not fixed. + expect(chunks.sort()).toEqual(['build err\n', 'build out\n']) + }) + it('kills a background command, also before it started', async () => { const { backend, killed } = fakeSandbox() const job = backend.spawn?.('serve') From 1dacc9934f92e820205fbd157ae853d881f42680 Mon Sep 17 00:00:00 2001 From: Alem Tuzlak Date: Sun, 18 Oct 2026 10:12:57 +0200 Subject: [PATCH 06/17] docs(ai-harness, ai-sandbox, ai-sandbox-local-process): run the coding tools for many threads, with workspaces per thread, jobs, moving threads between hosts, and snapshots in the log --- .changeset/ai-harness-coding-tools.md | 12 +- .changeset/ai-sandbox-harness.md | 3 +- .changeset/sandbox-local-process-env.md | 5 + docs/config.json | 18 +- docs/harness/coding-agents.md | 2 + docs/harness/coding-tools-backends.md | 29 ++- docs/harness/coding-tools.md | 16 +- docs/harness/coding-workspaces.md | 230 ++++++++++++++++++++++++ docs/harness/deploy.md | 1 + docs/harness/durable-sessions.md | 1 + docs/harness/snapshots.md | 44 ++++- docs/sandbox/auth.md | 5 +- docs/sandbox/providers.md | 17 +- docs/sandbox/quick-start.md | 2 +- 14 files changed, 344 insertions(+), 41 deletions(-) create mode 100644 .changeset/sandbox-local-process-env.md create mode 100644 docs/harness/coding-workspaces.md diff --git a/.changeset/ai-harness-coding-tools.md b/.changeset/ai-harness-coding-tools.md index 2ed90ac76a..37ddf04add 100644 --- a/.changeset/ai-harness-coding-tools.md +++ b/.changeset/ai-harness-coding-tools.md @@ -5,9 +5,15 @@ Add coding tools in the new Node entry `@tanstack/ai-harness/plugins/coding`: - `workspaceTools({ root })` gives `read_file` (text pages, images, PDFs), `write_file`, `edit_file` or `patch` (`editStyle`), `list_files`, `grep`, `bash` (with `background`), `webfetch`, and `websearch` (with a `SearchProvider`). With `permissions()`, edits, `bash`, and `webfetch` ask first, and the rules also apply to the real path of a link. `outside: 'ask'` asks before a path outside `root`. The `cwd` thread setting gives each thread its own working folder. -- A `WorkspaceBackend` runs the tools. `hostBackend` is the default. On Windows, commands do not run a program from the working folder by its name. -- `snapshots()` adds `/undo` and `/redo`, which restore only the files of the turn (git 2.26 or newer), file restore for `session.revert()`, and a diff for each model step. -- `formatter()` formats each written file with the formatter of the project. With `permissions()`, each run asks first. +- `workspaceTools({ workspace: { bind, reattach? } })` gives each thread its own workspace. The root of the workspace is the boundary of every tool. The session log keeps where the workspace is (`harness.workspace.bound`). On another host, `reattach` gives the workspace back (for example a resumed sandbox), or `'pinned'` or `'gone'`. A gone workspace sends `tanstack/workspace-tools:stranded`. +- `createHarnessHost({ ownership: { hostId } })` gives a host a stable id. Then a thread that uses a folder on that host is pinned to it, and other hosts refuse its tool calls. Without `hostId`, threads are not pinned. +- Commands, `rg`, `git`, and the formatters get only a short list of variables of the server process, plus the `env` of the backend call. Git runs without the hooks and the `fsmonitor` of the project. +- The tools resolve links with `realpath` under the file lock, and fail when a backend cannot resolve a path. A write fails when the file changed after the tool read it. The file locks are shared by the sessions of a host. +- Background jobs: at most 8 per session (`maxJobs`), the newest 1 MiB of output per job, and the events `tanstack/workspace-tools:job-started`, `job-output`, and `job-ended`. A job that an earlier host left running is stopped at the next tool call, with a note to the model. +- `read_file` sends images and PDFs up to 256 KB. `webfetch` gives the model the first 256 KB of a page. +- A `WorkspaceBackend` runs the tools. `hostBackend` is the default. `realpath` is required. `spawn` can give a `jobId`, and `killJob`, `binding`, and `close` are optional. On Windows, commands do not run a program from the working folder by its name. +- `snapshots()` adds `/undo` and `/redo`, which restore only the files of the turn (git 2.26 or newer), file restore for `session.revert()`, and a diff for each model step. Each thread of each tenant has its own shadow repository, which never stores `.env` files. A step keeps only the files of its own run. The steps and the undo point are in the session log, so another host with the same repository can undo. A revert uses the last 500 steps. `purgeSnapshots()` removes the repository of a thread. A failed snapshot sends `tanstack/snapshots:failed`. +- `formatter()` formats each written file with the formatter of the project. With `permissions()`, each run asks first. A run that changes the file sends `tanstack/formatter:rewrote`. - `WorkspaceHooks` lets other plugins run code after a read or a write. `@vscode/ripgrep` and `turndown` are optional peer dependencies. diff --git a/.changeset/ai-sandbox-harness.md b/.changeset/ai-sandbox-harness.md index 10f2443916..89b9f4e848 100644 --- a/.changeset/ai-sandbox-harness.md +++ b/.changeset/ai-sandbox-harness.md @@ -5,6 +5,7 @@ Add `@tanstack/ai-sandbox/harness` for `@tanstack/ai-harness`: - `codingAgents({ sandbox, agents, workspace? })` is a harness plugin. The lead model gets one tool per coding agent (Claude Code, Codex, Grok Build, or any ACP agent). Each agent works in the sandbox, keeps its own session per thread (also after a restart), and starts read-only in the `plan` mode of `permissions()`. `/fresh [agent]` starts a new session. `workspace: 'shared'` (default) runs the agents of a thread one at a time in one sandbox. `'per-agent'` gives each agent its own sandbox. -- `sandboxWorkspaceBackend(sandbox)` runs the coding tools of `@tanstack/ai-harness/plugins/coding` (`workspaceTools()`, `snapshots()`, `formatter()`) in a sandbox. +- `codingAgents` runs agents only on the host that owns the log of the thread, so `'shared'` runs one agent at a time also with many hosts. +- `sandboxWorkspaceBackend(sandbox, { owned? })` runs the coding tools of `@tanstack/ai-harness/plugins/coding` (`workspaceTools()`, `snapshots()`, `formatter()`) in a sandbox. It resolves links in the sandbox, reads one page of a file in the sandbox, and stops background jobs by their id, also from another host. `binding()` gives the provider and the sandbox id, so another host can resume the sandbox. `close()` stops the jobs of the backend, and with `owned: true` destroys the sandbox when the session closes. `@tanstack/ai-harness` is an optional peer dependency. diff --git a/.changeset/sandbox-local-process-env.md b/.changeset/sandbox-local-process-env.md new file mode 100644 index 0000000000..eb3957d401 --- /dev/null +++ b/.changeset/sandbox-local-process-env.md @@ -0,0 +1,5 @@ +--- +'@tanstack/ai-sandbox-local-process': minor +--- + +A command in `localProcessSandbox` gets only a few variables of the host process: `PATH`, `HOME`, `LANG`, `LC_ALL`, `TMPDIR`, `TEMP`, `TMP`, and on Windows `SystemRoot`, `ComSpec`, `PATHEXT`, `WINDIR`, and `USERPROFILE`. The workspace secrets and the env of the call still reach it. Provider keys and other secrets of the server do not. If a command needs another variable of the host, set it as a workspace secret or in the env of the call. `scrubEnv` still removes a variable from the env of each command. diff --git a/docs/config.json b/docs/config.json index 855e19e2d1..f49e926b5b 100644 --- a/docs/config.json +++ b/docs/config.json @@ -957,7 +957,7 @@ "label": "Deploy a harness", "to": "harness/deploy", "addedAt": "2026-10-13", - "updatedAt": "2026-10-17" + "updatedAt": "2026-10-18" }, { "label": "Send inputs safely", @@ -974,7 +974,7 @@ "label": "Durable sessions", "to": "harness/durable-sessions", "addedAt": "2026-10-13", - "updatedAt": "2026-10-17" + "updatedAt": "2026-10-18" }, { "label": "Run background work on workers", @@ -1056,19 +1056,24 @@ "label": "Give your agent coding tools", "to": "harness/coding-tools", "addedAt": "2026-10-13", - "updatedAt": "2026-10-17" + "updatedAt": "2026-10-18" }, { "label": "Sandbox and extend the coding tools", "to": "harness/coding-tools-backends", "addedAt": "2026-10-13", - "updatedAt": "2026-10-16" + "updatedAt": "2026-10-18" + }, + { + "label": "Run the coding tools for many threads", + "to": "harness/coding-workspaces", + "addedAt": "2026-10-18" }, { "label": "Undo a turn", "to": "harness/snapshots", "addedAt": "2026-10-13", - "updatedAt": "2026-10-17" + "updatedAt": "2026-10-18" }, { "label": "Build a coding agent", @@ -1079,7 +1084,8 @@ { "label": "Delegate to coding agents", "to": "harness/coding-agents", - "addedAt": "2026-10-13" + "addedAt": "2026-10-13", + "updatedAt": "2026-10-18" }, { "label": "Switch between agent profiles", diff --git a/docs/harness/coding-agents.md b/docs/harness/coding-agents.md index 125577e14c..e264d7f401 100644 --- a/docs/harness/coding-agents.md +++ b/docs/harness/coding-agents.md @@ -108,6 +108,8 @@ To start over, run `/fresh claude_code`, or `/fresh` for every agent. 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. +Only the host that owns the log of the thread runs its agents. So with many servers, `'shared'` still runs one agent at a time. + ## Plan mode When the `permissions()` plugin is in [`plan` mode](./permissions#pick-a-mode), the agents start read-only: diff --git a/docs/harness/coding-tools-backends.md b/docs/harness/coding-tools-backends.md index 04ffd2ae84..c93d939d04 100644 --- a/docs/harness/coding-tools-backends.md +++ b/docs/harness/coding-tools-backends.md @@ -62,12 +62,12 @@ export const sandboxed = defineHarness({ 3. When you are done, stop the sandbox with `await sandbox.destroy()`. -Every provider works. Get the handle from `create()`, or from `resume()` for a sandbox that runs already. See [Providers](../sandbox/providers). +Every provider works. Get the handle from `create()`, or from `resume()` for a sandbox that runs already. See [Providers](../sandbox/providers). To create a sandbox for each thread, and resume it on another server, see [Run the coding tools for many threads](./coding-workspaces). In a sandbox, some things are different: - Paths are POSIX paths in the sandbox, like `/workspace/src/a.ts`, also on a Windows host. -- The tools do not check where links lead. The sandbox backend has no `realpath`. +- The tools check where links lead with `lstat` and `readlink` in the sandbox. A sandbox without `lstat` cannot find links, so the tools refuse every path. - `background: true` needs a sandbox with `capabilities.backgroundProcesses`. - A sandbox with `capabilities.killableProcesses` off cannot stop a command after the time limit. The call returns, and the command continues in the sandbox. - `list_files` and `grep` use `rg` in the sandbox. The `@vscode/ripgrep` binary is a host file, so the sandbox does not use it. @@ -82,11 +82,15 @@ The backend must have: | --- | --- | | `readFile(path)` | Gives the bytes of a file. Throws when there is no file. | | `writeFile(path, data)` | Creates or replaces a file. Creates the missing parent folders. | -| `stat(path)` | Gives `{ type, size, mtimeMs }`, or `undefined` when nothing is at `path`. | +| `stat(path)` | Gives `{ type, size, mtimeMs }`, or `undefined` when nothing is at `path`. It does not read the file. | +| `realpath(path)` | Gives the path with every link resolved, or `undefined` when nothing is at `path`. Throws when it cannot resolve a path. | | `readdir(path)` | Gives the entries of a folder. A link has the type `'link'`. | | `exec(command, options)` | Runs a shell command and waits. Gives `{ exitCode, stdout, stderr }`. | -`exec` resolves with the exit code of a command that fails. It does not throw. After `options.timeoutMs`, it stops the command with every process that it started, and gives exit code 124. +- `exec` resolves with the exit code of a command that fails. It does not throw. +- After `options.timeoutMs`, `exec` stops the command with every process that it started, and gives exit code 124. `options.signal` stops it too. +- `options.env` holds the variables to add. Give a command only these and the few variables that it needs to start. Do not give it the variables of your server. +- When `realpath` throws, the tool call fails. So a path with a link that the backend cannot resolve is never used. The backend can also have: @@ -95,8 +99,17 @@ The backend can also have: | `shell` | `'sh'` or `'cmd'`. Sets the path style and how arguments are quoted. | `sh` quotes, and the path style of this machine. | | `remove(path)` | Removes one file. | `patch` cannot delete or move files. | | `readLines(path, { offset, limit })` | Gives one page of a file: `{ head, size, lines, total }`. `head` is the first 8 KiB. | `read_file` reads the whole file for each page. | -| `realpath(path)` | Gives the path with every link resolved. | The tools do not check links. | -| `spawn(command, options)` | Starts a command in the background. | `background: true` gives an error. | + +For background jobs and for threads that move to another server: + +| Member | What it does | Without it | +| --- | --- | --- | +| `spawn(command, options)` | Starts a command in the background. Gives `{ wait, kill, output, jobId? }`. `jobId` resolves with the id of the process, for `killJob`. | `background: true` gives an error. | +| `killJob(jobId)` | Stops a job that `spawn` started, also on an earlier process or machine. | A job that a crash left running is not stopped. | +| `binding()` | Gives `{ kind: 'sandbox', provider, sandboxId }` for a sandbox that another machine can resume. | The workspace is a folder on this host. | +| `close({ reason })` | Releases the backend. `'closed'`: the session closed. `'moved'`: the session continues on another machine, so keep what it can resume. | Nothing to release. | + +`close` must never destroy a sandbox that the backend did not create. The fastest start is to wrap `hostBackend`. This backend logs each command before it runs: @@ -213,7 +226,7 @@ export const withToml = formatter({ - `timeoutMs` limits each run. The default is 20 seconds. - In a sandbox or on your own backend, give `formatter()` the same `backend` as `workspaceTools()`. -A formatter that fails, stops at the time limit, or is not installed does not fail the tool call. The file stays as the tool wrote it, and the plugin sends `FormatFailed`: +A formatter that fails, stops at the time limit, or is not installed does not fail the tool call. The file stays as the tool wrote it, and the plugin sends `FormatFailed`. A formatter that changes the file sends `FormatterRewrote` with `{ path, formatter }`: ```ts group=harness-coding-tools-backends import { FormatFailed } from '@tanstack/ai-harness/plugins/coding' @@ -228,7 +241,7 @@ export const formatAlerts = definePlugin({ }) ``` -Clients get the same event as a `harness.plugin.event` with the name `tanstack/formatter:failed`. +Clients get the same events as a `harness.plugin.event` with the name `tanstack/formatter:failed` or `tanstack/formatter:rewrote`. ### Each run asks first diff --git a/docs/harness/coding-tools.md b/docs/harness/coding-tools.md index 5e60cde431..af23ddc114 100644 --- a/docs/harness/coding-tools.md +++ b/docs/harness/coding-tools.md @@ -134,7 +134,7 @@ These tools run without a question. | `list_files` | Lists the files in a folder. An optional glob, like `src/**/*.ts`, filters them. | | `grep` | Searches the file contents with a regular expression. Gives `file:line: text`. | -- `read_file` gives images (PNG, JPEG, GIF, WebP) and PDFs up to 5 MB to the model as media. +- `read_file` gives images (PNG, JPEG, GIF, WebP) and PDFs up to 256 KB to the model as media. For a bigger file, the model gets an error that tells it to make a smaller copy. - `read_file` refuses a binary file with a short note. For a missing file, it names up to 3 files with a close name. - `list_files` and `grep` skip the files that `.gitignore` names, the `.git` and `node_modules` folders, and links. - `list_files` shows at most 1000 files. `grep` shows at most 200 matches. @@ -156,6 +156,8 @@ A model gets `edit_file` or `patch`, not both. Every model gets `write_file`. Th `edit_file` keeps the line endings and the byte order mark of the file. It also finds text that is different only in Unicode form or in the spaces at the end of a line. +When a file changed after the model read it, for example by another agent, the three tools fail. The model reads the file again, then edits it. + To let the user undo the changes of a turn, add [snapshots](./snapshots). ### Commands and the web @@ -170,8 +172,13 @@ For `bash`: - The model gets the exit code and the last 1000 lines (20 KiB) of the output. - A command stops after 2 minutes, with exit code 124. Every process that it started stops too. Set another limit with `bashTimeoutMs`. The model can ask for up to 10 minutes. -- With `background: true`, the call returns at once with a job id. When the job ends, the model gets a note with the output, and a new turn starts. The time limit does not apply to a background job. - With `spillDir`, long output goes to a file in that folder, and the model gets the path. `spillDir` is relative to `root`, or absolute. +- A command gets only a few variables of your process, like `PATH` and `HOME`. Your provider keys stay out of it. See [Keep server secrets out of commands](./coding-workspaces#2-keep-server-secrets-out-of-commands). + +With `background: true`, the call returns at once with a job id. When the job ends, the model gets a note with the output, and a new turn starts. + +- The time limit does not apply to a background job. +- At most 8 background jobs run at a time in a session. Set another limit with `maxJobs`. Each job keeps the newest 1 MiB of its output. - When the session closes, the background jobs stop. - The log keeps each background job until it ends. When a crash stops a job, the next host adds the note "Background job bash-1 stopped when the host restarted." @@ -182,7 +189,8 @@ For `webfetch`: - It reads `http` and `https` URLs only, and follows at most 5 redirects. - It refuses `localhost`, private network addresses, and cloud metadata addresses. It checks the DNS answer of each connection. - HTML comes back as Markdown, or as text without `turndown`. The model can ask for `format: 'html'`. Text and JSON come back as they are. -- It reads at most 5 MB. It stops after 30 seconds, and the model can ask for up to 120 seconds. +- The model gets the first 256 KB of a page, with a note when there is more. It reads at most 5 MB. +- It stops after 30 seconds, and the model can ask for up to 120 seconds. Change the web tools with the `web` option: @@ -301,4 +309,4 @@ await session.configure({ settings: { cwd: 'web-app' } }) - A question before each edit, command, and web fetch. - A working folder for each thread, and a question before a path outside `root`. -Next: run the same tools in a sandbox, or format each file after a write. See [Sandbox and extend the coding tools](./coding-tools-backends). +Next: run the same tools in a sandbox, or format each file after a write. See [Sandbox and extend the coding tools](./coding-tools-backends). To give each thread its own folder on a server, see [Run the coding tools for many threads](./coding-workspaces). diff --git a/docs/harness/coding-workspaces.md b/docs/harness/coding-workspaces.md new file mode 100644 index 0000000000..a2a051fb6e --- /dev/null +++ b/docs/harness/coding-workspaces.md @@ -0,0 +1,230 @@ +--- +title: Run the coding tools for many threads +id: harness-coding-workspaces +order: 7 +description: "Give each thread of a coding agent its own folder or sandbox, keep server secrets out of commands, let agents share files safely, watch background jobs, and move a thread to another server." +keywords: + - tanstack ai + - harness + - coding agent + - workspace tools + - sandbox + - multi-tenant + - background jobs +--- + +Your server runs coding agents for many users. Each thread must work in its own folder or sandbox. A command must not see the API keys of your server. Two agents in one folder must not lose each other's edits. And a thread must continue when another server takes it. + +`workspace: { bind }` gives each thread its own workspace. The session log keeps where that workspace is, so the next server finds it again. + +## 1. Give each thread its own workspace + +Pass `workspace` to `workspaceTools()` in place of `root`. `bind` gets the run, and gives a backend with its `root`: + +```ts group=harness-coding-workspaces +import { mkdir } from 'node:fs/promises' +import { join } from 'node:path' +import { defineHarness } from '@tanstack/ai-harness' +import { permissions } from '@tanstack/ai-harness/plugins' +import { hostBackend, workspaceTools } from '@tanstack/ai-harness/plugins/coding' +import { openaiText } from '@tanstack/ai-openai' + +export const coder = defineHarness({ + name: 'acme/coder', + adapter: openaiText('gpt-6.1-sol'), + plugins: () => [ + permissions(), + workspaceTools({ + workspace: { + bind: async ({ run }) => { + const root = join('/srv/workspaces', run.tenantId ?? 'default', run.threadId) + await mkdir(root, { recursive: true }) + return { ...hostBackend, root } + }, + }, + }), + ], +}) +``` + +- The tools call `bind` once for each thread, before the first turn or at the first tool call. If `bind` throws, the next turn calls it again. +- The `root` that `bind` gives is the boundary of every tool. A path outside it is refused, or asks first with `outside: 'ask'`. The `cwd` thread setting must stay inside it. +- The tools resolve every link with `realpath` before they read or write, while they hold the lock on the file. So a link in one thread cannot lead into the folder of another thread. +- The session log keeps a `harness.workspace.bound` record with the place of the workspace. [Move a thread to another server](#5-move-a-thread-to-another-server) uses it. + +For a sandbox per thread, create it in `bind`, and give the backend `owned: true`. Then the sandbox is destroyed when the session closes: + +```ts group=harness-coding-workspaces +import { sandboxWorkspaceBackend } from '@tanstack/ai-sandbox/harness' +import { dockerSandbox } from '@tanstack/ai-sandbox-docker' + +const docker = dockerSandbox({ image: 'node:22' }) + +export const sandboxed = workspaceTools({ + workspace: { + bind: async () => ({ + ...sandboxWorkspaceBackend(await docker.create({}), { owned: true }), + root: '/workspace', + }), + }, +}) +``` + +## 2. Keep server secrets out of commands + +A command gets only a few variables of the server process: + +- `PATH`, `HOME`, `LANG`, `LC_ALL`, `TMPDIR`, `TEMP`, and `TMP`. +- On Windows also `SystemRoot`, `ComSpec`, `PATHEXT`, `WINDIR`, and `USERPROFILE`. +- `AGENT=1`, so a tool can see that an agent runs it. + +Provider keys, `DATABASE_URL`, and every other variable stay in the server. The same rule applies to `bash`, `rg`, `git`, and the formatters. Git also runs without the hooks and the `fsmonitor` of the project. + +To give commands a variable, add it to the `env` of the backend calls: + +```ts group=harness-coding-workspaces +import type { WorkspaceBackend } from '@tanstack/ai-harness/plugins/coding' + +const extra = { NPM_TOKEN: 'npm_read_only_token' } +const { spawn } = hostBackend + +const withToken: WorkspaceBackend = { + ...hostBackend, + exec: (command, options) => + hostBackend.exec(command, { ...options, env: { ...options?.env, ...extra } }), + spawn: + spawn && + ((command, options) => + spawn(command, { ...options, env: { ...options?.env, ...extra } })), +} + +export const withTokenTools = workspaceTools({ root: '/srv/repo', backend: withToken }) +``` + +`exec` runs the foreground commands. `spawn` runs the background commands. + +## 3. Let agents share a folder + +Two agents, or two sessions, can work in the same folder. No edit is lost: + +- A write holds a lock on the file. The lock is shared by every session on this server, and it is keyed by the root and the real path of the file. +- `write_file`, `edit_file`, and `patch` fail when the file changed after the tool read it. The model gets `src/app.ts changed since it was read. Read it again before you edit it.`, reads the file again, and tries again. +- The check compares the size and the time of the last change. A sandbox gives no change time, so there only a change of size is seen. + +## 4. Watch background jobs + +With `background: true`, `bash` returns at once with a job id like `bash-1`. Two limits keep jobs small: + +- At most 8 background jobs run at a time in a session. One more gives the model a tool error. Set another limit with `maxJobs`. +- Each job keeps the newest 1 MiB of its output. + +Each job sends events. Listen for them in a plugin: + +```ts group=harness-coding-workspaces +import { definePlugin } from '@tanstack/ai-harness' +import { JobEnded, JobOutput, JobStarted } from '@tanstack/ai-harness/plugins/coding' + +export const jobLog = definePlugin({ + name: 'acme/job-log', + setup: (ctx) => { + ctx.on(JobStarted, (job) => console.log(`${job.jobId} started: ${job.command}`)) + ctx.on(JobOutput, (job) => process.stdout.write(job.output)) + ctx.on(JobEnded, (job) => console.log(`${job.jobId} ended: ${job.exitCode}`)) + }, +}) +``` + +A client gets each event as a `CUSTOM` event with the name `harness.plugin.event`. Its value is `{ plugin, name, value }`: + +```ts group=harness-coding-workspaces-client +import { createHarnessClient } from '@tanstack/ai-harness/client' +import type { coder } from './harness' + +const client = createHarnessClient({ + url: '/api/harness', + threadId: 'fix-login', +}) + +for await (const { event } of client.events()) { + if (event.type !== 'CUSTOM' || event.name !== 'harness.plugin.event') continue + const { value } = event + if (typeof value !== 'object' || value === null) continue + if (!('name' in value) || value.name !== 'tanstack/workspace-tools:job-output') continue + if ('value' in value) console.log(value.value) +} +``` + +The coding plugins send these events: + +| Name | When | +| --- | --- | +| `tanstack/workspace-tools:job-started` | A background job started. | +| `tanstack/workspace-tools:job-output` | New output of a job. At most one event each 250 ms, with at most the newest 64 KiB. | +| `tanstack/workspace-tools:job-ended` | A job ended, after its last output event. | +| `tanstack/workspace-tools:stranded` | The workspace of the thread is gone. See below. | +| `tanstack/snapshots:failed` | A snapshot failed. The turn goes on. | + +The formatter also sends `tanstack/formatter:rewrote` and `tanstack/formatter:failed`. See [Format each file after a write](./coding-tools-backends#format-each-file-after-a-write). + +## 5. Move a thread to another server + +One server owns a thread at a time. When it stops, another server takes the thread (see [One host drives a thread](./durable-sessions#one-host-drives-a-thread)). What happens to the workspace depends on where it is. + +### A sandbox + +The log keeps the provider and the id of the sandbox. The next server calls `reattach` in place of `bind`. Resume the sandbox there: + +```ts group=harness-coding-workspaces +export const movable = workspaceTools({ + workspace: { + bind: async () => ({ + ...sandboxWorkspaceBackend(await docker.create({}), { owned: true }), + root: '/workspace', + }), + reattach: async ({ binding }) => { + if (binding.kind === 'host') return 'pinned' + if (binding.provider !== 'docker') return 'gone' + const handle = await docker.resume({ id: binding.sandboxId }) + if (!handle) return 'gone' + return { ...sandboxWorkspaceBackend(handle, { owned: true }), root: binding.root } + }, + }, +}) +``` + +- Give the workspace back to continue in the same files. +- Give `'gone'` when the sandbox is gone. Then each tool call of the thread fails with `The workspace of this thread, /workspace in the docker sandbox , is gone, so the tools cannot work on it.`, and the tools send `tanstack/workspace-tools:stranded`. +- Without `reattach`, a sandbox counts as `'gone'`. + +### A folder on one server + +A folder on the disk of one server cannot move. Give each server a stable id: + +```ts group=harness-coding-workspaces +import { createHarnessHost } from '@tanstack/ai-harness' +import { memoryPersistence } from '@tanstack/ai-persistence' + +const host = createHarnessHost({ + persistence: memoryPersistence(), + ownership: { hostId: process.env.HOSTNAME ?? 'web-1' }, +}) +``` + +- With `hostId`, a thread that uses a local folder is pinned to this server. Workers on other servers do not take it. +- Another server that opens the thread refuses its tool calls with `The workspace of this thread is on the host web-1. Only that host can continue the thread.` +- Keep the id the same after a restart, for example the machine or pod name. +- Without `hostId`, threads are not pinned, and every server calls `bind` again. Use this when every server sees the same folder, for example on a shared disk. + +### After a move + +- Background jobs do not move. At the next tool call of the thread, the new server stops each job that the old server left running. The model gets the note `Background job bash-1 was stopped: it ran on a host before the session moved to this one.` +- `/undo`, `/redo`, and a revert work on the new server when it reaches the same snapshot repository. See [Undo a turn](./snapshots#on-more-than-one-server). +- [Coding agents](./coding-agents) run only on the server that owns the log. So two servers never run the agents of one shared workspace at the same time. + +## What you have now + +- A folder or a sandbox for each thread, with its root as the boundary of every tool. +- Commands that never see the secrets of your server. +- Agents that share a folder and never lose an edit. +- Events for every background job, and a limit on how many run. +- Threads that move to another server with their sandbox, or stay on the server that has their folder. diff --git a/docs/harness/deploy.md b/docs/harness/deploy.md index 7d41029e56..711a7ba31e 100644 --- a/docs/harness/deploy.md +++ b/docs/harness/deploy.md @@ -29,6 +29,7 @@ Give every server the same log store. Then any server can take a request for any - A [worker](./workers) runs the threads that no server has open: submitted inputs, schedules, and the work that a deploy cut. - Set the same `media: { secret }` on the handler of each server, so each server checks the media URLs of the others. See [Serve the session over HTTP](./connect#serve-the-session-over-http). - Stop each server with `host.close({ recoverable: true })`. A worker takes its cut work at once. See [Stop a host for a deploy](./durable-sessions#stop-a-host-for-a-deploy). +- If your agents use the [coding tools](./coding-tools) on the disk of a server, set `ownership: { hostId }` on each host. Then a thread stays on the server that has its folder. See [Move a thread to another server](./coding-workspaces#5-move-a-thread-to-another-server). ## Build a worker artifact diff --git a/docs/harness/durable-sessions.md b/docs/harness/durable-sessions.md index d6c83bd257..9de802bc88 100644 --- a/docs/harness/durable-sessions.md +++ b/docs/harness/durable-sessions.md @@ -290,6 +290,7 @@ const ownedHost = createHarnessHost({ - `ttlMs`: how long the ownership lasts after each renewal. The default is 30 000. After it, another host can take the thread. - `renewMs`: how often the timer renews. The default is 10 000. +- `hostId`: a stable id for this host, for example the machine or pod name. With it, a thread whose coding tools use a folder on this host stays on this host. Other hosts refuse its tool calls. See [A folder on one server](./coding-workspaces#a-folder-on-one-server). An input that waits for the owner keeps its JSON overrides (`model`, `modelOptions`, `reasoning`, and `promptCache`). The `adapter` and `tools` overrides are objects, so the owner runs the input without them. diff --git a/docs/harness/snapshots.md b/docs/harness/snapshots.md index 8e1e7cd6e3..985bcdfc5d 100644 --- a/docs/harness/snapshots.md +++ b/docs/harness/snapshots.md @@ -26,10 +26,8 @@ import { openaiText } from '@tanstack/ai-openai' import { memoryPersistence } from '@tanstack/ai-persistence' const root = process.cwd() -const history = snapshots({ - root, - dataDir: join(homedir(), '.acme-coder', 'snapshots'), -}) +const dataDir = join(homedir(), '.acme-coder', 'snapshots') +const history = snapshots({ root, dataDir }) const coder = defineHarness({ name: 'acme/coder', @@ -44,7 +42,7 @@ const session = await host.open(coder, { threadId: 'fix-login' }) ``` - `root` is the workspace folder, an absolute path. -- `dataDir` is a folder for the snapshot repositories, an absolute path. Put it outside `root`. Each workspace gets its own repository in it. +- `dataDir` is a folder for the snapshot repositories, an absolute path. Put it outside `root`. Each thread of each tenant gets its own repository in it. - The plugin needs the `git` command. `/undo` and `/redo` need git 2.26 or later. - Your project does not need to be a git repository, and the plugin never uses the `.git` folder of your project. @@ -59,10 +57,15 @@ In a terminal app, the user types `/undo` and `/redo`. See [Build a coding agent - `/undo` puts back each file that the turn changed in `root`. This includes the edits of the file tools and the changes of `bash` commands. - `/undo` and `/redo` do not touch other files. This is also true for a file that you add or edit after the turn. -- If you edit a file that the turn changed, `/undo` and `/redo` still write over it, and your edit is lost. -- The snapshots skip the files that `.gitignore` names, like `node_modules`. `/undo` does not touch them. +- The snapshots skip the files that `.gitignore` names, like `node_modules`, and `.env` files. `/undo` does not touch them. - A new turn clears `/redo`. +When someone else edits the same files: + +- If you edit a file that the turn changed, `/undo` and `/redo` still write over it, and your edit is lost. +- A step keeps only the files that its own run changed. So when another agent or session edits a file during the turn, `/undo` keeps that edit. +- A `bash` command does not tell which files it changed. So its step keeps every file that changed while the command ran. + In your own UI, add `expose: { commands: ['undo', 'redo'] }` to `defineHarness`. A client can run only the commands that the harness exposes. Then call `client.command({ name: 'undo' })`. The result text arrives as a `harness.command.result` event. ## See what a step changed @@ -86,6 +89,7 @@ for (const { record } of await log.read({ logId: logIdOf({ logId: 'fix-login' }) ``` - Without `to`, `diff` compares with the files now. `history.diff(record.to)` shows what changed after that step. +- `diff` finds a snapshot only in the repository of a session that is open. - A step that changed no file is not saved. ## Go back to an earlier message @@ -101,6 +105,7 @@ if (firstAnswer?.id) await session.revert(firstAnswer.id) await session.unrevert() ``` +- A revert uses the last 500 steps of the thread. - A revert touches only the files that the tool calls after that message changed. Each file goes back to its state before the first of these calls. A file that did not exist then is removed. - Other files stay as they are. Your own edits between the hidden turns stay too. - `unrevert()` touches only the same files. It puts them back as they were when the revert started, with your edits. @@ -131,14 +136,35 @@ export const sandboxHistory = snapshots({ Git then runs in the sandbox, and `dataDir` is a path in the sandbox. The sandbox image needs `git`. +## On more than one server + +The undo point and the steps are in the session log. So when another server takes the thread, `/undo`, `/redo`, and a revert work there too, if that server reaches the same snapshot repository: + +- In a sandbox, the repository is in the sandbox. A server that uses the same sandbox reaches it. See [A sandbox](./coding-workspaces#a-sandbox). +- On the host, give every server the same `dataDir` on a shared disk. Or pin each thread to the server that has its folder with `ownership: { hostId }`. See [A folder on one server](./coding-workspaces#a-folder-on-one-server). + +## Remove the snapshots of a thread + +When you delete a thread, remove its snapshot repository too. Call `purgeSnapshots` with the options of the plugin and the thread: + +```ts group=harness-snapshots +import { purgeSnapshots } from '@tanstack/ai-harness/plugins/coding' + +await purgeSnapshots({ root, dataDir }, { tenantId: 'acme', threadId: 'fix-login' }) +``` + +- Call it after the sessions of the thread close. See [Lifecycle and retention](./lifecycle). +- Leave out `tenantId` for a thread in the default tenant. +- It waits for the git commands of that repository in this process. +- It resolves also when there is no repository. + ## Known limits - One level of undo. A second `/undo` answers `Nothing to undo.` After a `/redo`, `/undo` works again. - `/undo` and `/redo` work only when the session is idle. During a turn, they answer `Wait until the turn ends, then try again.` - While a revert stands, `/undo` and `/redo` answer `A revert stands. Run unrevert, or send a prompt, first.` -- The undo point is kept in memory. After a restart, or when the session opens again, `/undo` works again after the next turn. - If you sent a message while the turn ran, `/undo` removes the conversation from that message on. The files still go back to the start of the turn. -- Without `git`, the plugin prints `snapshots: git was not found. Snapshots are off.` once, and adds no commands. `diff()` throws. The turns run as usual. +- Without `git`, the plugin sends a `tanstack/snapshots:failed` event (`SnapshotFailed`) with `snapshots: git was not found. Snapshots are off.`, and adds no commands. `diff()` throws. The turns run as usual. A snapshot that fails for another reason sends the same event, and the turn goes on. ## What you have now diff --git a/docs/sandbox/auth.md b/docs/sandbox/auth.md index 5f0ef5b5ee..ddf0204be8 100644 --- a/docs/sandbox/auth.md +++ b/docs/sandbox/auth.md @@ -51,8 +51,9 @@ const stream = chat({ }) ``` -`scrubEnv` removes keys the host process inherited. If the CLI sees -`XAI_API_KEY`, it can prefer that key over your login. +`scrubEnv` removes keys from the env of each command, for example a key that +a workspace secret sets. If the CLI sees `XAI_API_KEY`, it can prefer that key +over your login. The keys of the host process never reach the command. ## API key diff --git a/docs/sandbox/providers.md b/docs/sandbox/providers.md index 8eaa138c18..9faff0a499 100644 --- a/docs/sandbox/providers.md +++ b/docs/sandbox/providers.md @@ -78,10 +78,13 @@ import { localProcessSandbox } from '@tanstack/ai-sandbox-local-process' const dev = localProcessSandbox() ``` -- **Isolation:** none. The harness runs directly on your host, inheriting your - host environment. Use it for trusted or dev work only. There is no boundary - between the agent and your machine. -- **Auth / env:** inherits the host environment. Set `authMode` on the harness +- **Isolation:** none. The harness runs directly on your host. Use it for + trusted or dev work only. There is no boundary between the agent and your + machine. +- **Auth / env:** a command gets only a few variables of the host process + (`PATH`, `HOME`, `LANG`, `LC_ALL`, `TMPDIR`, `TEMP`, `TMP`, and on Windows + `SystemRoot`, `ComSpec`, `PATHEXT`, `WINDIR`, `USERPROFILE`), plus the + workspace secrets and the env of the call. Set `authMode` on the harness (`'host'` or `'api-key'`). The provider does not pick this. See [Harness Auth](./auth). - **Snapshot / resume:** no snapshots and no durable resume-by-id; each run @@ -95,9 +98,9 @@ The provider is where the agent runs, not how it signs in. The default login. A local-process run can be your laptop or a GitHub runner. See [Harness Auth](./auth). -`localProcessSandbox` inherits the host environment, including any API keys -exported there. If you set `authMode: 'host'`, pass `scrubEnv` so those keys -do not override the CLI login: +An API key that the host process exports does not reach a command. A key +that a workspace secret or the call sets does. If you set `authMode: 'host'`, +pass `scrubEnv` so those keys do not override the CLI login: ```ts import { localProcessSandbox } from '@tanstack/ai-sandbox-local-process' diff --git a/docs/sandbox/quick-start.md b/docs/sandbox/quick-start.md index ab43c418b3..0b93610b22 100644 --- a/docs/sandbox/quick-start.md +++ b/docs/sandbox/quick-start.md @@ -184,7 +184,7 @@ const stream = chat({ ``` Set `authMode: 'host'` so the adapter uses `grok login`. `scrubEnv` removes -keys the host process inherited. Those keys would override the login. +keys from the env of each command. Those keys would override the login. A local-process run can also be a CI runner. That machine has no browser login. Set `authMode: 'api-key'`. Then inject `XAI_API_KEY` as a workspace From f73c34eba51f5305074fa99342cf86a631e1e1db Mon Sep 17 00:00:00 2001 From: Alem Tuzlak Date: Sun, 18 Oct 2026 10:13:08 +0200 Subject: [PATCH 07/17] feat(ai-harness): ownership.hostId pins threads that use a local folder to their host, with E2E specs for env, roots, moving a sandbox thread, stranded, and pinning --- packages/ai-harness/src/host.ts | 17 +- packages/ai-harness/src/ownership.ts | 13 +- .../tests/workspace-binding.test.ts | 100 +++- testing/e2e/src/routeTree.gen.ts | 21 + .../e2e/src/routes/api.coding-workspace.ts | 445 ++++++++++++++++++ .../tests/tools-test/coding-workspace.spec.ts | 145 ++++++ 6 files changed, 737 insertions(+), 4 deletions(-) create mode 100644 testing/e2e/src/routes/api.coding-workspace.ts create mode 100644 testing/e2e/tests/tools-test/coding-workspace.spec.ts diff --git a/packages/ai-harness/src/host.ts b/packages/ai-harness/src/host.ts index e0879c41e0..f4b4eaee49 100644 --- a/packages/ai-harness/src/host.ts +++ b/packages/ai-harness/src/host.ts @@ -162,8 +162,10 @@ export interface HarnessHostOptions

{ * its first write until the thread has no work. `ttlMs`: how long the * ownership lasts after each renewal (default 30000 ms). After it, another * host can take the thread. `renewMs`: how often one timer renews every - * thread that the host owns (default 10000 ms). Leave it out for the - * defaults. + * thread that the host owns (default 10000 ms). `hostId`: a stable id + * for this host, for example the machine or pod name. With it, a thread + * that uses a local folder is pinned to this host, and another host + * refuses the thread. Leave it out for the defaults: no pins. */ ownership?: OwnershipOptions /** How the sessions of this host live. Leave it out for the defaults. */ @@ -1222,6 +1224,15 @@ export function createHarnessHost

( 'Host access needs stores.sessions: each session keeps its owner and members there.', ) } + const affinity = options.ownership?.hostId + if ( + affinity !== undefined && + (typeof affinity !== 'string' || affinity.trim() === '') + ) { + throw new Error( + 'ownership.hostId must be a non-empty string, for example the machine or pod name. Leave it out when the host has no stable id.', + ) + } const policy = access && membersOnlyPolicyOf(access.authorize) const memory = memoryPersistence().stores // ponytail: memory credentials when none are given. Pass `credentials` @@ -1291,6 +1302,8 @@ export function createHarnessHost

( const ownership = new HostOwnership({ store: logStore, ownerId: hostId, + // With a stable host id, folder bindings pin their logs to this host. + ...(affinity !== undefined ? { affinity } : {}), ...(options.ownership ? { timing: options.ownership } : {}), }) /** The shared log of each log with an open session, by its store log id. */ diff --git a/packages/ai-harness/src/ownership.ts b/packages/ai-harness/src/ownership.ts index bf28606af2..26145c865d 100644 --- a/packages/ai-harness/src/ownership.ts +++ b/packages/ai-harness/src/ownership.ts @@ -4,8 +4,19 @@ import type { LogStore } from '@tanstack/ai-persistence' /** How long a host owns a log after each renewal, and how often it renews. */ export const OWNERSHIP = { ttlMs: 30_000, renewMs: 10_000 } -/** Ownership timing for a host. A missing field uses {@link OWNERSHIP}. */ +/** + * Ownership timing for a host, and its stable id. A missing timing field + * uses {@link OWNERSHIP}. + */ export interface OwnershipOptions { + /** + * A stable id for this host, for example the machine or pod name. Keep it + * the same after a restart. When you set it, a thread that uses a local + * folder is pinned to this host, and another host refuses the thread with + * the pinned error. When you do not set it, threads are not pinned, and + * every host can continue every thread. + */ + hostId?: string /** * How long the host owns a log after each renewal, in milliseconds of the * store clock. Default 30 000. diff --git a/packages/ai-harness/tests/workspace-binding.test.ts b/packages/ai-harness/tests/workspace-binding.test.ts index 564c6a8a04..9f520dee34 100644 --- a/packages/ai-harness/tests/workspace-binding.test.ts +++ b/packages/ai-harness/tests/workspace-binding.test.ts @@ -71,10 +71,15 @@ function open(options: { adapter: AnyTextAdapter plugins: Array crashes?: boolean + /** The stable id of the host, `ownership.hostId`. */ + hostId?: string }) { const host = createHarnessHost({ persistence: options.persistence, - ...(options.crashes ? { ownership: CRASHING_OWNERSHIP } : {}), + ownership: { + ...(options.crashes ? CRASHING_OWNERSHIP : {}), + ...(options.hostId !== undefined ? { hostId: options.hostId } : {}), + }, }) const opened = host.open( defineHarness({ @@ -124,12 +129,14 @@ async function readOnce( persistence: HarnessPersistence, options: WorkspaceToolsOptions, id: string, + hostId?: string, ) { const { adapter, calls } = readTurn(id) const { host, opened } = open({ persistence, adapter, plugins: [workspaceTools({ ...options, web: false })], + ...(hostId !== undefined ? { hostId } : {}), }) const session = await opened await session.send('read') @@ -254,6 +261,97 @@ describe('the workspace binding in the session log', () => { expect(await canOtherHostTake(log, LOG)).toBe(false) }) + it('pins a folder thread to the host with ownership.hostId, across restarts', async () => { + const persistence = memoryPersistence() + const { log } = persistence.stores + const inFolder = { root: '/ws', backend: backend() } + + expect(await readOnce(persistence, inFolder, 'c0', 'pod-a')).toBe(READ) + expect(await recordsOf(persistence, 'harness.workspace.bound')).toEqual([ + { + type: 'harness.workspace.bound', + binding: { kind: 'host', hostId: 'pod-a', root: '/ws' }, + }, + ]) + // The log is pinned to pod-a: no other host gets it. + expect(await canOtherHostTake(log, LOG)).toBe(false) + + // Host B does not run the thread. Its input waits for pod-a. + const b = readTurn('b0') + const other = open({ + persistence, + adapter: b.adapter, + hostId: 'pod-b', + plugins: [workspaceTools({ ...inFolder, web: false })], + }) + const sent = (await other.opened).send('read') + await vi.waitFor( + async () => { + expect(await log.offered({ logId: LOG })).toHaveLength(1) + }, + { timeout: 5000 }, + ) + expect(b.calls).toHaveLength(0) + + // Host A after a restart, with the same hostId, runs the input of B. + const a = readTurn('c1') + const again = open({ + persistence, + adapter: a.adapter, + hostId: 'pod-a', + plugins: [workspaceTools({ ...inFolder, web: false })], + }) + await again.opened + await sent + expect(resultOf(a.calls, 'c1')).toBe(READ) + expect(b.calls).toHaveLength(0) + await other.host.close() + await again.host.close() + }) + + it('refuses a folder thread of another hostId with the pinned error', async () => { + const persistence = memoryPersistence() + const { log } = persistence.stores + const inFolder = { root: '/ws', backend: backend() } + // A binding to pod-a, and no pin yet: host B gets the log. + await appendToLog(log, { threadId: THREAD }, [ + { + type: 'harness.workspace.bound', + binding: { kind: 'host', hostId: 'pod-a', root: '/ws' }, + }, + ]) + + expect(await readOnce(persistence, inFolder, 'c0', 'pod-b')).toContain( + 'The workspace of this thread is /ws on the host pod-a. Only that host can continue the thread.', + ) + expect(await canOtherHostTake(log, LOG)).toBe(false) + + // pod-a after a restart continues the thread. + expect(await readOnce(persistence, inFolder, 'c1', 'pod-a')).toBe(READ) + }) + + it('does not pin a folder thread when the hosts have no hostId', async () => { + const persistence = memoryPersistence() + const inFolder = { root: '/ws', backend: backend() } + + expect(await readOnce(persistence, inFolder, 'c0')).toBe(READ) + expect(await readOnce(persistence, inFolder, 'c1')).toBe(READ) + + expect(await recordsOf(persistence, 'harness.workspace.bound')).toEqual([ + { + type: 'harness.workspace.bound', + binding: { kind: 'host', root: '/ws' }, + }, + ]) + expect(await canOtherHostTake(persistence.stores.log, LOG)).toBe(true) + }) + + it('refuses an empty ownership.hostId', () => { + expect(() => createHarnessHost({ ownership: { hostId: ' ' } })).toThrow( + 'ownership.hostId must be a non-empty string', + ) + }) + it('stops a job that a crashed host left running, and notes it', async () => { const persistence = memoryPersistence() const spawn = () => ({ diff --git a/testing/e2e/src/routeTree.gen.ts b/testing/e2e/src/routeTree.gen.ts index 6fe275eef9..c444ba1298 100644 --- a/testing/e2e/src/routeTree.gen.ts +++ b/testing/e2e/src/routeTree.gen.ts @@ -78,6 +78,7 @@ import { Route as ApiChatCompletionsIncompleteStreamRouteImport } from './routes import { Route as ApiClientToolResumeStoreRouteImport } from './routes/api.client-tool-resume-store' import { Route as ApiCloudflareBindingWireRouteImport } from './routes/api.cloudflare-binding-wire' import { Route as ApiCodingToolsRouteImport } from './routes/api.coding-tools' +import { Route as ApiCodingWorkspaceRouteImport } from './routes/api.coding-workspace' import { Route as ApiCompactionUsageWireRouteImport } from './routes/api.compaction-usage-wire' import { Route as ApiCompactionWireRouteImport } from './routes/api.compaction-wire' import { Route as ApiContextOverflowRouteImport } from './routes/api.context-overflow' @@ -558,6 +559,11 @@ const ApiCodingToolsRoute = ApiCodingToolsRouteImport.update({ path: '/api/coding-tools', getParentRoute: () => rootRouteImport, } as any) +const ApiCodingWorkspaceRoute = ApiCodingWorkspaceRouteImport.update({ + id: '/api/coding-workspace', + path: '/api/coding-workspace', + getParentRoute: () => rootRouteImport, +} as any) const ApiCompactionUsageWireRoute = ApiCompactionUsageWireRouteImport.update({ id: '/api/compaction-usage-wire', path: '/api/compaction-usage-wire', @@ -1234,6 +1240,7 @@ export interface FileRoutesByFullPath { '/api/client-tool-resume-store': typeof ApiClientToolResumeStoreRoute '/api/cloudflare-binding-wire': typeof ApiCloudflareBindingWireRoute '/api/coding-tools': typeof ApiCodingToolsRoute + '/api/coding-workspace': typeof ApiCodingWorkspaceRoute '/api/compaction-usage-wire': typeof ApiCompactionUsageWireRoute '/api/compaction-wire': typeof ApiCompactionWireRoute '/api/context-overflow': typeof ApiContextOverflowRoute @@ -1421,6 +1428,7 @@ export interface FileRoutesByTo { '/api/client-tool-resume-store': typeof ApiClientToolResumeStoreRoute '/api/cloudflare-binding-wire': typeof ApiCloudflareBindingWireRoute '/api/coding-tools': typeof ApiCodingToolsRoute + '/api/coding-workspace': typeof ApiCodingWorkspaceRoute '/api/compaction-usage-wire': typeof ApiCompactionUsageWireRoute '/api/compaction-wire': typeof ApiCompactionWireRoute '/api/context-overflow': typeof ApiContextOverflowRoute @@ -1609,6 +1617,7 @@ export interface FileRoutesById { '/api/client-tool-resume-store': typeof ApiClientToolResumeStoreRoute '/api/cloudflare-binding-wire': typeof ApiCloudflareBindingWireRoute '/api/coding-tools': typeof ApiCodingToolsRoute + '/api/coding-workspace': typeof ApiCodingWorkspaceRoute '/api/compaction-usage-wire': typeof ApiCompactionUsageWireRoute '/api/compaction-wire': typeof ApiCompactionWireRoute '/api/context-overflow': typeof ApiContextOverflowRoute @@ -1798,6 +1807,7 @@ export interface FileRouteTypes { | '/api/client-tool-resume-store' | '/api/cloudflare-binding-wire' | '/api/coding-tools' + | '/api/coding-workspace' | '/api/compaction-usage-wire' | '/api/compaction-wire' | '/api/context-overflow' @@ -1985,6 +1995,7 @@ export interface FileRouteTypes { | '/api/client-tool-resume-store' | '/api/cloudflare-binding-wire' | '/api/coding-tools' + | '/api/coding-workspace' | '/api/compaction-usage-wire' | '/api/compaction-wire' | '/api/context-overflow' @@ -2172,6 +2183,7 @@ export interface FileRouteTypes { | '/api/client-tool-resume-store' | '/api/cloudflare-binding-wire' | '/api/coding-tools' + | '/api/coding-workspace' | '/api/compaction-usage-wire' | '/api/compaction-wire' | '/api/context-overflow' @@ -2360,6 +2372,7 @@ export interface RootRouteChildren { ApiClientToolResumeStoreRoute: typeof ApiClientToolResumeStoreRoute ApiCloudflareBindingWireRoute: typeof ApiCloudflareBindingWireRoute ApiCodingToolsRoute: typeof ApiCodingToolsRoute + ApiCodingWorkspaceRoute: typeof ApiCodingWorkspaceRoute ApiCompactionUsageWireRoute: typeof ApiCompactionUsageWireRoute ApiCompactionWireRoute: typeof ApiCompactionWireRoute ApiContextOverflowRoute: typeof ApiContextOverflowRoute @@ -2959,6 +2972,13 @@ declare module '@tanstack/react-router' { preLoaderRoute: typeof ApiCodingToolsRouteImport parentRoute: typeof rootRouteImport } + '/api/coding-workspace': { + id: '/api/coding-workspace' + path: '/api/coding-workspace' + fullPath: '/api/coding-workspace' + preLoaderRoute: typeof ApiCodingWorkspaceRouteImport + parentRoute: typeof rootRouteImport + } '/api/compaction-usage-wire': { id: '/api/compaction-usage-wire' path: '/api/compaction-usage-wire' @@ -3906,6 +3926,7 @@ const rootRouteChildren: RootRouteChildren = { ApiClientToolResumeStoreRoute: ApiClientToolResumeStoreRoute, ApiCloudflareBindingWireRoute: ApiCloudflareBindingWireRoute, ApiCodingToolsRoute: ApiCodingToolsRoute, + ApiCodingWorkspaceRoute: ApiCodingWorkspaceRoute, ApiCompactionUsageWireRoute: ApiCompactionUsageWireRoute, ApiCompactionWireRoute: ApiCompactionWireRoute, ApiContextOverflowRoute: ApiContextOverflowRoute, diff --git a/testing/e2e/src/routes/api.coding-workspace.ts b/testing/e2e/src/routes/api.coding-workspace.ts new file mode 100644 index 0000000000..1df9337c10 --- /dev/null +++ b/testing/e2e/src/routes/api.coding-workspace.ts @@ -0,0 +1,445 @@ +import { tmpdir } from 'node:os' +import { join, resolve, sep } from 'node:path' +import { createFileRoute } from '@tanstack/react-router' +import { EventType } from '@tanstack/ai' +import { createAnthropicChat } from '@tanstack/ai-anthropic' +import { + createHarnessHost, + defineHarness, + definePlugin, + logIdOf, +} from '@tanstack/ai-harness' +import { + WorkspaceStranded, + hostBackend, + workspaceTools, +} from '@tanstack/ai-harness/plugins/coding' +import { memoryPersistence } from '@tanstack/ai-persistence' +import { sandboxWorkspaceBackend } from '@tanstack/ai-sandbox/harness' +import { localProcessSandbox } from '@tanstack/ai-sandbox-local-process' +import type { AnyTextAdapter } from '@tanstack/ai' +import type { HarnessPersistence, HarnessPlugin } from '@tanstack/ai-harness' +import type { + WorkspaceBinder, + WorkspaceBinding, +} from '@tanstack/ai-harness/plugins/coding' + +/** + * The workspace of the coding tools for many threads and many hosts. Body: + * `{ scenario, root }`. `root` must be a folder in the temp folder that the + * test made. The answer is JSON with what the test checks. + * + * - `env`: `bash` prints a server env var. The command must not see it. + * - `other-root`: each thread gets its own folder in `root`. The thread + * `beta` tries to read a file of the thread `alpha`. + * - `sandbox-move`: a thread works in a local-process sandbox on `root`. + * Its host closes, and a second host continues the thread. The second + * host gets the sandbox from `reattach`, and does not call `bind`. + * - `stranded`: the same, but the second host has no `reattach`. The tool + * call fails, and the stranded event comes. + * - `pin`: hosts with `ownership.hostId`. A folder thread runs on `pod-a`. + * `pod-b` does not run the next input: it waits for `pod-a`, which runs + * it after a restart. + * + * A custom `fetch` answers with canned Anthropic streams, so no aimock + * fixture is needed. + */ + +const DUMMY_KEY = 'sk-ant-e2e-test-dummy-key' +const MODEL = 'claude-sonnet-5-5' +/** The server env var of the `env` scenario. */ +const SECRET_NAME = 'E2E_CODING_SERVER_SECRET' +const SECRET_VALUE = 'server-only-value' + +type AnthropicEvent = { type: string } & Record + +const sseOf = (events: Array) => + events + .map((event) => `event: ${event.type}\ndata: ${JSON.stringify(event)}\n\n`) + .join('') + +const messageStart = (id: string): AnthropicEvent => ({ + type: 'message_start', + message: { + id, + type: 'message', + role: 'assistant', + content: [], + model: MODEL, + stop_reason: null, + stop_sequence: null, + usage: { input_tokens: 5, output_tokens: 0 }, + }, +}) + +const ending = (stopReason: string): Array => [ + { type: 'content_block_stop', index: 0 }, + { + type: 'message_delta', + delta: { stop_reason: stopReason, stop_sequence: null }, + usage: { output_tokens: 1 }, + }, + { type: 'message_stop' }, +] + +/** An answer that calls the tool `name` with `input`. */ +const call = (name: string, input: Record) => (id: string) => + sseOf([ + messageStart(id), + { + type: 'content_block_start', + index: 0, + content_block: { type: 'tool_use', id: `toolu_${id}`, name, input: {} }, + }, + { + type: 'content_block_delta', + index: 0, + delta: { + type: 'input_json_delta', + partial_json: JSON.stringify(input), + }, + }, + ...ending('tool_use'), + ]) + +/** An answer with the text `text`. */ +const reply = (text: string) => (id: string) => + sseOf([ + messageStart(id), + { + type: 'content_block_start', + index: 0, + content_block: { type: 'text', text: '' }, + }, + { + type: 'content_block_delta', + index: 0, + delta: { type: 'text_delta', text }, + }, + ...ending('end_turn'), + ]) + +/** + * A model that answers with `answers`, in order. `calls` counts its + * requests. Each model needs its own `prefix`: the log keeps the result of + * each tool call id. + */ +function scriptedModel(prefix: string, answers: Array<(id: string) => string>) { + const counter = { calls: 0 } + const scriptedFetch: typeof fetch = async () => { + counter.calls += 1 + const answer = answers[counter.calls - 1] ?? reply('No more answers.') + return new Response(answer(`msg_${prefix}_${counter.calls}`), { + headers: { 'Content-Type': 'text/event-stream' }, + }) + } + const adapter = createAnthropicChat(MODEL, DUMMY_KEY, { + fetch: scriptedFetch, + }) + return { adapter, counter } +} + +/** A model that reads `path` with `read_file`, then answers `text`. */ +const reader = (prefix: string, path: string, text: string) => + scriptedModel(prefix, [call('read_file', { path }), reply(text)]) + +/** A host on `persistence`, with `ownership.hostId` when `hostId` is set. */ +const hostOn = (persistence: HarnessPersistence, hostId?: string) => + createHarnessHost({ + persistence, + ...(hostId !== undefined ? { ownership: { hostId } } : {}), + }) + +/** A harness with `adapter` and `plugins`. */ +const harnessOf = (adapter: AnyTextAdapter, plugins: Array) => + defineHarness({ + name: 'e2e/coding-workspace', + adapter, + plugins: () => plugins, + }) + +/** Send `message` and read the turn. Gives the tool results that the client gets. */ +async function turnOf( + session: { + send: (message: string) => { stream: () => AsyncIterable } + }, + message: string, +) { + const results: Array = [] + for await (const chunk of session.send(message).stream()) { + if (typeof chunk !== 'object' || chunk === null) continue + if (Reflect.get(chunk, 'type') === EventType.TOOL_CALL_RESULT) { + results.push(String(Reflect.get(chunk, 'content'))) + } + } + return { results } +} + +/** `bash` prints the server env var {@link SECRET_NAME}. */ +async function envScenario(root: string) { + // The server has the var. A command must not get it. + process.env[SECRET_NAME] = SECRET_VALUE + const model = scriptedModel('env', [ + call('bash', { + command: `node -e "process.stdout.write('secret=' + (process.env.${SECRET_NAME} ?? 'unset'))"`, + }), + reply('The command ran.'), + ]) + const host = hostOn(memoryPersistence()) + try { + const session = await host.open( + harnessOf(model.adapter, [workspaceTools({ root, web: false })]), + { threadId: 'coding-env' }, + ) + const turn = await turnOf(session, 'print the secret') + return { + ...turn, + serverHasSecret: process.env[SECRET_NAME] === SECRET_VALUE, + } + } finally { + await host.close() + } +} + +/** + * Each thread works in its own folder of `root`. The thread `alpha` reads + * its file. The thread `beta` tries the same file by its full path and by + * a relative path. + */ +async function otherRootScenario(root: string) { + const secret = join(root, 'alpha', 'secret.txt') + const workspace: WorkspaceBinder = { + bind: async ({ run }) => ({ + ...hostBackend, + root: join(root, run.threadId), + }), + } + const alpha = reader('alpha', 'secret.txt', 'Read my file.') + const beta = scriptedModel('beta', [ + call('read_file', { path: secret }), + call('read_file', { path: '../alpha/secret.txt' }), + reply('I cannot read it.'), + ]) + const host = hostOn(memoryPersistence()) + try { + const plugins = [workspaceTools({ workspace, web: false })] + const a = await host.open(harnessOf(alpha.adapter, plugins), { + threadId: 'alpha', + }) + const b = await host.open(harnessOf(beta.adapter, plugins), { + threadId: 'beta', + }) + return { + alpha: await turnOf(a, 'read your file'), + beta: await turnOf(b, 'read the file of alpha'), + } + } finally { + await host.close() + } +} + +/** + * A thread in a local-process sandbox on `root` moves to a second host. + * With `withReattach`, the second host resumes the sandbox. Without it, the + * sandbox counts as gone. + */ +async function sandboxMoveScenario(root: string, withReattach: boolean) { + const provider = localProcessSandbox({ dir: root }) + const persistence = memoryPersistence() + const threadId = 'coding-sandbox-move' + const bindOf = (counter: { calls: number }) => async () => { + counter.calls += 1 + return { + ...sandboxWorkspaceBackend(await provider.create({})), + root: '/workspace', + } + } + const strandedOf = (list: Array) => + definePlugin({ + name: 'e2e/stranded', + setup: (ctx) => { + ctx.on(WorkspaceStranded, (event) => list.push(event.binding)) + }, + }) + + // The first host binds the sandbox, and closes. + const firstBind = { calls: 0 } + const first = hostOn(persistence) + let firstTurn: Awaited> + try { + const model = reader('first', 'notes.txt', 'Read notes.txt.') + const session = await first.open( + harnessOf(model.adapter, [ + workspaceTools({ workspace: { bind: bindOf(firstBind) }, web: false }), + ]), + { threadId }, + ) + firstTurn = await turnOf(session, 'read notes.txt') + } finally { + await first.close() + } + + // The second host continues the thread. + const bind = { calls: 0 } + const reattached: Array = [] + const stranded: Array = [] + const workspace: WorkspaceBinder = { + bind: bindOf(bind), + ...(withReattach + ? { + reattach: async ({ binding }) => { + reattached.push(binding) + if (binding.kind !== 'sandbox') return 'gone' + const handle = await provider.resume({ id: binding.sandboxId }) + return handle + ? { ...sandboxWorkspaceBackend(handle), root: binding.root } + : 'gone' + }, + } + : {}), + } + const second = hostOn(persistence) + try { + const model = reader('second', 'notes.txt', 'Read notes.txt again.') + const session = await second.open( + harnessOf(model.adapter, [ + strandedOf(stranded), + workspaceTools({ workspace, web: false }), + ]), + { threadId }, + ) + const secondTurn = await turnOf(session, 'read notes.txt again') + return { + first: { ...firstTurn, bindCalls: firstBind.calls }, + second: { + ...secondTurn, + bindCalls: bind.calls, + reattached, + stranded, + }, + } + } finally { + await second.close() + } +} + +/** + * A folder thread on hosts with `ownership.hostId`. `pod-a` runs the first + * turn. `pod-b` does not run the next input: the input waits for `pod-a`, + * and `pod-a` runs it after a restart with the same id. + */ +async function pinScenario(root: string) { + const persistence = memoryPersistence() + const threadId = 'coding-pin' + const tools = () => [workspaceTools({ root, web: false })] + + const a = reader('pod_a', 'notes.txt', 'Read notes.txt on pod-a.') + const first = hostOn(persistence, 'pod-a') + try { + const session = await first.open(harnessOf(a.adapter, tools()), { + threadId, + }) + await session.send('read notes.txt') + } finally { + await first.close() + } + + const b = reader('pod_b', 'notes.txt', 'Read notes.txt on pod-b.') + const other = hostOn(persistence, 'pod-b') + const restarted = hostOn(persistence, 'pod-a') + try { + const session = await other.open(harnessOf(b.adapter, tools()), { + threadId, + }) + const sent = Promise.resolve(session.send('read notes.txt again')) + // pod-b offers the input to the host of the thread. + const logId = logIdOf({ logId: threadId }) + const deadline = Date.now() + 5000 + let offers = 0 + while (Date.now() < deadline) { + offers = (await persistence.stores.log.offered({ logId })).length + if (offers > 0) break + await new Promise((done) => setTimeout(done, 50)) + } + const podBCallsBefore = b.counter.calls + + // pod-a after a restart takes the input and runs it. + const again = reader( + 'pod_a2', + 'notes.txt', + 'Read notes.txt on pod-a again.', + ) + const view = await restarted.open(harnessOf(again.adapter, tools()), { + threadId, + }) + // The send of pod-b ends when pod-a ran the input. + const podBSendEnded = await Promise.race([ + sent.then(() => true), + new Promise((done) => setTimeout(() => done(false), 10_000)), + ]) + const transcript = await view.transcript() + const records = await persistence.stores.log.read({ logId }) + return { + offers, + podBCalls: podBCallsBefore, + podBCallsAfter: b.counter.calls, + podBSendEnded, + restartedCalls: again.counter.calls, + answers: transcript + .filter((message) => message.role === 'assistant') + .map((message) => + typeof message.content === 'string' ? message.content : '', + ) + .filter((text) => text !== ''), + bindings: records + .map(({ record }) => record) + .filter((record) => record.type === 'harness.workspace.bound') + .map((record) => Reflect.get(record, 'binding')), + } + } finally { + await other.close() + await restarted.close() + } +} + +export const Route = createFileRoute('/api/coding-workspace')({ + server: { + handlers: { + POST: async ({ request }) => { + const body: unknown = await request.json() + const field = (name: string) => + typeof body === 'object' && body !== null && name in body + ? Reflect.get(body, name) + : undefined + const scenario = field('scenario') + const rootField = field('root') + const root = typeof rootField === 'string' ? resolve(rootField) : '' + if (!root.startsWith(resolve(tmpdir()) + sep)) { + return new Response('root must be in the temp folder', { + status: 400, + }) + } + try { + switch (scenario) { + case 'env': + return Response.json(await envScenario(root)) + case 'other-root': + return Response.json(await otherRootScenario(root)) + case 'sandbox-move': + return Response.json(await sandboxMoveScenario(root, true)) + case 'stranded': + return Response.json(await sandboxMoveScenario(root, false)) + case 'pin': + return Response.json(await pinScenario(root)) + default: + return new Response('unknown scenario', { status: 400 }) + } + } catch (error) { + return Response.json({ + error: error instanceof Error ? error.message : String(error), + }) + } + }, + }, + }, +}) diff --git a/testing/e2e/tests/tools-test/coding-workspace.spec.ts b/testing/e2e/tests/tools-test/coding-workspace.spec.ts new file mode 100644 index 0000000000..94059333d6 --- /dev/null +++ b/testing/e2e/tests/tools-test/coding-workspace.spec.ts @@ -0,0 +1,145 @@ +import { mkdir, mkdtemp, rm, writeFile } from 'node:fs/promises' +import { tmpdir } from 'node:os' +import { join } from 'node:path' +import { test, expect } from '../fixtures' +import type { APIRequestContext } from '@playwright/test' + +/** + * The workspace of `workspaceTools()` for many threads and many hosts. Each + * scenario of `/api/coding-workspace` runs in a temp folder that the test + * makes, with a scripted model, and answers with JSON. See the route for + * the scenarios. + */ + +/** The workspace of the running test. */ +let root = '' +const file = (path: string) => join(root, path) + +/** The message of a tool result that is an error, `{ "error": "..." }`. */ +const errorOf = (result: string | undefined): unknown => + JSON.parse(result ?? '{}').error + +/** Run `scenario` in `root`. Gives the JSON of the route. */ +async function run(request: APIRequestContext, scenario: string) { + const response = await request.post('/api/coding-workspace', { + data: { scenario, root }, + }) + const body = await response.text() + expect(response.ok(), body).toBe(true) + const json: Record = JSON.parse(body) + expect(json.error, body).toBeUndefined() + return json +} + +test.describe('Coding tools workspace', () => { + test.beforeEach(async () => { + root = await mkdtemp(join(tmpdir(), 'e2e-coding-workspace-')) + }) + + test.afterEach(async () => { + await rm(root, { recursive: true, force: true }) + }) + + test('a command does not see an env var of the server', async ({ + request, + }) => { + const result = await run(request, 'env') + + // The server has the var. The command gets only the allowlist. + expect(result.serverHasSecret).toBe(true) + expect(result.results).toEqual(['exit code: 0\nsecret=unset']) + }) + + test('a thread cannot read the folder of another thread', async ({ + request, + }) => { + await mkdir(file('alpha')) + await mkdir(file('beta')) + await writeFile(file('alpha/secret.txt'), 'alpha only\n') + + const result = await run(request, 'other-root') + + expect(result.alpha).toMatchObject({ results: ['1\talpha only'] }) + const beta = result.beta as { results: Array } + expect(beta.results.map(errorOf)).toEqual([ + `Path "${file('alpha/secret.txt')}" is outside the workspace.`, + 'Path "../alpha/secret.txt" is outside the workspace.', + ]) + }) + + test('a sandbox thread moves to a second host with reattach, not bind', async ({ + request, + }) => { + await writeFile(file('notes.txt'), 'alpha\nbeta\n') + + const result = await run(request, 'sandbox-move') + + expect(result.first).toMatchObject({ + results: ['1\talpha\n2\tbeta'], + bindCalls: 1, + }) + expect(result.second).toMatchObject({ + results: ['1\talpha\n2\tbeta'], + bindCalls: 0, + reattached: [ + { + kind: 'sandbox', + provider: 'local-process', + sandboxId: root, + root: '/workspace', + }, + ], + stranded: [], + }) + }) + + test('a sandbox thread on a host with no reattach is stranded', async ({ + request, + }) => { + await writeFile(file('notes.txt'), 'alpha\nbeta\n') + + const result = await run(request, 'stranded') + + const second = result.second as { + results: Array + bindCalls: number + stranded: Array + } + expect(second.bindCalls).toBe(0) + expect(second.results.map(errorOf)).toEqual([ + `The workspace of this thread, /workspace in the local-process sandbox ${root}, is gone, so the tools cannot work on it.`, + ]) + // The event comes each time the tools try the workspace: before the + // turn, and at the tool call. + expect(second.stranded.length).toBeGreaterThan(0) + for (const binding of second.stranded) { + expect(binding).toEqual({ + kind: 'sandbox', + provider: 'local-process', + sandboxId: root, + root: '/workspace', + }) + } + }) + + test('a folder thread is pinned to its ownership.hostId, also after a restart', async ({ + request, + }) => { + await writeFile(file('notes.txt'), 'alpha\nbeta\n') + + const result = await run(request, 'pin') + + expect(result.bindings).toEqual([{ kind: 'host', hostId: 'pod-a', root }]) + // pod-b offers the input, and its model never runs. + expect(result.offers).toBe(1) + expect(result.podBCalls).toBe(0) + expect(result.podBCallsAfter).toBe(0) + expect(result.podBSendEnded).toBe(true) + // pod-a after a restart runs it: a tool call and an answer. + expect(result.restartedCalls).toBe(2) + expect(result.answers).toEqual([ + 'Read notes.txt on pod-a.', + 'Read notes.txt on pod-a again.', + ]) + }) +}) From be54212f03ddd1322b84a56a65c901a4b71af672 Mon Sep 17 00:00:00 2001 From: Alem Tuzlak Date: Sun, 18 Oct 2026 14:13:25 +0200 Subject: [PATCH 08/17] chore(ai-harness): allow the ps binary in knip, and give polled checks room for slower CI coverage runs --- knip.json | 2 +- packages/ai-harness/tests/coding-formatter.test.ts | 8 +++++++- packages/ai-harness/tests/coding-locks.test.ts | 10 ++++++++-- packages/ai-harness/tests/coding-snapshots.test.ts | 10 ++++++++-- packages/ai-harness/tests/workspace-binding.test.ts | 8 +++++++- packages/ai-harness/tests/workspace-tools.test.ts | 12 +++++++++--- 6 files changed, 40 insertions(+), 10 deletions(-) diff --git a/knip.json b/knip.json index bcaf1c98b0..cc250d19a7 100644 --- a/knip.json +++ b/knip.json @@ -5,7 +5,7 @@ "@aws-sdk/client-bedrock", "@typescript/typescript6" ], - "ignoreBinaries": ["taskkill"], + "ignoreBinaries": ["taskkill", "ps"], "ignoreWorkspaces": [ "examples/**", "testing/**", diff --git a/packages/ai-harness/tests/coding-formatter.test.ts b/packages/ai-harness/tests/coding-formatter.test.ts index 9b3394a1a7..2e6471cc32 100644 --- a/packages/ai-harness/tests/coding-formatter.test.ts +++ b/packages/ai-harness/tests/coding-formatter.test.ts @@ -21,6 +21,12 @@ import type { FormatterOptions, } from '../src/first-party/coding/formatter' +// Coverage runs in CI are slow, so polled checks wait up to 5 s. +const waitFor = ( + check: () => T | Promise, + options?: { timeout?: number; interval?: number }, +) => vi.waitFor(check, { timeout: 5_000, ...options }) + const { join, posix, resolve } = nodePath const root = resolve(tmpdir(), 'harness-format') const encoder = new TextEncoder() @@ -458,7 +464,7 @@ describe('formatter with permissions()', () => { * for the turn that resumes. Returns its message. */ async function answer(session: HarnessSession, value: unknown) { - await vi.waitFor( + await waitFor( () => expect(session.snapshot().pendingWaits).toHaveLength(1), { timeout: 5_000 }, ) diff --git a/packages/ai-harness/tests/coding-locks.test.ts b/packages/ai-harness/tests/coding-locks.test.ts index 76a0d9136f..4ea1dd558d 100644 --- a/packages/ai-harness/tests/coding-locks.test.ts +++ b/packages/ai-harness/tests/coding-locks.test.ts @@ -7,6 +7,12 @@ import { createWorkspaceTools } from '../src/first-party/coding/workspace' import { gate } from './helpers' import type { WorkspaceBackend } from '../src/first-party/coding/backend' +// Coverage runs in CI are slow, so polled checks wait up to 5 s. +const waitFor = ( + check: () => T | Promise, + options?: { timeout?: number; interval?: number }, +) => vi.waitFor(check, { timeout: 5_000, ...options }) + let root = '' beforeEach(async () => { @@ -95,7 +101,7 @@ describe('file locks across sessions', () => { ) // `reach` resolves the path once. The lock resolves it again only when // the first edit is done. - await vi.waitFor(() => expect(second.calls.realpath).toBe(before + 1)) + await waitFor(() => expect(second.calls.realpath).toBe(before + 1)) await settle() expect(second.calls.realpath).toBe(before + 1) expect(secondEdit.state.settled).toBe(false) @@ -141,7 +147,7 @@ describe('file locks across sessions', () => { ) // One `stat` before the lock sees no file. The check under the lock // waits for the first write. - await vi.waitFor(() => expect(second.calls.stat).toBe(1)) + await waitFor(() => expect(second.calls.stat).toBe(1)) await settle() expect(secondWrite.state.settled).toBe(false) diff --git a/packages/ai-harness/tests/coding-snapshots.test.ts b/packages/ai-harness/tests/coding-snapshots.test.ts index 914ba7462b..602aac39c9 100644 --- a/packages/ai-harness/tests/coding-snapshots.test.ts +++ b/packages/ai-harness/tests/coding-snapshots.test.ts @@ -30,6 +30,12 @@ import type { HarnessPersistence, HarnessSession, Principal } from '../src' import type { WorkspaceBackend } from '../src/first-party/coding/backend' import type { Reply } from './helpers' +// Coverage runs in CI are slow, so polled checks wait up to 5 s. +const waitFor = ( + check: () => T | Promise, + options?: { timeout?: number; interval?: number }, +) => vi.waitFor(check, { timeout: 5_000, ...options }) + const THREAD = 't1' /** The id of the log of `THREAD` in the store. */ const LOG = logIdOf({ logId: THREAD }) @@ -677,7 +683,7 @@ describe.skipIf(!hasGit)('session.revert', () => { await session.send('Hello') const [message] = await session.transcript() const running = session.send('Again') - await vi.waitFor(() => expect(session.snapshot().status).toBe('running')) + await waitFor(() => expect(session.snapshot().status).toBe('running')) expect(await session.revert(message?.id ?? '')).toMatchObject({ status: 'rejected', @@ -778,7 +784,7 @@ describe.skipIf(!hasGit)('session.revert keeps the edits of the user', () => { await reached.opened const prompting = session.send('Four') - await vi.waitFor(() => expect(session.snapshot().queuedTurns).toBe(1)) + await waitFor(() => expect(session.snapshot().queuedTurns).toBe(1)) expect(session.snapshot().status).toBe('idle') blocking = false wait.open() diff --git a/packages/ai-harness/tests/workspace-binding.test.ts b/packages/ai-harness/tests/workspace-binding.test.ts index 9f520dee34..0cd02eff7b 100644 --- a/packages/ai-harness/tests/workspace-binding.test.ts +++ b/packages/ai-harness/tests/workspace-binding.test.ts @@ -24,6 +24,12 @@ import type { WorkspaceToolsOptions, } from '../src/first-party/coding/workspace' +// Coverage runs in CI are slow, so polled checks wait up to 5 s. +const waitFor = ( + check: () => T | Promise, + options?: { timeout?: number; interval?: number }, +) => vi.waitFor(check, { timeout: 5_000, ...options }) + const THREAD = 't' const LOG = logIdOf({ logId: THREAD }) const SANDBOX = { @@ -285,7 +291,7 @@ describe('the workspace binding in the session log', () => { plugins: [workspaceTools({ ...inFolder, web: false })], }) const sent = (await other.opened).send('read') - await vi.waitFor( + await waitFor( async () => { expect(await log.offered({ logId: LOG })).toHaveLength(1) }, diff --git a/packages/ai-harness/tests/workspace-tools.test.ts b/packages/ai-harness/tests/workspace-tools.test.ts index c50b0d24f0..8c361781bf 100644 --- a/packages/ai-harness/tests/workspace-tools.test.ts +++ b/packages/ai-harness/tests/workspace-tools.test.ts @@ -31,6 +31,12 @@ import type { WorkspaceToolsOptions, } from '../src/first-party/coding/workspace' +// Coverage runs in CI are slow, so polled checks wait up to 5 s. +const waitFor = ( + check: () => T | Promise, + options?: { timeout?: number; interval?: number }, +) => vi.waitFor(check, { timeout: 5_000, ...options }) + const { posix } = nodePath /** @@ -902,7 +908,7 @@ describe('background bash in a session', () => { workspaceTools({ root: '/workspace', backend }), ]) await session.send('build') - await vi.waitFor(() => expect(calls).toHaveLength(3)) + await waitFor(() => expect(calls).toHaveLength(3)) expect(JSON.stringify(calls[2].messages)).toContain( 'Background job bash-1 ended.\\nexit code: 0\\nbuilt', ) @@ -938,7 +944,7 @@ describe('background bash in a session', () => { listener, ]) await session.send('build') - await vi.waitFor(() => expect(events).toHaveLength(3)) + await waitFor(() => expect(events).toHaveLength(3)) expect(events).toEqual([ ['started', { jobId: 'bash-1', command: 'make' }], ['output', { jobId: 'bash-1', output: 'built\n' }], @@ -1141,7 +1147,7 @@ describe('bash', () => { const pid = Number( await readFile(nodePath.join(cwd, 'child.pid'), 'utf8'), ) - await vi.waitFor(() => expect(isRunning(pid)).toBe(false), { + await waitFor(() => expect(isRunning(pid)).toBe(false), { timeout: 5000, }) } finally { From ce5a93e629a11d60d85df2025fe9e1502d4c0ff3 Mon Sep 17 00:00:00 2001 From: Alem Tuzlak Date: Mon, 19 Oct 2026 13:29:49 +0200 Subject: [PATCH 09/17] feat(ai, ai-harness): protocol v2 with one wire encoding, healthy streams, run tags, lifecycle events, usage by source, and tracing (#1821) * feat(ai-harness, ai): protocol v2 frames with error codes and request ids, observability types, and pull-based SSE and socket senders * docs(ai-harness): first draft of the harness protocol page * feat(ai-harness): one wire encoding on every transport with hello, version, request ids, and heartbeats; hooks and status from other hosts; a gap event for evicted feeds * test(e2e): the access spec skips the hello frame of the session stream * feat(ai-harness): a grouped stream option for the heartbeat, the slow-consumer close, and the SSE and socket buffer limits * feat(ai-harness): run tags on every event frame, ctx.model.chat for plugins, and usage by source * feat(ai-harness): lifecycle events, counted router and transfer model calls, tracing with span links, and onOperation reports * docs(ai-harness): the harness protocol page, keeping the stream open, stopping a turn from useChat, and the changesets for reporting and observability * docs(ai-harness): run tags and the gap reload in the custom UI page * test(e2e): the wire, the gap, the run tree in one stream, lifecycle events, tracing, remote status, cancel from another host, and an unsupported version * test(ai-harness): give polled checks room for slower CI coverage runs * feat(ai-mcp, ai-acp, ai-code-mode, ai-skills, ai-harness-cli, ai-harness): integrations on the new harness, with MCP reconnects, ACP order, a skills watcher per host, and a CLI that cancels at once (#1822) * feat(ai-mcp): mcp() reconnects with backoff, lists tools again on list_changed, and takes needsApproval and toolName per server * fix(ai-acp, ai-code-mode, ai-harness): ACP keeps prompt order and releases idle sessions, code mode never moves approval tools, harnessText acts for the caller * perf(ai-skills): one folder watcher per host with a cache of the skills * feat(ai-harness-cli): cancel and steer act at once in a terminal, loadOptional reports real errors, and --dashboard goes away * chore(changeset): the integration changes for mcp, acp, code mode, skills, and harnessText * perf(ai-harness, ai-harness-cli): cheaper plugin, log, media, view, client, and host-event paths, no idle store calls, and context prompts after the history (#1823) * perf(ai-harness): cache plugin lists and permission patterns, read instruction files async, parse SSE as a stream with one abort listener and capped backoff, and give each control input an id * perf(ai-harness): stream media, cache base64 parts across turns, cap open-ended ranges, and add a byte-capped event ring * perf(ai-harness, ai-harness-cli): the view and the printer only redo what changed, children load on demand with view.loadAgent, operation.stream reads from its start * perf(ai-harness): ring buffer log tail, joins and rebuilt operations read from the run start, a per-tenant host event bus, no MESSAGES_SNAPSHOT in the log, and no idle store calls * perf(ai-harness): plugin prompts with placement 'context' go after the history, so the cached prompt prefix stays valid * test(ai-harness): one threadId in the expected gap frame * test(e2e): the first-party fixtures match the context prompt that now comes after the user message * fix(ai, ai-harness): a guard tier for tool calls, complete plugin removal, media and body limits, signing keys, and turn and config fixes (#1824) * fix(ai-harness): a settled operation keeps its status, a media signer with key rotation, and agents() rule docs that match the code * fix(ai-harness): every plugin contribution has an owner, ctx.on cleans up, inline media goes to the media store, and capabilities report the socket only when mounted * docs(ai-harness): say when capabilities report the WebSocket * feat(ai): onGuardToolCall middleware hook that runs on the final arguments of every tool call * test(e2e): onGuardToolCall denies on the final args after a transformArgs * fix(ai-harness): tool turns take no steers, a turn keeps its config, removing a plugin removes all of it, body and frame limits, signing keys, and a warning on tool name clashes * fix(ai-harness): the media signer's key check covers an empty secret * fix(ai-harness-cli): serve signs media URLs with a key derived from the token, so a short token meets the 32-byte key rule * feat(ai-harness, ai-mcp, ai-acp, ai-code-mode, ai-harness-cli): the final review of the harness, with its fixes, one object per call, and the known limits (#1825) * fix(ai-harness, ai-mcp, ai-acp, ai-code-mode): final review fixes for durability, access, plugins, transports, integrations, and coding tools * fix(ai-acp): one agent app serves one connection, with editor capabilities read from that connection, and no private SDK field * fix(ai-harness): a thread delete ends it in first-party plugins, the bash spill is scrubbed, removed schedules go at once, agent run receipts resolve, and session.append takes one object with an optional inputId * refactor(ai-harness, ai-mcp, ai-harness-cli): session, client, and view methods take one object, the client and view use putMedia and client.sessions, and internal fields leave the public types * docs(ai-harness, ai-mcp, ai-acp, ai-code-mode): the final review changes, the fixed durable-tools and server-tools text, and a known limits section * fix(ai-harness): revoke and unshare reject when work does not stop, thread delete reaches child logs and known harnesses, plugins declare kept records, and the review fixes - host.revoke and host.sessions.unshare reject when the work stop fails (the access is gone already). unshare ends the reads before it stops the work. - createHarnessHost takes harnesses, so a delete on a new host ends the thread in its plugins. An unknown harness warns. A child thread ends in the log of its root. - A plugin declares keepRecords. The core fold reads them from the plugins of the harness that folds the log, and counts per type. snapshots() keeps its 500 steps this way. - revert and unrevert open the workspace of the thread from the log when no turn opened it on this host. - The harness.revert and harness.job records keep the principal. - append with no inputId goes through accept(), so its id is real. - A formatter config below the root, or a write under node_modules, asks again. - The base64 cache takes the largest encodedBytes. The client takes idleTimeoutMs. - An events() read listens for a revoke from the call, not the first pull. * docs(ai-harness): the harnesses option, keepRecords, idleTimeoutMs, a revoke that rejects, and one known limit less * feat(ai-harness): ctx.session.keptRecords reads the kept records of a plugin, and session.revert reads them the same way * fix(ai-harness): an always answer saves no rule for an @unsaved: action * refactor(ai-harness): mounted plugins keep each contribution list with its owner, not in parallel owner arrays * fix(ai-harness): a plugins remove setting is rejected when a kept turn plugin or a harness middleware needs the removed capability * fix(ai-harness): client media checks every part before it stores any, refuses provider file handles unless media.allowProviderFiles, and the frame limit docs name the server maxPayload * fix(ai-harness): a retried client input with inline media reuses the stored media, so it stores nothing again and is a duplicate, not a conflict * docs(ai-harness): routed root-agent turns take steers and late joins, only session.tool() turns do not * fix(ai-harness-cli): --serve signs media URLs with --media-key or HARNESS_MEDIA_KEY, or a random key, never with the bearer token * fix(ai-harness): keep the usage source in the view, and a host events read that listens from the call * fix(ai-harness): send context prompts as the last system prompt, not as a user message * fix(ai-harness): check the snapshot shape in client.events(), and document its head-first default * feat(ai-harness): an onTurnStart plugin hook, used by projectInstructions and workspaceTools * fix(ai-harness): keep the interrupt messages snapshot in the log, so a rejoin from the log gets it * fix(ai-harness-cli): load the saved child agents in the TUI example, and document view.loadAgent for a ui * fix(ai-mcp): no reconnect timer after close, backoff that survives crash loops, and onDuplicateTool 'error' * fix(ai-harness-cli): tell a missing optional package by its resolve error code, not the message text * fix(ai-harness): check the inner session of harnessText through the outer thread, and pass the run principal on as is * fix(ai-harness): bind the harnessText inner thread to its outer thread before the outer check applies * fix(ai-acp): keep the sessions of the connection known on a host with no session index * refactor(ai-skills): one shared watcher per host that only triggers a sync, with no skill cache * fix(ai-harness): one harnessText inner thread id shape, /inner/, on every host * fix(ai-harness): use the real OpenTelemetry types, guard a tracer that throws, and trace a transfer inside it @opentelemetry/api is now an optional peer (type imports, and a lazy load for trace.setSpan when a host has a tracer). The copied span types and the private context key are gone. A tracer or span call that throws is logged and fails no run. The harness.transfer span is part of host.transfer, not a patch after it. * fix(ai-harness): trim the protocol to one version check, put the parent run on RUN_STARTED, and surface harness.error in the client - A requested version other than 2 gets unsupported_version from one equality check. negotiateVersion, SUPPORTED_PROTOCOL_VERSIONS, and the supported field are gone. - harness.event frames and session events no longer carry parentRunId. The AG-UI RUN_STARTED of a child agent run has it. - client.events() and client.hostEvents() connect again after slow_consumer and throw a HarnessProtocolError (code + message) for any other error frame. * fix(ai-harness): no usage event after a run finished, and bySource on old snapshots - A ctx.model.chat call that ends after its run settled (the title plugin) writes the usage record only: the run sends nothing after RUN_FINISHED. - Usage totals from a snapshot written before bySource get an empty map at load, so session.usage().bySource is always an object. * fix(ai-harness, ai-persistence): read takeover and open questions from real state, and clear the runs of a stopped owner - LogStore.acquire returns expired: true when it takes a log from another owner whose ownership expired with no release. harness.ownership.changed says taken_over from that, not from pending inputs. - The host status reads open questions from the folded log only. - The status rule drops the running operations of a stopped owner on taken_over, and the old operation of a resumed input, so a reader host does not show running forever after a takeover. * test(ai-harness): an append is an input, so a non-owner forwards it and the log has its input records - shared-log: check that the two app records land in one store append, or neither, among the input records that the append now writes. - ownership: the append of a host that does not own the log is queued and the owner writes it. A write that is not an input (a schedule) is still refused. - docs: an append on a host that does not own the thread gets a queued receipt. --- .changeset/ai-harness-cli.md | 6 +- .changeset/ai-harness-coding-tools.md | 8 +- .changeset/ai-harness-core.md | 35 +- .changeset/ai-harness-plugins.md | 7 +- .changeset/ai-harness-transport.md | 16 +- .changeset/ai-harness-view.md | 5 +- .changeset/ai-mcp-harness.md | 5 +- .changeset/ai-sandbox-harness.md | 2 +- .changeset/ai-tool-guard.md | 11 + .changeset/ai-websocket-buffered-amount.md | 5 + .changeset/code-mode-harness.md | 2 +- .changeset/harness-acp-agent.md | 4 +- .changeset/harness-cli-handler-options.md | 2 +- .changeset/harness-join-any-instance.md | 2 +- .changeset/harness-secure-defaults.md | 3 +- .changeset/harness-skills.md | 2 +- .changeset/harness-tool-ops.md | 2 +- .changeset/harness-turn-and-body-limits.md | 14 + .changeset/persistence-session-stores.md | 2 +- docs/advanced/middleware.md | 33 + docs/config.json | 49 +- docs/harness/access.md | 15 +- docs/harness/acp.md | 7 +- docs/harness/agents.md | 13 +- docs/harness/cli.md | 15 +- docs/harness/code-mode.md | 3 +- docs/harness/coding-tools-backends.md | 3 +- docs/harness/coding-tools.md | 4 +- docs/harness/coding-workspaces.md | 10 +- docs/harness/connect.md | 136 ++- docs/harness/custom-ui.md | 48 +- docs/harness/deploy.md | 76 ++ docs/harness/durable-sessions.md | 9 +- docs/harness/durable-tools.md | 2 +- docs/harness/goal.md | 2 + docs/harness/handoffs.md | 2 +- docs/harness/inputs.md | 12 +- docs/harness/lifecycle.md | 28 +- docs/harness/mcp.md | 69 +- docs/harness/media.md | 80 +- docs/harness/overview.md | 2 + docs/harness/permissions.md | 8 +- docs/harness/plugins.md | 99 +- docs/harness/protocol.md | 480 ++++++++++ docs/harness/session-log.md | 33 +- docs/harness/sessions.md | 6 +- docs/harness/shared-logs.md | 18 +- docs/harness/shared-threads.md | 6 +- docs/harness/skills.md | 6 +- docs/harness/snapshots.md | 14 +- docs/harness/subagents.md | 42 +- docs/harness/thread-settings.md | 9 +- docs/harness/turn-control.md | 4 + docs/harness/usage.md | 49 +- docs/persistence/store-reference.md | 5 +- docs/tools/server-tools.md | 2 +- examples/harness-cli/src/tui.tsx | 32 +- packages/ai-acp/src/agent/index.ts | 187 ++-- packages/ai-acp/tests/agent-handlers.test.ts | 126 ++- packages/ai-acp/tests/agent.test.ts | 82 ++ packages/ai-code-mode/src/harness.ts | 52 +- packages/ai-code-mode/tests/harness.test.ts | 99 +- packages/ai-harness-cli/src/args.ts | 10 +- packages/ai-harness-cli/src/commands.ts | 11 +- packages/ai-harness-cli/src/index.ts | 47 +- packages/ai-harness-cli/src/lines.ts | 57 +- packages/ai-harness-cli/src/optional.ts | 63 +- packages/ai-harness-cli/src/printer.ts | 13 +- packages/ai-harness-cli/src/serve.ts | 19 +- packages/ai-harness-cli/tests/attach.test.ts | 2 +- packages/ai-harness-cli/tests/cli.test.ts | 148 ++- packages/ai-harness-cli/tests/helpers.ts | 46 +- .../ai-harness-cli/tests/optional.test.ts | 70 +- packages/ai-harness-cli/tests/printer.test.ts | 35 + packages/ai-harness/package.json | 5 + packages/ai-harness/src/agents.ts | 26 +- packages/ai-harness/src/build.ts | 26 +- packages/ai-harness/src/client.ts | 570 +++++++++--- packages/ai-harness/src/commands.ts | 30 +- packages/ai-harness/src/define.ts | 24 +- packages/ai-harness/src/feed.ts | 61 +- packages/ai-harness/src/first-party/agents.ts | 59 +- .../ai-harness/src/first-party/coding/bash.ts | 29 +- .../src/first-party/coding/formatter.ts | 205 ++++- .../src/first-party/coding/snapshots.ts | 270 ++++-- .../src/first-party/coding/workspace.ts | 198 +++- packages/ai-harness/src/first-party/files.ts | 54 +- packages/ai-harness/src/first-party/goal.ts | 63 +- .../ai-harness/src/first-party/permissions.ts | 50 +- .../src/first-party/session-tools.ts | 19 +- packages/ai-harness/src/first-party/title.ts | 16 +- packages/ai-harness/src/first-party/todos.ts | 2 + packages/ai-harness/src/harness-text.ts | 351 +++++-- packages/ai-harness/src/host-event-bus.ts | 136 +++ packages/ai-harness/src/host.ts | 863 +++++++++++++----- packages/ai-harness/src/http.ts | 642 +++++++++---- packages/ai-harness/src/index.ts | 2 + packages/ai-harness/src/inline-media.ts | 239 +++++ packages/ai-harness/src/log.ts | 346 +++++-- packages/ai-harness/src/media-signing.ts | 114 +++ packages/ai-harness/src/media.ts | 121 ++- packages/ai-harness/src/observability.ts | 380 ++++++++ packages/ai-harness/src/operation.ts | 26 +- packages/ai-harness/src/ownership.ts | 75 +- packages/ai-harness/src/plugins.ts | 547 +++++++++-- packages/ai-harness/src/protocol.ts | 305 ++++++- packages/ai-harness/src/ring.ts | 150 +++ packages/ai-harness/src/server.ts | 1 + packages/ai-harness/src/session/agents.ts | 142 ++- packages/ai-harness/src/session/base.ts | 169 +++- packages/ai-harness/src/session/handle.ts | 162 +++- .../ai-harness/src/session/harness-session.ts | 68 ++ packages/ai-harness/src/session/helpers.ts | 139 ++- packages/ai-harness/src/session/inputs.ts | 242 ++++- .../src/session/observed-session.ts | 68 +- packages/ai-harness/src/session/requests.ts | 134 ++- packages/ai-harness/src/session/state.ts | 406 ++++++-- packages/ai-harness/src/session/turn-parts.ts | 170 +++- packages/ai-harness/src/session/turn.ts | 99 +- packages/ai-harness/src/session/types.ts | 3 + packages/ai-harness/src/sse.ts | 361 ++++++++ packages/ai-harness/src/types.ts | 178 +++- packages/ai-harness/src/usage.ts | 32 +- packages/ai-harness/src/view/index.ts | 329 +++++-- packages/ai-harness/src/view/reduce.ts | 37 +- packages/ai-harness/src/view/types.ts | 2 +- packages/ai-harness/src/wire.ts | 27 + packages/ai-harness/src/worker.ts | 77 +- packages/ai-harness/tests/access-work.test.ts | 385 ++++++++ packages/ai-harness/tests/agent-group.test.ts | 4 +- .../ai-harness/tests/agent-limits.test.ts | 4 +- .../ai-harness/tests/agent-resume.test.ts | 4 +- packages/ai-harness/tests/agent-runs.test.ts | 69 +- packages/ai-harness/tests/agents.test.ts | 140 ++- packages/ai-harness/tests/auth-wait.test.ts | 2 +- .../ai-harness/tests/background-tools.test.ts | 96 +- packages/ai-harness/tests/build-stub.test.ts | 10 + .../ai-harness/tests/client-media.test.ts | 12 +- .../ai-harness/tests/client-parse.test.ts | 345 +++++++ .../ai-harness/tests/client-reads.test.ts | 12 +- packages/ai-harness/tests/client.test.ts | 85 +- packages/ai-harness/tests/coding-bash.test.ts | 17 + .../ai-harness/tests/coding-formatter.test.ts | 120 ++- packages/ai-harness/tests/coding-jobs.test.ts | 34 + .../ai-harness/tests/coding-snapshots.test.ts | 132 ++- .../ai-harness/tests/discover-tools.test.ts | 104 +++ .../tests/durable-background.test.ts | 2 +- .../ai-harness/tests/durable-inputs.test.ts | 29 +- .../ai-harness/tests/durable-session.test.ts | 70 +- packages/ai-harness/tests/feed.test.ts | 47 +- packages/ai-harness/tests/first-party.test.ts | 40 +- packages/ai-harness/tests/fork.test.ts | 6 +- packages/ai-harness/tests/gap.test.ts | 162 ++++ packages/ai-harness/tests/goal.test.ts | 83 +- .../ai-harness/tests/harness-text.test.ts | 381 +++++++- packages/ai-harness/tests/helpers.ts | 46 +- .../ai-harness/tests/host-event-bus.test.ts | 120 +++ packages/ai-harness/tests/http-access.test.ts | 79 +- packages/ai-harness/tests/http-media.test.ts | 193 +++- packages/ai-harness/tests/idle-budget.test.ts | 76 ++ .../ai-harness/tests/inline-media.test.ts | 266 ++++++ packages/ai-harness/tests/input-lane.test.ts | 242 ++++- .../ai-harness/tests/lifecycle-events.test.ts | 312 +++++++ packages/ai-harness/tests/log-stores.test.ts | 30 +- packages/ai-harness/tests/log.test.ts | 101 +- .../ai-harness/tests/media-signing.test.ts | 82 ++ packages/ai-harness/tests/media.test.ts | 134 +++ .../ai-harness/tests/observability.test.ts | 130 +++ packages/ai-harness/tests/on-event.test.ts | 40 + .../tests/operation-settled.test.ts | 73 ++ packages/ai-harness/tests/ownership.test.ts | 66 +- packages/ai-harness/tests/permissions.test.ts | 74 +- .../ai-harness/tests/plugin-hotpaths.test.ts | 157 ++++ .../ai-harness/tests/plugin-model.test.ts | 268 ++++++ .../ai-harness/tests/plugin-owners.test.ts | 238 +++++ .../ai-harness/tests/prompt-placement.test.ts | 124 +++ packages/ai-harness/tests/protocol.test.ts | 322 ++++++- packages/ai-harness/tests/redaction.test.ts | 52 +- .../ai-harness/tests/remote-hooks.test.ts | 144 +++ .../ai-harness/tests/remote-worker.test.ts | 82 +- .../ai-harness/tests/resume-pending.test.ts | 21 +- packages/ai-harness/tests/retention.test.ts | 9 +- packages/ai-harness/tests/ring.test.ts | 116 +++ packages/ai-harness/tests/routing.test.ts | 51 ++ .../ai-harness/tests/run-route-hosts.test.ts | 67 +- packages/ai-harness/tests/run-route.test.ts | 29 + packages/ai-harness/tests/run-tags.test.ts | 207 +++++ .../ai-harness/tests/session-media.test.ts | 8 +- .../ai-harness/tests/session-plugins.test.ts | 75 +- .../ai-harness/tests/session-surface.test.ts | 21 + packages/ai-harness/tests/shared-log.test.ts | 68 +- packages/ai-harness/tests/snapshots.test.ts | 134 ++- packages/ai-harness/tests/sse.test.ts | 320 +++++++ packages/ai-harness/tests/status-feed.test.ts | 103 ++- .../tests/subagent-children.test.ts | 17 +- .../ai-harness/tests/thread-settings.test.ts | 92 +- packages/ai-harness/tests/title-usage.test.ts | 31 + packages/ai-harness/tests/tool-op.test.ts | 45 +- packages/ai-harness/tests/tracing.test.ts | 487 ++++++++++ packages/ai-harness/tests/transfer.test.ts | 17 +- .../ai-harness/tests/transport-edges.test.ts | 74 +- .../ai-harness/tests/transport-wire.test.ts | 382 ++++++++ .../ai-harness/tests/turn-control.test.ts | 8 +- packages/ai-harness/tests/types.test-d.ts | 8 + packages/ai-harness/tests/usage.test.ts | 44 + .../tests/view-client-tools.test.ts | 3 +- packages/ai-harness/tests/view-fixtures.ts | 2 +- packages/ai-harness/tests/view-media.test.ts | 2 +- packages/ai-harness/tests/view-paths.test.ts | 203 ++-- packages/ai-harness/tests/view.test.ts | 38 +- packages/ai-harness/tests/wire.test.ts | 73 ++ .../ai-harness/tests/workspace-tools.test.ts | 157 +++- packages/ai-mcp/src/client.ts | 22 +- packages/ai-mcp/src/connector.ts | 45 +- packages/ai-mcp/src/harness-plugin.ts | 406 ++++++-- packages/ai-mcp/src/harness.ts | 20 +- packages/ai-mcp/src/types.ts | 18 + packages/ai-mcp/tests/client.test.ts | 38 + .../ai-mcp/tests/connector-reconnect.test.ts | 97 ++ packages/ai-mcp/tests/harness-elicit.test.ts | 32 +- packages/ai-mcp/tests/harness-plugin.test.ts | 123 ++- .../ai-mcp/tests/harness-reconnect.test.ts | 268 ++++++ packages/ai-mcp/tests/harness.test.ts | 30 + packages/ai-persistence/src/memory.ts | 5 +- .../ai-persistence/src/testkit/conformance.ts | 2 +- packages/ai-persistence/src/types.ts | 4 +- packages/ai-persistence/tests/memory.test.ts | 2 +- packages/ai-skills/src/harness.ts | 596 ++++++------ packages/ai-skills/tests/harness.test.ts | 176 +++- packages/ai/src/activities/chat/index.ts | 8 + .../src/activities/chat/middleware/compose.ts | 38 + .../src/activities/chat/middleware/index.ts | 1 + .../src/activities/chat/middleware/types.ts | 38 +- .../src/activities/chat/tools/tool-calls.ts | 57 +- packages/ai/src/index.ts | 1 + packages/ai/src/stream-to-websocket.ts | 6 + packages/ai/tests/middleware.test.ts | 192 ++++ packages/ai/tests/stream-to-websocket.test.ts | 10 +- pnpm-lock.yaml | 3 + testing/e2e/src/routeTree.gen.ts | 22 + testing/e2e/src/routes/api.coding-tools.ts | 2 +- .../e2e/src/routes/api.harness-forwarding.ts | 241 ++++- .../src/routes/api.harness-observability.$.ts | 475 ++++++++++ .../e2e/src/routes/api.harness-protocol.$.ts | 2 +- testing/e2e/src/routes/api.middleware-test.ts | 29 + testing/e2e/src/routes/middleware-test.tsx | 1 + testing/e2e/tests/harness-access.spec.ts | 25 +- testing/e2e/tests/harness-forwarding.spec.ts | 46 + .../e2e/tests/harness-observability.spec.ts | 382 ++++++++ testing/e2e/tests/harness-protocol.spec.ts | 30 + testing/e2e/tests/middleware.spec.ts | 36 + 251 files changed, 21127 insertions(+), 2641 deletions(-) create mode 100644 .changeset/ai-tool-guard.md create mode 100644 .changeset/ai-websocket-buffered-amount.md create mode 100644 .changeset/harness-turn-and-body-limits.md create mode 100644 docs/harness/protocol.md create mode 100644 packages/ai-harness/src/host-event-bus.ts create mode 100644 packages/ai-harness/src/inline-media.ts create mode 100644 packages/ai-harness/src/media-signing.ts create mode 100644 packages/ai-harness/src/observability.ts create mode 100644 packages/ai-harness/src/ring.ts create mode 100644 packages/ai-harness/src/sse.ts create mode 100644 packages/ai-harness/src/wire.ts create mode 100644 packages/ai-harness/tests/access-work.test.ts create mode 100644 packages/ai-harness/tests/client-parse.test.ts create mode 100644 packages/ai-harness/tests/gap.test.ts create mode 100644 packages/ai-harness/tests/host-event-bus.test.ts create mode 100644 packages/ai-harness/tests/idle-budget.test.ts create mode 100644 packages/ai-harness/tests/inline-media.test.ts create mode 100644 packages/ai-harness/tests/lifecycle-events.test.ts create mode 100644 packages/ai-harness/tests/media-signing.test.ts create mode 100644 packages/ai-harness/tests/observability.test.ts create mode 100644 packages/ai-harness/tests/operation-settled.test.ts create mode 100644 packages/ai-harness/tests/plugin-hotpaths.test.ts create mode 100644 packages/ai-harness/tests/plugin-model.test.ts create mode 100644 packages/ai-harness/tests/plugin-owners.test.ts create mode 100644 packages/ai-harness/tests/prompt-placement.test.ts create mode 100644 packages/ai-harness/tests/remote-hooks.test.ts create mode 100644 packages/ai-harness/tests/ring.test.ts create mode 100644 packages/ai-harness/tests/run-tags.test.ts create mode 100644 packages/ai-harness/tests/sse.test.ts create mode 100644 packages/ai-harness/tests/tracing.test.ts create mode 100644 packages/ai-harness/tests/transport-wire.test.ts create mode 100644 packages/ai-harness/tests/wire.test.ts create mode 100644 packages/ai-mcp/tests/connector-reconnect.test.ts create mode 100644 packages/ai-mcp/tests/harness-reconnect.test.ts create mode 100644 testing/e2e/src/routes/api.harness-observability.$.ts create mode 100644 testing/e2e/tests/harness-observability.spec.ts diff --git a/.changeset/ai-harness-cli.md b/.changeset/ai-harness-cli.md index 330bd2d3ab..c8cc1c3e07 100644 --- a/.changeset/ai-harness-cli.md +++ b/.changeset/ai-harness-cli.md @@ -4,9 +4,9 @@ Add `@tanstack/ai-harness-cli`: `runCli(harness)` runs a harness from the terminal. -- Line mode reads one message or `/command` per line, hides a secret answer while the user types, and sends `@path` files. When a turn waits, each next line answers one wait: `y` or `n` for an approval (`once`, `always`, or `reject` for an approval of `permissions()`), and an answer read with the schema for a question. Pass `ui` to show your own screen on a session view instead. +- Line mode reads one message or `/command` per line, hides a secret answer while the user types, and sends `@path` files. In a terminal, `/cancel` and a message typed while a turn runs act at once: the message joins the running turn. When a turn waits, each next line answers one wait: `y` or `n` for an approval (`once`, `always`, or `reject` for an approval of `permissions()`), and an answer read with the schema for a question. Pass `ui` to show your own screen on a session view instead. - `-p "prompt"` runs one prompt for scripts and CI, and `--output ndjson` prints every AG-UI event. `-p` exits with code 3 (`EXIT.needsAction`) at the first wait of the turn: an approval, a question, or a sign-in. The exit code for this case changed from 2 to 3. A script that checks for exit code 2 must check for 3. - `--serve` serves the session protocol over HTTP behind a bearer token. -- `--acp`, `--mcp`, and `--dashboard` load `@tanstack/ai-acp`, `@tanstack/ai-mcp`, and `@tanstack/ai-dashboard` only when you use the flag. A missing package, or one without harness support, fails that flag with a clear message. +- `--acp` and `--mcp` load `@tanstack/ai-acp` and `@tanstack/ai-mcp` only when you use the flag. A missing package, or one without harness support, fails that flag with a clear message. Any other error while the package loads is reported, not hidden. The `--dashboard` flag was removed. - In every mode, the first SIGINT or SIGTERM cancels the running turn. A second one, or one while no turn runs, stops the CLI, and `runCli` resolves to 130 (SIGINT) or 143 (SIGTERM). The CLI closes a host that it made with `recoverable: true`, or releases its session handle on your own host. The `signals` option takes another signal source, for example an `EventEmitter` in a test. -- `host` and `principal` options let the CLI use your durable host and run as your user. Every mode except `--dashboard` opens its sessions as that principal. On a host with `access`, `--serve` reads the user of each request with `identify`, and `--mcp` does not start, because stdio has no login. +- `host` and `principal` options let the CLI use your durable host and run as your user. Every mode opens its sessions as that principal. On a host with `access`, `--serve` reads the user of each request with `identify`, and `--mcp` does not start, because stdio has no login. diff --git a/.changeset/ai-harness-coding-tools.md b/.changeset/ai-harness-coding-tools.md index 60394b4675..85f989ea92 100644 --- a/.changeset/ai-harness-coding-tools.md +++ b/.changeset/ai-harness-coding-tools.md @@ -5,15 +5,15 @@ Add coding tools in the new Node entry `@tanstack/ai-harness/plugins/coding`: - `workspaceTools({ root })` gives `read_file` (text pages, images, PDFs), `write_file`, `edit_file` or `patch` (`editStyle`), `list_files`, `grep`, `bash` (with `background`), `webfetch`, and `websearch` (with a `SearchProvider`). With `permissions()`, edits, `bash`, and `webfetch` ask first, and the rules also apply to the real path of a link. `outside: 'ask'` makes a path outside `root` need an approval. `permissions()` decides outside paths and the real paths of links before the call runs, so a tool never asks while it runs. The `cwd` thread setting gives each thread its own working folder. -- `workspaceTools({ workspace: { bind, reattach? } })` gives each thread its own workspace. The root of the workspace is the boundary of every tool. The session log keeps where the workspace is (`harness.workspace.bound`). On another host, `reattach` gives the workspace back (for example a resumed sandbox), or `'pinned'` or `'gone'`. A gone workspace sends `tanstack/workspace-tools:stranded`. +- `workspaceTools({ workspace: { bind, reattach? } })` gives each thread its own workspace. The root of the workspace is the boundary of every tool. The session log keeps where the workspace is (`harness.workspace.bound`). On another host, `reattach` gives the workspace back (for example a resumed sandbox), or `'pinned'` or `'gone'`. A gone workspace sends `tanstack/workspace-tools:stranded`. With this plugin first, `snapshots()` and `formatter()` use the workspace of each thread, and their `root` is optional. - `createHarnessHost({ ownership: { hostId } })` gives a host a stable id. Then a thread that uses a folder on that host is pinned to it, and other hosts refuse its tool calls. Without `hostId`, threads are not pinned. - Commands, `rg`, `git`, and the formatters get only a short list of variables of the server process, plus the `env` of the backend call. Git runs without the hooks and the `fsmonitor` of the project. - The tools resolve links with `realpath` under the file lock, and fail when a backend cannot resolve a path. A write fails when the file changed after the tool read it. The file locks are shared by the sessions of a host. -- Background jobs: at most 8 per session (`maxJobs`), the newest 1 MiB of output per job, and the events `tanstack/workspace-tools:job-started`, `job-output`, and `job-ended`. A job that an earlier host left running is stopped at the next tool call, with a note to the model. +- Background jobs: at most 8 per session (`maxJobs`), the newest 1 MiB of output per job, and the events `tanstack/workspace-tools:job-started`, `job-output`, and `job-ended`. A job that an earlier host left running is stopped at the next tool call, with a note to the model. The wake turn of a job runs as the user who started it. The `bash` spill file has `[redacted]` in place of each secret of the session. - `read_file` sends images and PDFs up to 256 KB. `webfetch` gives the model the first 256 KB of a page. - A `WorkspaceBackend` runs the tools. `hostBackend` is the default. `realpath` is required. `spawn` can give a `jobId`, and `killJob`, `binding`, and `close` are optional. On Windows, commands do not run a program from the working folder by its name. -- `snapshots()` adds `/undo` and `/redo`, which restore only the files of the turn (git 2.26 or newer), file restore for `session.revert()`, and a diff for each model step. Each thread of each tenant has its own shadow repository, which never stores `.env` files. A step keeps only the files of its own run. The steps and the undo point are in the session log, so another host with the same repository can undo. A revert uses the last 500 steps. `purgeSnapshots()` removes the repository of a thread. A failed snapshot sends `tanstack/snapshots:failed`. -- `formatter()` formats each written file with the formatter of the project. With `permissions()`, the rules decide the formatter with the call that writes, before that call runs: one approval covers both, and the formatter runs after the write. +- `snapshots()` adds `/undo` and `/redo`, which restore only the files of the turn (git 2.26 or newer), file restore for `session.revert()`, and a diff for each model step. Each thread of each tenant has its own shadow repository, which never stores `.env` files. A step keeps only the files of its own run. The steps and the undo point are in the session log, so another host with the same repository can undo. A revert uses the last 500 steps. `host.sessions.delete` removes the repository of the thread, and `purgeSnapshots()` does it from your own code. A failed snapshot sends `tanstack/snapshots:failed`. +- `formatter()` formats each written file with the formatter of the project. With `permissions()`, the rules decide the formatter with the call that writes, before that call runs: one approval covers both, and the formatter runs after the write. An `always` answer is tied to a hash of the formatter config files. A write of a formatter config below the root, or under `node_modules`, asks each time, and an `always` answer saves no rule for that formatter run. - `WorkspaceHooks` lets other plugins run code after a read or a write. `@vscode/ripgrep` and `turndown` are optional peer dependencies. diff --git a/.changeset/ai-harness-core.md b/.changeset/ai-harness-core.md index a82ababc2f..b49a7da12a 100644 --- a/.changeset/ai-harness-core.md +++ b/.changeset/ai-harness-core.md @@ -6,20 +6,22 @@ Add `@tanstack/ai-harness`: build an agent harness on TanStack AI. A harness kee - `defineHarness()` names an agent: its adapter, tools, system prompts, middleware, agents, subagents, routing, plugins, turn hooks, and durability options. - `createHarnessHost()` opens sessions. Each `host.open` gives its own handle to the live session of the thread. `session.close()`, or `await using`, releases the handle. A dropped handle is released by the garbage collector, with one warning. A call on a handle after its session closed opens the session again. -- `lifecycle` on `createHarnessHost()` sets how sessions live. A session that no handle holds closes after it was idle for `idleCloseMs` (default 60 000). `maxOpen` (default 500) first closes the least recently used idle sessions. `retention` (`settledOperations`, `settledAgentRuns`, `eventBytes`) bounds the settled operations, agent runs, receipts, events, and log tail that each session keeps in memory. `host.close({ recoverable })` stops the work of every session and refuses later opens. With `recoverable: true`, the host gives each log that has work back with a wake, so a worker takes it at once. +- `lifecycle` on `createHarnessHost()` sets how sessions live. A session that no handle holds closes after it was idle for `idleCloseMs` (default 60 000). `maxOpen` (default 500) first closes the least recently used idle sessions. A session is not idle while a background job of this host runs. `retention` (`settledOperations`, `settledAgentRuns`, `eventBytes`) bounds the settled operations, agent runs, receipts, events, and log tail that each session keeps in memory. The log keeps the `MESSAGES_SNAPSHOT` that core sends when a run pauses at an interrupt, so a rejoin from the log rebuilds the paused run. `host.close({ recoverable })` stops the work of every session and refuses later opens. With `recoverable: true`, the host gives each log that has work back with a wake, so a worker takes it at once. - `host.observe(harness, { threadId, principal })` reads a thread with no session: its events, transcript, snapshot, media, usage, and entry. It sets up no plugins, writes nothing, and never takes the log, so a read works while another host owns the thread. `view.logState()` reads the `durability.reduce` fold. -- `host.sessions.delete` closes the sessions of a thread and of its child threads. Then it deletes the log of the session tree (with the settings, config, plugin state, interrupted turn, runs, interrupts, usage, revert data, and the inputs that wait), the index entries, the uploaded media, the generated artifacts, and the log checkpoint. +- `host.sessions.delete` closes the sessions of a thread and of its child threads. Then it deletes the log of the session tree (with the settings, config, plugin state, interrupted turn, runs, interrupts, usage, revert data, and the inputs that wait), the index entries, the uploaded media, the generated artifacts, and the log checkpoint. The first-party plugins also end the thread: `workspaceTools()` closes an owned sandbox with `'closed'`, and `snapshots()` removes the shadow repository. This needs the index entry, and the harness on that host: a host that opened it, or one with it in the `harnesses` option of `createHarnessHost()`. Else the host warns. A child thread is ended in the log of its root. - A session takes `send` (a prompt, a steer into the running turn, a follow-up, or a message to an agent run with `to`), `continue`, `reset`, `resolve`, `cancel`, and `background` inputs. Each returns an operation with a receipt, and the session streams typed events. `host.isOpen(harness, { threadId })` says whether this host holds that session open. -- `createHarnessHost({ access })` decides who can use each session. `identify({ request })` reads the user of a request, and `authorize({ principal, action, session, run })` decides each action. The principal type comes from what `identify` returns. The host keeps each tenant apart, needs a `principal` on each call, and `host.open` returns a handle that acts only as its principal. A refused action throws `AccessDeniedError`. Without `access`, one user runs the host, with no checks. With `access`, keyed adapters and provider keys read env vars only when `access.envKeys` is `true`. -- `membersOnly()` is a ready `authorize` with the roles owner, member, and viewer. `host.sessions.share` and `host.sessions.unshare` change the members, and `host.sessions.fork` takes `members: 'self' | 'all'`. For apps with their own access tables, `host.revoke` takes the access of a user to a thread. `host.identify` reads the user of a request for your own routes. Each `host.sessions` method takes one object. +- The session methods take one object: `cancelInput({ inputId, principal? })`, `revert({ messageId, principal? })`, `background({ toolCallId?, principal? })`, `getMedia({ id })`, `mediaUrl({ id })`, and `agentTranscript({ subagentRunId })`. +- `session.append({ records, inputId?, principal? })` returns a `Receipt`. A second append with the same `inputId` writes nothing. Without `inputId`, the append gets a new input id, and the log keeps it as an input. The `harness.revert` record keeps the principal of `revert`, and the `harness.job` record keeps the principal of `background`. +- `createHarnessHost({ access })` decides who can use each session. `identify({ request })` reads the user of a request, and `authorize({ principal, action, session, run })` decides each action. The principal type comes from what `identify` returns. The host keeps each tenant apart, needs a `principal` on each call, and `host.open` returns a handle that acts only as its principal. A refused action throws `AccessDeniedError`. The operations and the agent runs that a handle returns check access too. Without `access`, one user runs the host, with no checks. With `access`, keyed adapters and provider keys read env vars only when `access.envKeys` is `true`. +- `membersOnly()` is a ready `authorize` with the roles owner, member, and viewer. `host.sessions.share` and `host.sessions.unshare` change the members, and `host.sessions.fork` takes `members: 'self' | 'all'`. For apps with their own access tables, `host.revoke` takes the access of a user to a thread, and returns a `Promise`. `host.revoke` and `unshare` cancel the running and queued work of the user on every host, and remove the schedules of the user at once. When that work cannot be stopped, they reject, and the access is gone already. A running `events()` read of a handle or a view ends with `AccessDeniedError`, and later view calls throw. `host.identify` reads the user of a request for your own routes. Each `host.sessions` method takes one object. - A principal with a `tenantId` keeps its sessions in its tenant: one thread id in two tenants is two sessions, with two index entries and two logs. `logIdOf({ tenantId, logId })` gives the id of a log in the store. -- Every run has a run context with its user, and that user's credentials and provider keys: `ctx.context.run` in tools and middleware (type them with `HarnessRunContext`), `ctx.turn.run` in a turn plugin, and `run` in a command. The `ctx.credentials` and `ctx.keys` of a session plugin have no user, so each call throws. `ctx.agents.start(agent, input, { principal })` starts an agent for a user. The log keeps the sender of every input. The run context also has `epoch`, the ownership epoch of the log, and `idempotencyKey(step, scope?)`, a key for one side effect that is the same when the run runs again on any host. -- `definePlugin()` adds tools, prompts, middleware, agents, commands, config options, state, and extension points to a session or to one turn. With `lifetime: 'host'`, the host sets a plugin up once per harness and shares it with every session. It is cleaned up when the last session that uses it closes, or at `host.close()`. +- Every run has a run context with its user, and that user's credentials and provider keys: `ctx.context.run` in tools and middleware (type them with `HarnessRunContext`), `ctx.turn.run` in a turn plugin, and `run` in a command. The `ctx.credentials` and `ctx.keys` of a session plugin have no user, so each call throws. `ctx.agents.start(agent, input, { principal })` starts an agent for a user. `AgentRun.send` takes `principal?`, and the `send` and `cancel` of a run act as the user of the run. `session.agents..run` honors `principal`. The log keeps the sender of every input. The run context also has `epoch`, the ownership epoch of the log, and `idempotencyKey(step, scope?)`, a key for one side effect that is the same when the run runs again on any host. +- `definePlugin()` adds tools, prompts, middleware, agents, commands, config options, state, and extension points to a session or to one turn. With `lifetime: 'host'`, the host sets a plugin up once per harness and shares it with every session. It is cleaned up when the last session that uses it closes, or at `host.close()`. A prompt with `placement: 'context'` goes in one system prompt after all the others, made again at each model call, so text that changes often (the todo list, the goal status, plan mode) leaves the system prompts before it in the provider prompt cache. The default `placement: 'system'` keeps a prompt in the system prompts. `prepareTools` gets `run`. `onTurnStart(turn)` runs once at the start of each turn, before the prompts, so a plugin can refresh the state that its prompts read. A plugin's `ctx.session.send` and `note` take `principal?`, and default to the session owner. A thread that removes a plugin cannot run the agents of that plugin by name. - `routing.router` sends a turn to a root agent. The log keeps the active agent, and the router reads it as `activeAgent`. A router that returns `undefined` gives the turn to the active agent, else to the main model. One agent name makes that agent active, `'main'` clears it, and `{ handoff, note }` is a handoff from the router. `routing.after: 'end' | 'main'` (default `'end'`) says whether the main model answers after the routed agents. - Handoffs move a conversation to another root agent, and the target answers in the same turn. With `routing`, the main model and each root agent get a `handoff` tool. Code hands off with `ctx.harness.handoff({ to, note })` in a root agent, and with `ctx.context.run.handoff(...)` (`RunContext.handoff`) in a tool of the main model. Each handoff writes a `harness.handoff` record and streams a `harness.handoff` event (`from`, `to`, `note`, `by`, `runId`). A message that an agent writes has `metadata.harness.agent`. A turn hands off 10 times at most. An app tool named `handoff` fails the turn. - With `routing.humans: true`, a handoff to `{ human: true }` opens a `human` wait with no default expiry, and the turn ends `interrupted`. User messages during the pause go in the transcript, and no model answers them. `session.send(text, { as: 'operator' })` adds an operator reply. It needs the `operator` access action, and it is rejected with `no_human_wait` when no handoff to a person is open. The operator ends the pause with an answer to the `human` wait: `{ resume: true }` or `{ handBack: '' | 'main' }`. - `host.transfer({ from, to, envelope, principal, transferId })` moves a conversation to another thread, also of another harness, in the same tenant. Only the envelope moves: a note, then no message, a summary that the source model writes, or the settled transcript with its media. The principal needs `transfer` on the source and `send` on a target that exists. A new target has the principal as its only owner. Both logs get a `harness.transferred` record, the index entries link the threads in `metadata`, `host.events()` yields a `transferred` event, and a retry with the same `transferId` is safe. A fork stays in its harness: `host.fork` with a thread of another harness throws. -- Agents run from code with `session.agents..run()` and `.start()`. In agent code, `ctx.harness` gives `step.do()` (a finished step returns its stored value when the run runs again), `step.sleep({ name, ms })`, `agents`, `state`, and `handoff` (in a root agent). +- Agents run from code with `session.agents..run()` and `.start()`. `agentRun.receipt` resolves. In agent code, `ctx.harness` gives `step.do()` (a finished step returns its stored value when the run runs again), `step.sleep({ name, ms })`, `agents`, `state`, and `handoff` (in a root agent). - `ctx.harness.agents` in agent code and `ctx.agents` in a plugin have one type, `AgentActions`: `start(agent, input, options?)`, `run(agent, input, options?)`, and `group({ concurrency, onFailure, key }, (g) => [g.start(...)])`. The input type comes from the input schema of the agent, and the result type from its output. `group` resolves to the results in order. With `onFailure: 'collect'`, each entry is an `AgentSettled` (`{ ok, value }` or `{ ok, error }`). The children of a group over its `concurrency` wait as `queued`. - Each start takes `key`, `limits: { timeoutMs, maxAttempts }`, `budget: { maxTokens, maxCost }`, and `onModelError`. A run over its budget fails with the code `budget`, after the first model call over it. A run past `limits.timeoutMs` fails with the code `timeout`. A child never has more budget than its parent has left. `onModelError` is not in the log, so a run that another host takes over runs without it. - The log keeps every run of a session as a `RunNode` of one run tree (`turn`, `agent`, `tool`, or `command`), with its parent, its user, its limits, its budget, and its result. `agentRuns()` reads the tree, and `AgentRunInfo` has `runId`, `kind`, `budget`, `result`, and `error`. Each `harness.usage` record has the `runId` of its run, and an agent run that ends keeps the usage of the run and its children. @@ -29,10 +31,12 @@ Add `@tanstack/ai-harness`: build an agent harness on TanStack AI. A harness kee - `ctx.harness.state({ key, initial })` gives the agents of a session a shared state by key, with `get` and `update`. The log keeps it, and a fork copies it. `STATE_SNAPSHOT` has `{ plugins, shared }`. - A routed agent takes steers at its next model call, and late joins in the same turn, and `turn.beforeFinish` and `turn.onModelError` run for it. - Every host keeps each session tree in one log, `stores.log` from `@tanstack/ai-persistence`. `memoryPersistence()` has a memory log, and a host with no `persistence` uses it. The log holds the inputs, events, transcripts, runs, interrupts, settings, config, plugin state, agent results, schedules, and tool steps, and the host folds them when a thread opens. A root thread and its child threads (agent runs and subagents) write to the log of the root. `session.logState()` reads the `durability.reduce` fold of the whole tree. With a log store that survives a restart, the next host recovers a turn that a crash stopped. -- One host owns a log at a time: the first write takes it, one timer per host renews it, and a host that loses it stops its work on it. A read never takes the log. An input that reaches another host waits in the input lane of the log, and the owner runs it (receipt `queued`). `ownership: { ttlMs, renewMs }` sets the timing. -- The owner writes a snapshot of the folded log to the log store after `log.snapshot.afterBytes` of new records (default 1 000 000) or `log.snapshot.afterMs` (default 300 000), and a new host opens from it. With `lifecycle.retention.logRecords: 'since-snapshot'` (default `'all'`), the owner truncates the records before each snapshot. A reader whose cursor is in the dropped part gets one `harness.gap` event (`EventGap`), then the stream goes on. With `'since-snapshot'`, a thread whose snapshot the host cannot use (a changed `reduce` version or `project`) cannot open, and rejects with `LogTruncatedError`. -- `host.submit` stores an input for a thread and runs nothing. `session.schedule({ at, input, scheduleId, principal })` runs an input later, and `session.cancelSchedule({ scheduleId })` stops it. `createHarnessWorker({ host, harnesses, concurrency, signal, pollMs })` runs the threads that have work and no live host, and returns `{ stop }`. `host.resumePending` does one sweep. -- `durableTool()` runs side effects once. Its `step.do` passes `{ idempotencyKey, epoch }` to `fn`. `step.ask({ name, message, schema })` asks the user: the log keeps the question, any host can take the answer, and a run again after a stop returns the logged answer. `step.sleep({ name, ms })` waits in the process, and the log keeps the wake time, so a worker continues the call after a stop. `durability.truncatedToolResult` sets the text of a tool call cut at the output limit. Without it, a cut answer ends the turn. +- One host owns a log at a time: the first write takes it, one timer per host renews it, and a host that loses it stops its work on it. A read never takes the log. An input that reaches another host waits in the input lane of the log, and the owner runs it (receipt `queued`). `ownership: { ttlMs, renewMs }` sets the timing. `renewMs` must be less than `ttlMs`, else `createHarnessHost` throws a `RangeError`. With only `ttlMs`, `renewMs` is a third of it. +- When the owner refuses an offered input (busy, or an id that the log has with another payload), `session.settled()` on the host that offered it rejects with `InputRejectedError`. The log gets a `harness.offer.rejected` record. An offer that comes while the owner appends is not missed. +- A host that does not own a thread opens it while the owner runs a background job. Settings, config values, and the interrupted turn follow the changes of other hosts. `configure` merges into the latest stored settings. +- The owner writes a snapshot of the folded log to the log store after `log.snapshot.afterBytes` of new records (default 1 000 000) or `log.snapshot.afterMs` (default 300 000), and a new host opens from it. With `lifecycle.retention.logRecords: 'since-snapshot'` (default `'all'`), the owner truncates the records before each snapshot. A reader whose cursor is in the dropped part gets one `harness.gap` event (`EventGap`), then the stream goes on. With `'since-snapshot'`, a thread whose snapshot the host cannot use (a changed `reduce` version or `project`) cannot open, and rejects with `LogTruncatedError`. `session.revert()` works with `'since-snapshot'`: the snapshot keeps the last 500 file steps. A plugin sets which of its records the snapshot keeps with `keepRecords: { [type]: count }` on its definition, and reads them, oldest first, with `ctx.session.keptRecords(type)`. +- `host.submit` stores an input for a thread and runs nothing. `session.schedule({ at, input, scheduleId, principal })` runs an input later, and `session.cancelSchedule({ scheduleId })` stops it. `createHarnessWorker({ host, harnesses, concurrency, signal, pollMs })` runs the threads that have work and no live host, and returns `{ stop }`. `host.resumePending` does one sweep. A swept thread that fails to open goes back with a wake in 30 s. +- `durableTool()` keeps the result of each step. A crash can cut a step, and then it runs again, so `step.do` passes `{ idempotencyKey, epoch }` to `fn`. Give the key to your system, so it refuses the second side effect. `step.ask({ name, message, schema })` asks the user: the log keeps the question, any host can take the answer, and a run again after a stop returns the logged answer. `step.sleep({ name, ms })` waits in the process, and the log keeps the wake time, so a worker continues the call after a stop. `durability.truncatedToolResult` sets the text of a tool call cut at the output limit. Without it, a cut answer ends the turn. - `turn.onModelError`, `retryTransientErrors()`, `turn.beforeFinish`, and `turn.canJoin` control how a turn retries, continues, and ends. - Every approval, question, sign-in, client tool, and human takeover that a run waits for is a `Wait` in the session log (`waitId`, `kind`, `runId`, `schema`, `expiresAt`, `message`). `WaitKind` is `approval`, `question`, `sign-in`, `client-tool`, or `human`. `waitKind`, `isApproval`, `waitOfInterrupt`, and `waitOfQuestion` come from the root and from `./view`. `snapshot().pendingWaits` lists the open waits, and `session.waitFor({ runId, signal })` resolves with the next one. A wait ends `answered`, `cancelled`, or `expired`. A wait with `expiresAt` expires on time. Only `sign-in` has a default limit: 15 minutes. - A wait that only the log holds keeps no host busy: the session goes idle, and an answer on any host continues the run. `resolve` answers the waits of a stopped turn and of a background agent run. `answer` answers the waits of `pendingQuestions` and a `human` wait that no turn holds. `host.submit` checks an answer against the wait that it answers. @@ -41,4 +45,11 @@ Add `@tanstack/ai-harness`: build an agent harness on TanStack AI. A harness kee - The `operator` access action answers a `human` wait of any run. `membersOnly()` gives it to the owner and to members. `authorize` gets the `wait` (`waitId`, `kind`) of an `answer` or `operator` action. - `session.settled({ inputId, signal })` takes one object. `signal` stops the wait for the settlement. - `session.onQuestion`, `session.onAuthRequired`, `session.onMedia`, and `session.onUsage` call a listener for the events that a user answers or sees, with a typed value. -- Sessions also have thread settings, fork and revert, usage totals, media files, the transcript of each child agent (`session.agentTranscript`), and `harnessText()` / `harnessAgent()` to use a harness as a text adapter or an agent. In a session, a `harnessAgent` child is a child run in the log of the root, with the user, keys, limits, and budget of the run that calls it. Its usage and its waits show in the parent session. +- Sessions also have thread settings, fork and revert, usage totals, media files, the transcript of each child agent (`session.agentTranscript`), and `harnessText()` / `harnessAgent()` to use a harness as a text adapter or an agent. In a session, a `harnessAgent` child is a child run in the log of the root, with the user, keys, limits, and budget of the run that calls it. Its usage and its waits show in the parent session. The base64 cache of media files is shared by all the sessions of a host. +- `harnessText()` uses the principal of each call. The inner thread is `/inner/` on every host, and `host.open` refuses that id from app code. When a host with access runs the outer thread too, the inner session is checked through the outer thread, so every user of a shared outer thread can use it. The index entry of the inner thread keeps the link to its outer thread. An abort on a remote harness sends a cancel. The usage of the inner run reaches the outer `RUN_FINISHED`. Its SSE reader takes multi-line `data:` and CRLF. `harnessText()` and `harnessAgent()` forward events on one shared path. +- `createHarnessHost({ observability: { tracer, onOperation } })` traces runs. `tracer` takes an OpenTelemetry `Tracer`. `@opentelemetry/api` is an optional peer dependency: the host loads it only when it has a tracer. The host writes `harness.operation` spans (turns, agent runs, commands), and `harness.recovery`, `harness.wait`, and `harness.transfer` spans. A child run on the same host is a child span: the host passes `trace.setSpan(context.active(), parent)` as the third argument of `startSpan`. A tracer that throws is logged and fails no run. `SendOptions.traceparent` stores a W3C `traceparent` with the input, and the run span on the host that runs it gets a link to it. +- `onOperation` gets a `started` report and a `settled` report (`OperationReport`) for each operation that this host runs, with `status`, `durationMs`, `error`, and `usage`. A listener that throws gets a `harness.plugin.warning` event, and the work goes on. +- More lifecycle events in the session stream: `harness.input.queued`, `harness.ownership.changed` (`reason: 'taken_over'` when the store says that the old owner expired, else `'acquired'`), `harness.stranded`, `harness.transferred`, `harness.wait.opened`, and `harness.wait.settled`. `HarnessEventValues` maps each `harness.*` event name to the type of its value. +- `onEvent`, `onQuestion`, `onUsage`, and `onMedia` also get the events that another host writes to the log of the session. An event fires once. The status of `host.events()` counts that work too, for the sessions that are open on the host. +- The router gets `model`: `model.chat({ adapter, options, label })` in `routing.router` (`HarnessRouterContext.model`) calls a model with the keys of the sender, and a cancel of the turn stops it. +- `session.usage()` has `bySource`: the model calls that are not a turn or an agent run, keyed by `plugin:`, `router`, or `transfer`. Each `UsageCall` and `harness.usage` event of such a call has `source`, and `label` when the call has one. The summary of `host.transfer` with `messages: 'summary'` counts in the source thread under `transfer`, with the label `summary`, and stops after 2 minutes. diff --git a/.changeset/ai-harness-plugins.md b/.changeset/ai-harness-plugins.md index e29e1d3b60..0c2ca1a70d 100644 --- a/.changeset/ai-harness-plugins.md +++ b/.changeset/ai-harness-plugins.md @@ -4,12 +4,15 @@ Add first-party plugins. `@tanstack/ai-harness/plugins` loads no Node module. The plugins that read files are in the Node entry `@tanstack/ai-harness/plugins/node`. -- `permissions()`: rules for tools, paths, and commands, and the `default`, `plan`, `acceptEdits`, and `bypass` modes. It decides every call of a model message before the calls run, and opens one approval wait for each call that asks, so the turn stops and no process waits. An answer is `true` or `false`, `'once'`, `'always'`, or `'reject'`, or `{ answer, message? }`. It also has saved "always" rules with `/permissions` and `/permissions forget `, and `listSavedPermissions` / `deleteSavedPermission` for server code. The rules apply in subagent runs too. -- `agents()` (in `plugins/node`): agent profiles with their own prompt, tools, model, rules, and step limit, in code or in Markdown files, `/agent` to switch the primary agent, and the built-in `build`, `plan`, `general`, and `explore` agents. +- `permissions()`: rules for tools, paths, and commands, and the `default`, `plan`, `acceptEdits`, and `bypass` modes. It decides every call of a model message before the calls run, and opens one approval wait for each call that asks, so the turn stops and no process waits. An answer is `true` or `false`, `'once'`, `'always'`, or `'reject'`, or `{ answer, message? }`. It also has saved "always" rules with `/permissions` and `/permissions forget `, and `listSavedPermissions` / `deleteSavedPermission` for server code. The rules apply in subagent runs too. It decides in `onGuardToolCall`, on the final checked arguments. +- `agents()` (in `plugins/node`): agent profiles with their own prompt, tools, model, rules, and step limit, in code or in Markdown files, `/agent` to switch the primary agent, and the built-in `build`, `plan`, `general`, and `explore` agents. A subagent run uses the tools and the model of the turn that started it. A run whose turn is not in the memory of this host fails with a clear error. A run that code starts needs a profile `model`. - `question()`: a `question` tool for the lead turn. Each question is a wait of the turn, answered with `resolve`. Agent runs do not get the tool. - `todos()`, `title()`, `modelPicker()`, `compact()`, and `usage()` (with your own prices for calls that have no provider cost). +- `goal()` runs each turn as the user who set the goal. - `projectInstructions()`, `fileCommands()`, and `boundToolOutput()`, in `plugins/node`. A plugin runs agents with `ctx.agents.start`, `ctx.agents.run`, and `ctx.agents.group({ concurrency }, (g) => [g.start(agent, input)])`, with the same types as `ctx.harness.agents` in agent code. `PermissionRules`, `PermissionResources`, `PermissionDecisionCapability`, and `decidePermission` are also exported from the package root. They load no Node module, so an edge-safe plugin can read the permission rules. + +`ctx.model.chat({ adapter, options, label, run, timeoutMs })` calls a model from a plugin, for example a judge or a title. It uses the keys of the user of the run, the budget of the run applies, and it stops when the run stops or after `timeoutMs` (default 2 minutes). Its usage counts in the session under `plugin:` in `session.usage().bySource`. `goal()`, `title()`, and `compact()` call their models this way, with the labels `judge`, `title`, and `compact`. diff --git a/.changeset/ai-harness-transport.md b/.changeset/ai-harness-transport.md index d7109ce2be..ed658500f0 100644 --- a/.changeset/ai-harness-transport.md +++ b/.changeset/ai-harness-transport.md @@ -4,9 +4,19 @@ Serve a harness session to clients. -- `createHarnessHandler()` from `@tanstack/ai-harness/server` serves a session over HTTP, as the user that the host `access` reads from each request. `POST run` is standard AG-UI, so `useChat` with `fetchServerSentEvents` works on a harness, with approvals, client tools, and a reload that joins the running turn. `POST control`, `GET events`, `GET transcript`, `GET agent-transcript` (the messages of a child agent run), `GET snapshot`, and `GET describe` cover one session. `sessions` (with the `share` and `unshare` ops), `GET host-events`, and the media routes cover the rest. A refused action answers 403. A read of a thread that does not exist answers 404 and starts no session. The read routes read a thread with `host.observe`, so they set up no plugins. `GET describe` opens the session. -- `handleHarnessSocket()`, from the same entry, serves the same session over a WebSocket. It takes the upgrade `request` and reads the user from it. +- `createHarnessHandler()` from `@tanstack/ai-harness/server` serves a session over HTTP, as the user that the host `access` reads from each request. `POST run` is standard AG-UI, so `useChat` with `fetchServerSentEvents` works on a harness, with approvals, client tools, and a reload that joins the running turn. `POST control`, `GET events`, `GET transcript`, `GET agent-transcript` (the messages of a child agent run), `GET snapshot`, and `GET describe` cover one session. `sessions` (with the `share` and `unshare` ops), `GET host-events`, and the media routes cover the rest. A `host.events()` reader listens from the call, and one that falls more than 1 000 changes behind stops, and reads again for the current statuses. A refused action answers 403. A read of a thread that does not exist answers 404 and starts no session. The read routes read a thread with `host.observe`, so they set up no plugins. `GET describe` opens the session. +- `handleHarnessSocket()`, from the same entry, serves the same session over a WebSocket. It takes the upgrade `request` and reads the user from it. It handles the frames of one connection in order. +- A `command` or an `agent` input over HTTP or the socket runs once for one `inputId`, and the `inputId` of its receipt is the id of the client. - `expose` on `defineHarness()` names the agents, settings, config keys, and commands that a client can use. Any other input of those kinds gets `rejected` with `not_exposed`. -- `createHarnessClient()` from `@tanstack/ai-harness/client` is a typed, browser-safe client. It sends inputs, reads the session, and streams events that resume from a cursor. The `prompt` input of the protocol and `client.send` take `as: 'operator'` for the reply of an operator after a handoff to a person. +- `createHarnessClient()` from `@tanstack/ai-harness/client` is a typed, browser-safe client. It sends inputs, reads the session, and streams events that resume from a cursor. Without `from`, `client.events()` starts at the newest event (the cursor of the snapshot), so it does not send the whole history first. This is the default: a reader that needs the history (to rebuild a transcript, for example) must pass `from: '0'` or a saved cursor. When the snapshot answer has no string `cursor`, `client.events()` throws and does not retry. `client.events()` connects again after `idleTimeoutMs` (default 60 000) with no bytes. The `prompt` input of the protocol and `client.send` take `as: 'operator'` for the reply of an operator after a handoff to a person. +- The client methods take one object: `cancelInput`, `revert`, `agentTranscript`, `mediaUrl`, `append`, and `putMedia`. `client.sessions.list`, `rename`, `delete`, `update`, and `fork` manage the sessions. - `oauthConnector()` and the OAuth helpers from `@tanstack/ai-harness/oauth` (`loopbackLogin`, `deviceLogin`, `exchangeCode`, `refreshCredential`, and more) keep a credential per user. A tool calls `token({ run: ctx.context.run })` for the token of the user of its run. `token({ run, wait: true })` stops the turn with an `auth_required` interrupt until the user signs in. `deviceLogin`, `loopbackLogin`, `startLoopbackReceiver`, `exchangeCode`, and `refreshCredential(config, credential, { fetch, signal })` stop at once when their `signal` aborts, so a cancel or a session close stops a running `/connect`. `OAuthConfig.issuer` makes `loopbackLogin` refuse a code from another issuer (RFC 9207). The loopback deadline (`timeoutMs`) does not keep the process alive. - `run.credentials.refresh(id, update, { signal })` refreshes a saved credential once for all callers on a host, keyed by tenant, user, and id. It saves with `setIf` when the store has `getVersioned` and `setIf`, and uses the other host's token on a conflict. Without that pair, it does a plain `set`. It saves to the scope where it found the credential, so a tenant token stays the tenant's. `update` gets a signal that only `host.close()` aborts. An abort of `signal` rejects only that call. `oauthConnector` refreshes this way, and `refreshCredential` keeps the `client` of the old credential. +- The session stream speaks protocol version 2 on every transport: `GET events`, the socket, and the stdio worker. Each event is an AG-UI wire chunk in a `harness.event` frame, with `cursor`, `threadId`, and `operationId`. The AG-UI `RUN_STARTED` of a child agent run has `parentRunId`, so a client can rebuild the run tree. +- The first frame of a stream is `harness.hello` with the version `v`. A client asks for a version with `?v=` or with `v` in `harness.subscribe`. Another version gets a `harness.error` frame with the code `unsupported_version`, and the stream closes (socket code 4400). +- A `harness.error` frame has a `code` (`unauthorized`, `forbidden`, `invalid_frame`, `unsupported_version`, `slow_consumer`, or `internal`) and a `requestId`. On the socket and the stdio worker, the `requestId` is the id of the frame that failed. Every answer of the handler has an `X-Request-Id` header: the id that the client sent when it is a short safe id, else a new UUID. +- Streams stay open and protect the server. SSE streams send a `: ping` comment and the socket sends a `harness.ping` frame when idle. SSE reads the log only as fast as the client reads, and the socket waits while its `bufferedAmount` is over the limit. A client that does not read in time gets a `slow_consumer` error frame, and the stream closes (socket code 4408). A slow client loses no events: it connects again from its last cursor. `createHarnessClient` does this in `events()` and `hostEvents()`, and ends the read with a `HarnessProtocolError` (with its `code`) for any other error frame. +- `stream: { heartbeatMs, slowConsumerMs, sseQueueBytes, socketBufferBytes }` on `createHarnessHandler` and `handleHarnessSocket` changes these values. The defaults are 20 s, 30 s, 1 MiB, and 8 MiB. A value that is not a finite number above 0 throws at the start. `HarnessStreamOptions` is the type. +- A reader whose events are gone gets one `harness.gap` event, with `reason: 'truncated'` (a truncate of the log) or `'evicted'` (the memory feed of a session dropped them). The stream then goes on from `resumeAt`. +- A W3C `traceparent` goes with a message: as the `traceparent` header of `POST control` and `POST run`, or as the `traceparent` field of a `harness.input` frame. `ClientSendOptions.traceparent` makes `client.send` send the header. +- A `cancel` input stops a turn from any host that reads the same log, so `client.cancel({ operationId })` works behind a load balancer. A closed connection still never cancels work. To stop a `useChat` turn, call `client.cancel` with the `runId` of `useChat` in the stop handler. diff --git a/.changeset/ai-harness-view.md b/.changeset/ai-harness-view.md index f8d9234694..ffeeb53dfb 100644 --- a/.changeset/ai-harness-view.md +++ b/.changeset/ai-harness-view.md @@ -6,8 +6,9 @@ Add `createSessionView()` in `@tanstack/ai-harness/view`. It turns a session (lo - The state: messages with their parts, tool calls, approvals, client tools, other interrupts, questions, sign-ins, `waits` (every open wait, with its kind), child agents with their results and a `status` of `running` or `waiting`, media with signed URLs, waiting inputs, usage, models, settings, commands, and plugin state. - The waits of background agent runs show in `approvals`, `clientTools`, and `interrupts`. The view answers each one at once: with `resolve`, or with `answer` for a `human` wait that no turn holds. -- Actions: `send` (a prompt, a steer, or a `/command`), `followUp`, `continue`, `upload`, `cancel`, `cancelInput`, `setDelivery`, `configure`, `reset`, `revert`, and the answers to approvals (with edited arguments), client tools, interrupts, and questions. +- Actions: `send` (a prompt, a steer, or a `/command`), `followUp`, `continue`, `putMedia`, `cancel`, `cancelInput`, `setDelivery`, `configure`, `reset`, `revert`, `loadAgent`, and the answers to approvals (with edited arguments), client tools, interrupts, and questions. `cancelInput`, `revert`, `approve({ id, ...options })`, `reject({ id, payload? })`, `notice({ text })`, and `putMedia` take one object, and so does the `reject({ payload? })` of an approval item. - Typed events: `view.on('approval' | 'interrupt' | 'turnEnd' | ..., handler)` and plugin events. - In a harness with `routing`, an assistant message has `agent`, the root agent that wrote it. After a handoff, the next agent writes a new message in the same turn. - `view.send(text, { as: 'operator' })` sends the reply of an operator while a handoff to a person is open. The reply shows in each open view at once. -- A view that opens while a turn runs shows the turn from its start. After a reload, each child agent shows its own messages again. +- A view that opens while a turn runs shows the turn from its start. After a reload, each child agent shows its card, and `view.loadAgent({ id })` fills it with its own messages. +- When the stream sends a `harness.gap` event, the view stops its read, loads the snapshot and the transcript again, then reads on from the new cursor. No turn shows twice. diff --git a/.changeset/ai-mcp-harness.md b/.changeset/ai-mcp-harness.md index 9085ac6e05..4f95c025c4 100644 --- a/.changeset/ai-mcp-harness.md +++ b/.changeset/ai-mcp-harness.md @@ -5,9 +5,10 @@ Use MCP in a TanStack AI harness. - `@tanstack/ai-mcp/harness` adds `createHarnessMcpServer`. It serves a harness as an MCP server, so Claude Code, Claude Desktop, Cursor, or another agent can use it. The tools are `chat`, `steer`, `cancel`, `approve`, `reject`, `resolve`, `answer`, and `status`, plus `agent_` for each exposed agent and `command_` for each exposed command. Every tool takes an optional `threadId`. With `approvals: 'ask'` (the default), the client asks the user about each approval when it supports elicitation. Otherwise the approvals come back in the result. `approvals: 'auto'` approves every tool call. `resolve` answers every interrupt of a stopped turn in one call, or the waits of a background agent run. Every result lists the open waits of its work in `waits`, each with a `kind` from `waitKind` (`approval`, `question`, `sign-in`, `client-tool`, or `human`). The `status` tool lists every wait, and each background agent with `status: 'running'` or `'waiting'`. -- On a host with `access`, `createHarnessMcpServer` reads the user of each HTTP request with `identify`. The host checks each tool call of that user. A host without `access` needs `openAccess: true` (no user) or `openAccess: { principal }` (one local user), else the server throws at the start. Each tool call releases its session handle when it ends, so the session closes on idle. The `status` tool and the media reads read the thread with `host.observe`, so they set up no plugins. +- On a host with `access`, `createHarnessMcpServer` reads the user of each HTTP request with `identify`. The host checks each tool call of that user. A host without `access` needs `openAccess: true` (no user) or `openAccess: { principal }` (one local user), else the server throws at the start. Each tool call releases its session handle when it ends, so the session closes on idle. The `status` tool and the media reads read the thread with `host.observe`, so they set up no plugins. The server forgets a wait when its operation settles. - `@tanstack/ai-mcp/harness` also adds the `mcp()` harness plugin. It connects stdio and HTTP servers by name, gives the model their tools as `_`, and adds a `/mcp` status command. A stdio server starts once per host, and every session of the host shares it. `oauth: true` signs in to a server with `/connect `, for each user. +- `mcp()` connects again with backoff when a server closes, and lists the tools again on `list_changed`. A server that fails its first connect also tries again with backoff. `mcpConnector` and `oauth: true` servers connect again at the next turn, and list their tools again. Each server takes `needsApproval` and `toolName`. The backoff starts again only after a connection stayed open for 30 s, so a server that crashes right after it connects waits longer each time. When two servers give a tool with the same name, `mcp()` warns once. With `onDuplicateTool: 'error'`, the later server fails instead, with an error that names both servers and the tool. - `@tanstack/ai-mcp/connector` adds `mcpConnector`: one OAuth MCP server as a harness plugin, with `needsApproval`, `prefix`, `scopes`, and `clientName`. The sign-in passes the `iss` of the callback on (RFC 9207). A cancel or a session close stops a running sign-in at once: the loopback stops listening, and the OAuth requests abort. When the MCP SDK refreshes a token, the connector saves it with `credentials.refresh` of the harness: to the scope of the sign-in, and with compare-and-set when the store has it. Each session still sends its own refresh request. When the server refuses a token, the connector deletes the saved sign-in only if it is the token that this session used. -- `mcp()` and `mcpConnector` answer the `elicitation/create` requests of MCP spec 2025 servers, and the input requests of spec 2026-07-28 servers, as session questions. One tool call can ask up to 5 times. Each question is a live wait of the run that made the tool call: it ends `cancelled` when that run stops. +- `mcp()` and `mcpConnector` answer the `elicitation/create` requests of MCP spec 2025 servers, and the input requests of spec 2026-07-28 servers, as session questions. One tool call can ask up to 5 times. Each question is a live wait of the run that made the tool call: it ends `cancelled` when that run stops. `mcp()` tracks the live questions for each server, so each question goes to the run that called the tool. A spec 2025 request goes to the user only when all the running calls of that server belong to one run. Else the server gets `decline`. `@tanstack/ai-harness` is an optional peer dependency. Only the `./harness` and `./connector` entries need it. diff --git a/.changeset/ai-sandbox-harness.md b/.changeset/ai-sandbox-harness.md index 89b9f4e848..ea295c6dcf 100644 --- a/.changeset/ai-sandbox-harness.md +++ b/.changeset/ai-sandbox-harness.md @@ -6,6 +6,6 @@ Add `@tanstack/ai-sandbox/harness` for `@tanstack/ai-harness`: - `codingAgents({ sandbox, agents, workspace? })` is a harness plugin. The lead model gets one tool per coding agent (Claude Code, Codex, Grok Build, or any ACP agent). Each agent works in the sandbox, keeps its own session per thread (also after a restart), and starts read-only in the `plan` mode of `permissions()`. `/fresh [agent]` starts a new session. `workspace: 'shared'` (default) runs the agents of a thread one at a time in one sandbox. `'per-agent'` gives each agent its own sandbox. - `codingAgents` runs agents only on the host that owns the log of the thread, so `'shared'` runs one agent at a time also with many hosts. -- `sandboxWorkspaceBackend(sandbox, { owned? })` runs the coding tools of `@tanstack/ai-harness/plugins/coding` (`workspaceTools()`, `snapshots()`, `formatter()`) in a sandbox. It resolves links in the sandbox, reads one page of a file in the sandbox, and stops background jobs by their id, also from another host. `binding()` gives the provider and the sandbox id, so another host can resume the sandbox. `close()` stops the jobs of the backend, and with `owned: true` destroys the sandbox when the session closes. +- `sandboxWorkspaceBackend(sandbox, { owned? })` runs the coding tools of `@tanstack/ai-harness/plugins/coding` (`workspaceTools()`, `snapshots()`, `formatter()`) in a sandbox. It resolves links in the sandbox, reads one page of a file in the sandbox, and stops background jobs by their id, also from another host. `binding()` gives the provider and the sandbox id, so another host can resume the sandbox. `close()` stops the jobs of the backend. A session close closes it with `'moved'`, so an owned sandbox stays for the next host. `host.sessions.delete` closes it with `'closed'`, and with `owned: true` that destroys the sandbox. `@tanstack/ai-harness` is an optional peer dependency. diff --git a/.changeset/ai-tool-guard.md b/.changeset/ai-tool-guard.md new file mode 100644 index 0000000000..30087006ea --- /dev/null +++ b/.changeset/ai-tool-guard.md @@ -0,0 +1,11 @@ +--- +'@tanstack/ai': minor +--- + +Add the `onGuardToolCall` middleware hook for checks that must always run, such as permissions. + +- It runs for every tool call, after all `onBeforeToolCall` hooks and after the `inputSchema` check. `hookCtx.args` holds the final checked arguments that the tool gets. +- The guard of every middleware runs, in order, until one returns a decision. An earlier `onBeforeToolCall` decision does not skip it. +- Return `{ type: 'skip', result }` to deny the call. The tool does not run, and the model gets `result`, the same as an `onBeforeToolCall` skip. Return `{ type: 'abort', reason }` to abort the run. +- It runs before a server tool runs and before a client tool goes to the client. A guard cannot change the arguments. +- The new `GuardToolCallDecision` type is exported. diff --git a/.changeset/ai-websocket-buffered-amount.md b/.changeset/ai-websocket-buffered-amount.md new file mode 100644 index 0000000000..5391cc79e7 --- /dev/null +++ b/.changeset/ai-websocket-buffered-amount.md @@ -0,0 +1,5 @@ +--- +'@tanstack/ai': patch +--- + +`WebSocketLike` has an optional `bufferedAmount`. A server can read it to stop sending to a slow socket. A socket without it works as before. diff --git a/.changeset/code-mode-harness.md b/.changeset/code-mode-harness.md index 6b2b43211d..c9c649240c 100644 --- a/.changeset/code-mode-harness.md +++ b/.changeset/code-mode-harness.md @@ -2,4 +2,4 @@ '@tanstack/ai-code-mode': minor --- -Add `@tanstack/ai-code-mode/harness` with `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. By default, only tools that are safe to run without a question move into code mode: server tools with no `needsApproval` that the permission rules allow in plan mode, and never a tool that declares permission resources. `include` picks the tools yourself, a safe tool with `metadata.codeMode: true` always moves, and `lazy: true` lists only the tool names and adds `discover_tools`. Tool names that are not JavaScript identifiers get a safe name inside the program. `@tanstack/ai-harness` is an optional peer dependency. +Add `@tanstack/ai-code-mode/harness` with `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. By default, only tools that are safe to run without a question move into code mode: server tools with no `needsApproval` that the permission rules allow in plan mode, and never a tool that declares permission resources. `include` picks the tools yourself, but it never moves a tool that needs approval or that the rules deny or ask about, a safe tool with `metadata.codeMode: true` always moves, and `lazy: true` lists only the tool names and adds `discover_tools`. Tool names that are not JavaScript identifiers get a safe name inside the program. The rules decide in the permission mode of the session, again at each turn, also for `include`. `@tanstack/ai-harness` is an optional peer dependency. diff --git a/.changeset/harness-acp-agent.md b/.changeset/harness-acp-agent.md index bc91c478a5..2364e1e7cc 100644 --- a/.changeset/harness-acp-agent.md +++ b/.changeset/harness-acp-agent.md @@ -2,4 +2,6 @@ '@tanstack/ai-acp': minor --- -`@tanstack/ai-acp/agent` adds `createAcpAgent` and `serveAcp`, which serve a harness as an ACP v2 agent (experimental, like the draft protocol). Each wait of a turn goes to the editor by its kind: an approval becomes a permission request (with `allow_always` only when the answer schema takes `always`), a question or a `human` wait becomes a form elicitation, and a sign-in becomes a URL elicitation, when the editor declares them in `initialize`. A wait that the editor cannot answer is cancelled. Cancel cancels the running turn, and image, audio, and embedded file blocks go to the media store of the session. `session/list` lists the sessions of the harness from the session index of the host. The `principal` option is the user of the editor: each session opens as that user, and the list shows the sessions of that user. It is required on a host with `access`. The agent holds one session handle per ACP session. `session/close` releases it, and the close of the connection releases them all, so the host closes each session on idle. `@tanstack/ai-acp` now uses `@agentclientprotocol/sdk` 1.8, pinned to `~1.8.0` because the agent uses its experimental v2 entry point. The existing ACP client adapters keep working. +`@tanstack/ai-acp/agent` adds `createAcpAgent` and `serveAcp`, which serve a harness as an ACP v2 agent (experimental, like the draft protocol). Each wait of a turn goes to the editor by its kind: an approval becomes a permission request (with `allow_always` only when the answer schema takes `always`), a question or a `human` wait becomes a form elicitation, and a sign-in becomes a URL elicitation, when the editor declares them in `initialize`. A wait that the editor cannot answer is cancelled. Cancel cancels the running turn, and image, audio, and embedded file blocks go to the media store of the session. A prompt takes its place in the queue before its media is saved, and gives the place back when the session is idle. `session/list` lists the sessions of the harness from the session index of the host. The `principal` option is the user of the editor: each session opens as that user, and the list shows the sessions of that user. It is required on a host with `access`. The agent holds one session handle per ACP session. `session/close` releases it, and the close of the connection releases them all, so the host closes each session on idle. `@tanstack/ai-acp` now uses `@agentclientprotocol/sdk` 1.8, pinned to `~1.8.0` because the agent uses its experimental v2 entry point. The existing ACP client adapters keep working. + +One agent app serves one connection, and reads the editor capabilities from that connection. A second connection gets a clear error. For more connections, make one app for each: `new AcpServer({ createAgent: () => createAcpAgent(options) })`. A `request_permission` names the real tool call id. A session that the agent made or opened in its connection takes prompts after the agent released its handle, also on a host with no `stores.sessions`. diff --git a/.changeset/harness-cli-handler-options.md b/.changeset/harness-cli-handler-options.md index a53abd427c..bd1a0fc1f0 100644 --- a/.changeset/harness-cli-handler-options.md +++ b/.changeset/harness-cli-handler-options.md @@ -2,4 +2,4 @@ '@tanstack/ai-harness-cli': patch --- -`--serve` passes `media: { secret }` to the harness handler. On a host without `access`, it also passes `openAccess`: the `principal` of `runCli`, or `true` without one. So each request with the token runs as the user of the CLI. On a host with `access`, `identify` reads the user of each request. +`--serve` signs media URLs with its own key, never with the token: `--media-key` or `HARNESS_MEDIA_KEY` (at least 32 bytes), or a new random key at each start. Without a set key, signed media URLs stop working after a restart. On a host without `access`, it also passes `openAccess`: the `principal` of `runCli`, or `true` without one. So each request with the token runs as the user of the CLI. On a host with `access`, `identify` reads the user of each request. diff --git a/.changeset/harness-join-any-instance.md b/.changeset/harness-join-any-instance.md index 07a836dcd6..dc47ad6053 100644 --- a/.changeset/harness-join-any-instance.md +++ b/.changeset/harness-join-any-instance.md @@ -5,5 +5,5 @@ Any server instance that reads the same log store can serve a run. - `POST run` sends the log cursor of each event as its SSE `id`. A client that is cut off sends the same request with `Last-Event-ID`, and any instance continues the stream from there. The turn does not run again. -- `GET run?runId=&threadId=` finds a turn that another instance runs. It streams the turn from the log of `threadId`, from its start or from `Last-Event-ID`, until the turn ends. Without `threadId`, only the instance that runs the turn finds it, and another instance answers `400`. The join request of a `ChatClient` sends only the run id, so put `threadId` in the URL of the connection. +- `GET run?runId=&threadId=` finds a turn that another instance runs. It streams the turn from the log of `threadId`, from its start or from `Last-Event-ID`, until the turn ends. It reads the log from the start of the turn, not from the start of the log. A turn rebuilt after a restart (a retried `POST run`) also streams from its start. Without `threadId`, only the instance that runs the turn finds it, and another instance answers `400`. The join request of a `ChatClient` sends only the run id, so put `threadId` in the URL of the connection. - `GET run?threadId=` shows a turn that another instance runs as `activeRun`, so a reloaded `useChat` joins it. diff --git a/.changeset/harness-secure-defaults.md b/.changeset/harness-secure-defaults.md index cbe404cdcb..73d5aedf09 100644 --- a/.changeset/harness-secure-defaults.md +++ b/.changeset/harness-secure-defaults.md @@ -5,5 +5,6 @@ Make the transport defaults safe for a server with many instances. - `createHarnessHandler` and `handleHarnessSocket` read the user of each request with the host `access`. On a host without `access`, pass `openAccess: true` (no user) or `openAccess: { principal }` (one local user), else they throw at the start. So a server never serves every thread to everyone by mistake. -- `media: { secret }` signs the media URLs of `createHarnessHandler`. Set the same secret on every instance, so a media URL that one instance signs works on every instance. Without it, the handler signs with a random key of its own, and a URL works only on that instance. A signed URL names its user, and the handler checks that user again when it serves the file. +- `media: { secret }` signs the media URLs of `createHarnessHandler`. Set the same secret on every instance, so a media URL that one instance signs works on every instance. Without it, the handler signs with a random key of its own, and a URL works only on that instance. A signed URL names its harness and its user, and the handler checks that user again when it serves the file. A URL of a deleted thread answers 404. +- `media: { keys: [newKey, oldKey] }` rotates the key: the first key signs, and every key checks. Each key must be at least 32 bytes, else the handler throws at the start. - `oauthConnector` needs `login`. Use `'device'` on a server, and `'loopback'` only for a harness on the user's own machine, such as a CLI. diff --git a/.changeset/harness-skills.md b/.changeset/harness-skills.md index de7f2585e5..50ac03741d 100644 --- a/.changeset/harness-skills.md +++ b/.changeset/harness-skills.md @@ -2,4 +2,4 @@ '@tanstack/ai-skills': minor --- -`@tanstack/ai-skills/harness`: the new `skills({ dirs })` plugin gives the harness model the skills of a list of folders, and makes each skill a `/` command that starts a turn with that skill. A skill with the name of another command is `/skill:`, and `/skills` lists them. The plugin watches the folders, so a skill that you add or remove while the session runs changes the commands at once. +`@tanstack/ai-skills/harness`: the new `skills({ dirs })` plugin gives the harness model the skills of a list of folders, and makes each skill a `/` command that starts a turn with that skill. A skill with the name of another command is `/skill:`, and `/skills` lists them. The plugin watches the folders, so a skill that you add or remove while the session runs changes the commands at once. The watcher lives as long as the host, every session of the host shares it, and a change makes each session read the folders again. A skill that does not load gets a fallback message. diff --git a/.changeset/harness-tool-ops.md b/.changeset/harness-tool-ops.md index 5dbfbc10bd..4e0078cb69 100644 --- a/.changeset/harness-tool-ops.md +++ b/.changeset/harness-tool-ops.md @@ -20,7 +20,7 @@ Run a harness tool with no model turn, add host records from a client, and add a other tool is refused with `not_exposed`. The capabilities document marks each tool with `exposed`. - `expose.records` lists the host record types a client may add with an - `{ op: 'append', records }` input, or with the new `client.append(records)`. + `{ op: 'append', records }` input, or with the new `client.append({ records })`. A record of another type, or with a `thread` field, is refused with `not_exposed`. The session log keeps the records. - `createHarnessMcpServer` adds a `tool_` tool for each tool in diff --git a/.changeset/harness-turn-and-body-limits.md b/.changeset/harness-turn-and-body-limits.md new file mode 100644 index 0000000000..c3a8dc216c --- /dev/null +++ b/.changeset/harness-turn-and-body-limits.md @@ -0,0 +1,14 @@ +--- +'@tanstack/ai-harness': patch +--- + +Fix how a turn takes steers, settings, and plugin removals, and limit the size of client requests. + +- A steer that arrives during a `session.tool()` turn waits, and then runs as its own turn. Before, the tool turn took it, did not answer it, and settled it `completed`. +- A turn keeps the config values that it started with. A config change while it runs, for example `/mode bypass`, applies from the next turn. +- `configure({ settings: { plugins: { remove } } })` removes everything that those plugins give the thread: also their agents, subagents, commands, agent and generation middleware, and the items that they add for other plugins (for example the permission rules of `agents()`). A remove that leaves a plugin (of any lifetime) or a harness middleware without a capability that it needs is rejected. +- When two plugins give a tool with the same name from `discoverTools`, the first one stays, and clients get one `harness.plugin.warning` event that names both plugins. +- `createHarnessHandler` reads each JSON body up to `media.maxBodyBytes` (default 32 MiB). A bigger body answers 413. `handleHarnessSocket` closes the socket with code 1009 on a bigger frame. The WebSocket server reads the whole frame first, so also set its own frame limit (for example `maxPayload` of `ws`). +- Client media is checked part by part before any part is stored, so a refused input stores nothing. A retried input with the same `inputId` stores its media no second time, and gets the first receipt. +- A client input with a provider file handle (`source.type: 'file'`) is refused with `media_url_not_allowed` unless the harness sets `media.allowProviderFiles: true`. Server code can still send file handles. +- `POST run` moves the inline bytes of the last user message to the media store, as `POST control` does, so the transcript keeps a `harness-media:` reference. A refused part answers 413 (`media_too_large`), 415 (`media_not_allowed`), or 400 (`media_invalid`, `media_url_not_allowed`). diff --git a/.changeset/persistence-session-stores.md b/.changeset/persistence-session-stores.md index 7e368dfae6..9b1fa164b9 100644 --- a/.changeset/persistence-session-stores.md +++ b/.changeset/persistence-session-stores.md @@ -4,7 +4,7 @@ Add opt-in session stores. The chat and generation middleware work as before without them. -- `log` (`LogStore`): an append-only log for each session tree, with 14 methods, each with one object argument. Records: `append`, `read`, and `subscribe`. Ownership: `acquire`, `renew`, `release`, `setAffinity`, and `listOrphaned`. The input lane: `offer` and `offered`, consumed in the same write with `append({ consumes })`. Snapshots and lifecycle: `getSnapshot`, `putSnapshot`, `truncate`, and `deleteLog`. Each fenced write carries the epoch from `acquire`, and a write with an old epoch rejects with `OwnershipLostError`. A writer that never calls `acquire` uses `UNOWNED_EPOCH`. `append` rejects with `LogConflictError` when another writer took the position, and `read` rejects with `LogTruncatedError` before the truncation point. `memoryLogStore()` is the reference store. `memoryLogBackend()` lets two store instances share one log. The conformance kit needs `twin`, a second store set on the same backend, for the `log` cases across instances. +- `log` (`LogStore`): an append-only log for each session tree, with 14 methods, each with one object argument. Records: `append`, `read`, and `subscribe`. Ownership: `acquire`, `renew`, `release`, `setAffinity`, and `listOrphaned`. The input lane: `offer` and `offered`, consumed in the same write with `append({ consumes })`. Snapshots and lifecycle: `getSnapshot`, `putSnapshot`, `truncate`, and `deleteLog`. `acquire` returns `expired: true` when it takes the log from another owner whose ownership expired with no release. Each fenced write carries the epoch from `acquire`, and a write with an old epoch rejects with `OwnershipLostError`. A writer that never calls `acquire` uses `UNOWNED_EPOCH`. `append` rejects with `LogConflictError` when another writer took the position, and `read` rejects with `LogTruncatedError` before the truncation point. `memoryLogStore()` is the reference store. `memoryLogBackend()` lets two store instances share one log. The conformance kit needs `twin`, a second store set on the same backend, for the `log` cases across instances. - `sessions` (`SessionIndexStore`): one index entry per tenant and thread. The key is `{ tenantId?, threadId }` (`SessionIndexKey`), so one thread id in two tenants is two entries. An entry has an `ownerId` and `members` (`SessionMember`, with the role `owner`, `member`, or `viewer`). The store has `get(key)`, `getVersioned(key)`, `upsert(entry)`, `upsertIf({ entry, revision })` for compare-and-set writes, `list`, and `delete(key)`, a hard delete. `list` reads one tenant, filters by member, parent thread, title `search`, `harness` name, and exact `metadata` values, and pages with a cursor. `defineLogStore` and `defineSessionIndexStore` type an adapter. `memoryPersistence()` now also has `log` and `sessions`. `runPersistenceConformance` runs the cases of each session store when the store is present, so an existing backend needs no `skip` entry. diff --git a/docs/advanced/middleware.md b/docs/advanced/middleware.md index c52d5e90ea..6a843a8a9f 100644 --- a/docs/advanced/middleware.md +++ b/docs/advanced/middleware.md @@ -734,6 +734,38 @@ The `hookCtx` provides: | `toolName` | `string` | Tool name | | `toolCallId` | `string` | Tool call ID | +### onGuardToolCall + +A permission check in `onBeforeToolCall` can miss calls. An earlier middleware can return a decision first, and then your hook does not run. Put checks that must always run in `onGuardToolCall`: + +```typescript +import { type ChatMiddleware } from "@tanstack/ai"; + +const noDeletes: ChatMiddleware = { + name: "no-deletes", + onGuardToolCall: (ctx, hookCtx) => { + if (hookCtx.toolName === "deleteFile") { + return { type: "skip", result: { error: "Deleting files is not allowed." } }; + } + }, +}; +``` + +When the guard runs: + +- It runs for every tool call, after all `onBeforeToolCall` hooks and after the `inputSchema` check. +- `hookCtx.args` holds the final checked arguments. The tool gets these same arguments. +- The guard of every middleware runs, in array order, until one returns a decision. +- It runs just before a server tool runs, and before a client tool goes to the client. A call that waits for approval does not run it. + +| Decision | Effect | +|----------|--------| +| `void` / `undefined` | Allow the call. The next guard runs. | +| `{ type: 'skip', result }` | Deny the call. The tool does not run, and the model gets `result`, the same as an `onBeforeToolCall` skip. | +| `{ type: 'abort', reason? }` | Abort the entire chat run | + +A guard cannot change the arguments. To change them, use `onBeforeToolCall`. + ### onAfterToolCall Called after each tool execution (or skip). Every middleware runs, in array order. @@ -1025,6 +1057,7 @@ const stream = chat({ | `onStart` | Sequential | All run in order | | `onChunk` | **Piped** — chunks flow through each middleware | If first drops a chunk, later middleware never see it | | `onBeforeToolCall` | **First-win** — first non-void decision wins | Earlier middleware has priority | +| `onGuardToolCall` | Sequential. Every guard runs until one returns a decision | An earlier `onBeforeToolCall` decision never skips a guard | | `onShouldContinue` | **AND** — any explicit `false` stops the loop | Order only affects which middleware runs first when short-circuiting | | `onModelError` | Sequential. If any middleware calls `ctx.retry()` and the call has retries left, the model is called again | All run in order, so each one sees the error | | `onAfterToolCall` | **Piped**: each gets the result of the previous `replaceResult` | The last replacement wins | diff --git a/docs/config.json b/docs/config.json index 8871af0e04..fc516be022 100644 --- a/docs/config.json +++ b/docs/config.json @@ -353,7 +353,7 @@ "label": "Middleware", "to": "advanced/middleware", "addedAt": "2026-04-15", - "updatedAt": "2026-10-14" + "updatedAt": "2026-10-17" }, { "label": "Built-in Middleware", @@ -761,7 +761,7 @@ "label": "Store Reference", "to": "persistence/store-reference", "addedAt": "2026-08-04", - "updatedAt": "2026-10-17" + "updatedAt": "2026-10-19" }, { "label": "How Persistence Works", @@ -927,37 +927,43 @@ "label": "Build your first harness", "to": "harness/overview", "addedAt": "2026-10-13", - "updatedAt": "2026-10-17" + "updatedAt": "2026-10-18" }, { "label": "Send and show media", "to": "harness/media", "addedAt": "2026-10-13", - "updatedAt": "2026-10-16" + "updatedAt": "2026-10-19" }, { "label": "Connect clients", "to": "harness/connect", "addedAt": "2026-10-13", - "updatedAt": "2026-10-17" + "updatedAt": "2026-10-19" + }, + { + "label": "Write a client in any language", + "to": "harness/protocol", + "addedAt": "2026-10-17", + "updatedAt": "2026-10-19" }, { "label": "Access and tenants", "to": "harness/access", "addedAt": "2026-10-17", - "updatedAt": "2026-10-18" + "updatedAt": "2026-10-19" }, { "label": "Run in the terminal", "to": "harness/cli", "addedAt": "2026-10-13", - "updatedAt": "2026-10-18" + "updatedAt": "2026-10-19" }, { "label": "Deploy a harness", "to": "harness/deploy", "addedAt": "2026-10-13", - "updatedAt": "2026-10-18" + "updatedAt": "2026-10-19" }, { "label": "Send inputs safely", @@ -990,7 +996,8 @@ { "label": "Lifecycle and retention", "to": "harness/lifecycle", - "addedAt": "2026-10-17" + "addedAt": "2026-10-17", + "updatedAt": "2026-10-19" }, { "label": "Run side effects once", @@ -1007,7 +1014,7 @@ "label": "Build on the session log", "to": "harness/session-log", "addedAt": "2026-10-13", - "updatedAt": "2026-10-18" + "updatedAt": "2026-10-19" }, { "label": "Control how a turn ends", @@ -1019,7 +1026,7 @@ "label": "Store settings per thread", "to": "harness/thread-settings", "addedAt": "2026-10-13", - "updatedAt": "2026-10-17" + "updatedAt": "2026-10-18" }, { "label": "Compact a long session", @@ -1043,7 +1050,7 @@ "label": "Write a plugin", "to": "harness/plugins", "addedAt": "2026-10-13", - "updatedAt": "2026-10-18" + "updatedAt": "2026-10-19" }, { "label": "Track usage and cost", @@ -1072,13 +1079,14 @@ { "label": "Run the coding tools for many threads", "to": "harness/coding-workspaces", - "addedAt": "2026-10-18" + "addedAt": "2026-10-18", + "updatedAt": "2026-10-19" }, { "label": "Undo a turn", "to": "harness/snapshots", "addedAt": "2026-10-13", - "updatedAt": "2026-10-18" + "updatedAt": "2026-10-19" }, { "label": "Build a coding agent", @@ -1120,17 +1128,19 @@ "label": "Build your own UI", "to": "harness/custom-ui", "addedAt": "2026-10-13", - "updatedAt": "2026-10-18" + "updatedAt": "2026-10-19" }, { "label": "Code mode in a harness", "to": "harness/code-mode", - "addedAt": "2026-10-13" + "addedAt": "2026-10-13", + "updatedAt": "2026-10-18" }, { "label": "Work until a goal is met", "to": "harness/goal", - "addedAt": "2026-10-13" + "addedAt": "2026-10-13", + "updatedAt": "2026-10-18" }, { "label": "Run agents from a harness", @@ -1146,13 +1156,14 @@ { "label": "Add skills", "to": "harness/skills", - "addedAt": "2026-10-13" + "addedAt": "2026-10-13", + "updatedAt": "2026-10-19" }, { "label": "Use MCP servers", "to": "harness/mcp", "addedAt": "2026-10-13", - "updatedAt": "2026-10-18" + "updatedAt": "2026-10-19" }, { "label": "Use from any MCP client", diff --git a/docs/harness/access.md b/docs/harness/access.md index 4d433dd1fd..c58ec2be33 100644 --- a/docs/harness/access.md +++ b/docs/harness/access.md @@ -106,7 +106,7 @@ Each action covers these calls: - For `answer` and `operator`, `authorize` also gets `wait`: the `waitId` and the `kind` of the [wait](./waits) that the input answers, when the host knows it. A `human` wait takes either action: when `answer` is refused, the host asks for `operator`. - For `cancelSchedule`, `run` is the schedule: the `runId` is the schedule id, and the `principalId` is the user who made it. - `host.submit` and `schedule` check an answer or a cancel when you send it, with `run` from the log. A run that the log does not have is refused with `unknown_run`: `submit` returns a `rejected` receipt, and `schedule` throws `InputRejectedError`. -- A handle checks `read` once, when it opens. Each action that changes the session asks `authorize` again. +- A handle checks `read` once, when it opens. Each action that changes the session asks `authorize` again. This includes `cancel` and `send` on an operation or an agent run that the handle returns. ## What the host checks for you @@ -181,19 +181,22 @@ export const host = createHarnessHost({ }, }) -export function removeFromTeam(threadId: string, user: { id: string; tenantId: string }) { +export async function removeFromTeam(threadId: string, user: { id: string; tenantId: string }) { teams.get(threadId)?.delete(user.id) - host.revoke({ threadId, principal: user }) + await host.revoke({ threadId, principal: user }) } ``` -After `host.revoke`: +When the promise of `host.revoke` resolves: -- The streams of the user on the thread close. A socket gets a `forbidden` error frame. -- The open handles of the user on the thread throw `AccessDeniedError`. A new `host.open` asks `authorize` again. +- The running and queued work of the user on the thread is cancelled, on every host. The schedules of the user on the thread are removed. +- The streams of the user on the thread close. A socket gets a `forbidden` error frame. A running `events()` read of a handle or a view ends with `AccessDeniedError`. +- The open handles and views of the user on the thread throw `AccessDeniedError`. A new `host.open` asks `authorize` again. - The signed media URLs of the user stop working. A signed URL names its user, and the handler checks that user again when it serves the file. - `host.events()` yields a `revoked` event. +When the cancels cannot be sent (for example, the log store fails), the promise rejects. The access is gone already. Call `host.revoke` again to stop the work. `host.sessions.unshare` rejects the same way, after it removed the member. + `host.sessions.share` and `unshare` work only with `membersOnly()`. With your own tables, change the tables, then call `host.revoke`. ## Serve with no login diff --git a/docs/harness/acp.md b/docs/harness/acp.md index ee59f5bdf4..8b9574b571 100644 --- a/docs/harness/acp.md +++ b/docs/harness/acp.md @@ -50,7 +50,8 @@ serveAcp({ host, harness: assistant, version: '1.0.0' }) - Each ACP session is one harness session. The ACP session id is the thread id. - `principal` is the user of the editor, for example `serveAcp({ host, harness: assistant, principal: { id: 'ada' } })`. Every session opens as that user, so `/connect` saves keys for them. On a host with [`access`](./access), `principal` is required. -- The prompts of a session run one after another. The harness `busy` setting does not apply. +- The prompts of a session run one after another, in the order that the editor sent them. A prompt with a large file keeps its place while the file uploads. The harness `busy` setting does not apply. +- The agent holds a session only while a prompt runs or waits. An idle session closes after the host `lifecycle.idleCloseMs`, and the next prompt opens it again. - `version` is the version that `initialize` reports. Default: `'0.0.0'`. - `cwd` is the working directory that the session list shows for each session. Default: `process.cwd()`. - Use a persistent store instead of `memoryPersistence()` if a session must stay after the process stops. @@ -73,7 +74,7 @@ The agent sends each wait to the editor by its kind: | Kind | What the editor shows | | --- | --- | -| `approval` | A permission request. **Allow always** shows only for an approval of `permissions()`, because its answer schema takes `always`. | +| `approval` | A permission request with the id of the real tool call. **Allow always** shows only for an approval of `permissions()`, because its answer schema takes `always`. | | `question`, `human` | A form elicitation of the answer schema. | | `sign-in` | A URL elicitation with the sign-in page. | | `client-tool` | Nothing. The editor cannot run it, so the agent cancels it. | @@ -108,6 +109,8 @@ import { createAcpAgent } from '@tanstack/ai-acp/agent' export const app = createAcpAgent({ host, harness: assistant }) ``` +One agent app serves one connection. It reads the editor capabilities from that connection, and a second connection to it fails with an error. To serve more editors, for example over HTTP, make one app for each connection. With the `AcpServer` of `@agentclientprotocol/sdk/experimental/server`, pass `createAgent: () => createAcpAgent({ host, harness: assistant })`. + ## What you have now - Your harness as an agent in any ACP v2 editor. diff --git a/docs/harness/agents.md b/docs/harness/agents.md index 79fea8ae72..2d4e140054 100644 --- a/docs/harness/agents.md +++ b/docs/harness/agents.md @@ -143,10 +143,10 @@ export const withReview = agents({ | Field | What it does | | --- | --- | -| `model` | A model id. `adapter` turns it into an adapter. Without it, a primary agent keeps the harness model, and a subagent uses the model of the main turn. | +| `model` | A model id. `adapter` turns it into an adapter. Without it, a primary agent keeps the harness model, and a subagent uses the model of the turn that started it. | | `system` | Text for the system prompt of each turn of this agent. | | `tools` | The tools that this agent can call, by name. A `*` matches any text, for example `todo_*`. Without it, the agent gets every tool. | -| `permissions` | [Permission rules](./permissions#write-rules) that count while this agent is the primary agent. | +| `permissions` | [Permission rules](./permissions#write-rules) that count while this agent is the primary agent. At that time, they also apply in each subagent run. | | `steps` | The step limit. See [Limit the steps](#limit-the-steps). | | `hidden` | `true` leaves the agent out of the `agent` command and setting. The model can still start it as a subagent. | @@ -208,19 +208,20 @@ A profile with the mode `subagent` or `all` is also a subagent. The main model g A subagent run gets these parts of its profile: -- The tools of the last main turn, filtered by its `tools`. +- The tools of the turn that started it, filtered by its `tools`. - Its `system` prompt. -- Its `model`, or the model of the main turn. +- Its `model`, or the model of the turn that started it. - Its step limit. A subagent run stops after 50 model calls, or when an answer is cut at the output limit. The harness also limits the depth and the number of children. See [Stay within limits](./subagents#stay-within-limits). To give one run from code its own budget and time limit, see [Give one run a budget and a time limit](./subagents#give-one-run-a-budget-and-a-time-limit). ## Known limits -- The `permissions` of a profile count only while it is the primary agent. When the profile runs as a subagent, its own rules do not apply. Put rules that must hold in every run in `permissions({ rules })`. See [Rules in subagents](./permissions#rules-in-subagents). +- A subagent run uses the `permissions` of the primary agent, not the rules of its own profile. For example, when `build` is the primary agent and starts `review`, the `build` rules apply and the `review` rules do not. Put rules that must hold in every run in `permissions({ rules })`. See [Rules in subagents](./permissions#rules-in-subagents). - A Markdown agent cannot set `permissions`. Put those rules in a profile in code. - On Amazon Bedrock, the model can still call a tool on the last call. The tool does not run, but the turn can end with no text answer. -- A subagent that starts before the first turn of the session has no tools. It also needs its own `model`. +- A subagent run gets the tools and the model of its turn from the memory of the host that ran that turn. On another host, or after a restart, the run fails with `Agent cannot go on here`. Start it again from a turn. +- A subagent run that your code starts, not a turn, has no tools of a turn. Give its profile a `model`. ## What you have now diff --git a/docs/harness/cli.md b/docs/harness/cli.md index 88509cf4a7..58461aaa0c 100644 --- a/docs/harness/cli.md +++ b/docs/harness/cli.md @@ -59,7 +59,13 @@ To use the harness from another program: `--mcp` and `--serve` run only the commands in `expose.commands` of the harness. `--serve` also sets only the config keys in `expose.config`. Add the names to `expose` in `defineHarness`, for example `expose: { commands: ['connect:github'] }`. Line mode and `-p` run every command. See [Choose what clients can change](./connect#choose-what-clients-can-change). -Line mode reads one message or command per line and waits for each turn. It works the same in a terminal and with piped input. In a terminal, it also opens sign-in links in the browser. +Line mode reads one message or command per line. In a terminal, you can type while a turn runs: + +- `/cancel` cancels the running turn at once. +- A message joins the running turn as a steer. The model reads it at its next call. +- A message that you type when no turn runs starts a new turn. + +With piped input, each line waits for the turn that the line before it started. A script runs in order. In a terminal, line mode also opens sign-in links in the browser. ## 3. Use it in CI @@ -111,7 +117,7 @@ console.log(await exit) // 143 `--serve` is the quick way to put a harness on HTTP, so it makes the trust decisions for you: one bearer token, one user, every thread. Anyone with the token acts as that one user. They can open or drive any thread on the host, including the tools it runs (shell commands, file edits) on the machine the CLI runs on. Treat the token like a shell login. - It binds to `127.0.0.1:8787`, reachable only from the same machine. Change the bind with `--host` and `--port`. `--host 0.0.0.0` puts it on the network, so do that only behind TLS and a proxy you trust. -- The CLI signs media links with the token, so rotating the token stops any link a turn already handed out. +- The CLI signs media links with its own key, never with the token. By default it makes a new random key at each start, so media links stop working after a restart. To keep them working, pass `--media-key` or set `HARNESS_MEDIA_KEY`, with a key of at least 32 bytes. Change the key to stop every link a turn already handed out. Need per-user access, your own sign-in check, or a rule for which threads each user can open? Use [`createHarnessHandler`](./connect) on a host with [`access`](./access) instead. The host reads the user of each request and checks each action, so no thread is open by accident. @@ -144,7 +150,7 @@ await host.close() - Give `host`, or `persistence` and `credentials`, not both. With both, `runCli` writes `Give runCli a host, or persistence and credentials, not both.` and returns code 1. - Line mode, `-p`, your own `ui`, `--acp`, and `--mcp` open their sessions as `principal`, so `/connect` saves keys for that user. - `--serve` checks the token first. On a host without `access`, each request then runs as `principal`. On a host with [`access`](./access), `identify` reads the user of each request. -- On a host with `access`, give `principal` too. `--mcp` does not start on such a host, because stdio has no login. `--dashboard` opens its sessions without the principal. +- On a host with `access`, give `principal` too. `--mcp` does not start on such a host, because stdio has no login. ## Commands in line mode @@ -215,8 +221,9 @@ process.exitCode = await runCli(assistant, { }) ``` -- `ui` runs only in an interactive terminal. Piped input uses line mode. `-p`, `--acp`, `--mcp`, `--serve`, and `--dashboard` do not use `ui`. +- `ui` runs only in an interactive terminal. Piped input uses line mode. `-p`, `--acp`, `--mcp`, and `--serve` do not use `ui`. - When `ui` resolves, the CLI disposes the view and `runCli` returns. +- In a reopened thread, a saved child agent shows only its card. To show what it wrote, call `view.loadAgent({ id })`, for example when the screen draws the card. See [Show subagents](./custom-ui#show-subagents). - To write `Screen`, read [Build your own UI](./custom-ui). For an Ink screen with approvals, questions, sign-ins, and child agents, copy [`examples/harness-cli/src/tui.tsx`](https://github.com/TanStack/ai/blob/main/examples/harness-cli/src/tui.tsx). diff --git a/docs/harness/code-mode.md b/docs/harness/code-mode.md index 266c022390..b6a66b42ed 100644 --- a/docs/harness/code-mode.md +++ b/docs/harness/code-mode.md @@ -99,7 +99,8 @@ export const onlyIssues = codeMode({ }) ``` -- A tool that `include` picks moves even when it needs approval, and its calls inside the program do not ask. Pick only tools that are safe to run without a question. +- `include` never moves a tool with `needsApproval`, or a tool that an `ask` or `deny` rule matches. That tool stays a normal tool call with its approval, and code mode warns once. The rules decide in the permission mode of the session, again at each turn. So in `plan` mode, a tool that edits stays a normal tool call. +- Any other tool that `include` picks moves, also when a rule asks for it. Its calls inside the program do not ask, so pick only tools that are safe to run without a question. - With `include`, a call of `execute_typescript` gets the permission rules like any tool, for example a question with `default: 'ask'`. - A safe tool with `metadata.codeMode: true` moves even when `include` does not pick it. diff --git a/docs/harness/coding-tools-backends.md b/docs/harness/coding-tools-backends.md index c68126ed95..ca6b849f48 100644 --- a/docs/harness/coding-tools-backends.md +++ b/docs/harness/coding-tools-backends.md @@ -225,6 +225,7 @@ export const withToml = formatter({ - `builtins: false` turns the built-in formatters off. - `timeoutMs` limits each run. The default is 20 seconds. - In a sandbox or on your own backend, give `formatter()` the same `backend` as `workspaceTools()`. +- With [`workspaceTools({ workspace: { bind } })`](./coding-workspaces) before it, `formatter()` runs in the workspace of each thread. Leave out `root` and `backend`. A formatter that fails, stops at the time limit, or is not installed does not fail the tool call. The file stays as the tool wrote it, and the plugin sends `FormatFailed`. A formatter that changes the file sends `FormatterRewrote` with `{ path, formatter }`: @@ -251,7 +252,7 @@ The rules decide before the call runs, and the formatter runs after the write. T - When the formatter asks, the call asks one time for the write and the formatter, also in `acceptEdits` mode. Answer that [approval](./permissions#answer-an-approval) with `resolve`. - A `reject` stops the call. Nothing is written, and the formatter does not run. -- Answer `always` once, and the formatter runs without a question in this project. +- Answer `always` once, and the formatter runs without a question in this project. The answer is tied to a hash of the formatter config files, like `prettier.config.js` and `package.json`. When they change, the next write asks again. - A rule that denies the formatter keeps the write. The tool result says that the formatter did not run. - A rule can allow a formatter: `{ tool: 'formatter:prettier', decision: 'allow' }`. Without `permissions()`, the formatter runs without a question. diff --git a/docs/harness/coding-tools.md b/docs/harness/coding-tools.md index 0aa89709ae..a00faa5463 100644 --- a/docs/harness/coding-tools.md +++ b/docs/harness/coding-tools.md @@ -173,10 +173,10 @@ For `bash`: - The model gets the exit code and the last 1000 lines (20 KiB) of the output. - A command stops after 2 minutes, with exit code 124. Every process that it started stops too. Set another limit with `bashTimeoutMs`. The model can ask for up to 10 minutes. -- With `spillDir`, long output goes to a file in that folder, and the model gets the path. `spillDir` is relative to `root`, or absolute. +- With `spillDir`, long output goes to a file in that folder, and the model gets the path. `spillDir` is relative to `root`, or absolute. The file has `[redacted]` in place of each secret of the session, the same as the tool result. See [Keep tokens out of the log and the model](./auth#keep-tokens-out-of-the-log-and-the-model). - A command gets only a few variables of your process, like `PATH` and `HOME`. Your provider keys stay out of it. See [Keep server secrets out of commands](./coding-workspaces#2-keep-server-secrets-out-of-commands). -With `background: true`, the call returns at once with a job id. When the job ends, the model gets a note with the output, and a new turn starts. +With `background: true`, the call returns at once with a job id. When the job ends, the model gets a note with the output, and a new turn starts. That turn runs as the user whose run started the job. - The time limit does not apply to a background job. - At most 8 background jobs run at a time in a session. Set another limit with `maxJobs`. Each job keeps the newest 1 MiB of its output. diff --git a/docs/harness/coding-workspaces.md b/docs/harness/coding-workspaces.md index 6b4b71dcf3..7fb9eba82f 100644 --- a/docs/harness/coding-workspaces.md +++ b/docs/harness/coding-workspaces.md @@ -51,8 +51,9 @@ export const coder = defineHarness({ - The `root` that `bind` gives is the boundary of every tool. A path outside it is refused. With `outside: 'ask'`, a rule that allows the path, or `bypass` mode, lets a call use it. No approval can cover it, because the root is not known before the call. The `cwd` thread setting must stay inside it. - The tools resolve every link with `realpath` before they read or write, while they hold the lock on the file. So a link in one thread cannot lead into the folder of another thread. - The session log keeps a `harness.workspace.bound` record with the place of the workspace. [Move a thread to another server](#5-move-a-thread-to-another-server) uses it. +- Put `workspaceTools()` before [`snapshots()`](./snapshots) and [`formatter()`](./coding-tools-backends#format-each-file-after-a-write). Then they work in the workspace of each thread, and their `root` is optional. -For a sandbox per thread, create it in `bind`, and give the backend `owned: true`. Then the sandbox is destroyed when the session closes: +For a sandbox per thread, create it in `bind`, and give the backend `owned: true`. Then the thread owns the sandbox: ```ts group=harness-coding-workspaces import { sandboxWorkspaceBackend } from '@tanstack/ai-sandbox/harness' @@ -70,6 +71,9 @@ export const sandboxed = workspaceTools({ }) ``` +- A session close (an idle close, a deploy, or a move to another server) closes the backend with `'moved'`. The sandbox stays, so the next server can [reattach](#a-sandbox) it. +- [`host.sessions.delete`](./lifecycle#delete-a-thread-for-good) ends the thread. The host reattaches the sandbox and closes it with `'closed'`, so it is destroyed. The host that deletes must know the harness: it opened the harness, or it has the harness in its [`harnesses`](./lifecycle#delete-a-thread-for-good) option. Else it warns, and the sandbox stays. + ## 2. Keep server secrets out of commands A command gets only a few variables of the server process: @@ -118,6 +122,8 @@ With `background: true`, `bash` returns at once with a job id like `bash-1`. Two - At most 8 background jobs run at a time in a session. One more gives the model a tool error. Set another limit with `maxJobs`. - Each job keeps the newest 1 MiB of its output. +When a job ends, the model gets a note, and the wake turn runs as the user whose run started the job. + Each job sends events. Listen for them in a plugin: ```ts group=harness-coding-workspaces @@ -193,7 +199,7 @@ export const movable = workspaceTools({ ``` - Give the workspace back to continue in the same files. -- Give `'gone'` when the sandbox is gone. Then each tool call of the thread fails with `The workspace of this thread, /workspace in the docker sandbox , is gone, so the tools cannot work on it.`, and the tools send `tanstack/workspace-tools:stranded`. +- Give `'gone'` when the sandbox is gone. Then each tool call of the thread fails with `The workspace of this thread, /workspace in the docker sandbox , is gone, so the tools cannot work on it.`, and the tools send `tanstack/workspace-tools:stranded`. The session also sends a `harness.stranded` event, for clients that do not know the plugin. See the [event catalog](./protocol#event-catalog). - Without `reattach`, a sandbox counts as `'gone'`. ### A folder on one server diff --git a/docs/harness/connect.md b/docs/harness/connect.md index a13e2d84c9..59dd9bb415 100644 --- a/docs/harness/connect.md +++ b/docs/harness/connect.md @@ -50,6 +50,27 @@ export const handler = createHarnessHandler({ host, harness: assistant }) When more than one instance serves the same URL, also set `media: { secret }`. The secret signs media URLs, so every instance must use the same value. Without it, the handler signs with a random key of its own, so a signed URL works only on this instance, until it restarts. +```ts group=harness-connect +const mediaKey = process.env.MEDIA_KEY +const oldMediaKey = process.env.OLD_MEDIA_KEY +if (!mediaKey || !oldMediaKey) { + throw new Error('Set MEDIA_KEY and OLD_MEDIA_KEY.') +} + +export const rotatingHandler = createHarnessHandler({ + host, + harness: assistant, + // The first key signs. Every key checks the URLs. + media: { keys: [mediaKey, oldMediaKey] }, +}) +``` + +- Each key must be at least 32 bytes. A shorter key throws at the start. +- To change the key, put the new key first in `keys` and keep the old key after it. A URL that the old key signed works until it expires (1 hour). Then remove the old key. +- A signed URL works only for this harness and for the user that it was signed for. When the thread is deleted, its URLs answer `404`. + +A request body over `media.maxBodyBytes` (default 32 MiB) answers `413`. See [Files inside a client message](./media#files-inside-a-client-message). + The handler answers these paths under your route. For chat and events: - `GET capabilities`: the AG-UI capabilities, with the agents in `expose.agents`. @@ -84,13 +105,24 @@ A cursor is a position in the log of the thread. It works on every instance that To run more than one instance behind one URL, give every instance the same log store. -A host with `logRecords: 'since-snapshot'` drops old records (see [Drop the records before a snapshot](./lifecycle#drop-the-records-before-a-snapshot)). A reader whose cursor is in the dropped part gets one `CUSTOM` event first: +Sometimes the events after a cursor are gone. A host with `logRecords: 'since-snapshot'` drops old records (see [Drop the records before a snapshot](./lifecycle#drop-the-records-before-a-snapshot)), and before its log opens, a session keeps only its newest events in memory. A reader whose cursor is in the lost part gets one `CUSTOM` event first: - Its name is `harness.gap` (`HARNESS_EVENTS.gap`). -- Its value is an `EventGap`: `{ from, resumeAt, reason: 'truncated' }`. Its cursor is `resumeAt`. +- Its value is an `EventGap`: `{ from, resumeAt, reason }`. `reason` is `truncated` or `evicted`. Its cursor is `resumeAt`. - Then the stream goes on from `resumeAt`. -The session view loads the snapshot and the transcript again when it sees this event. A custom client must do the same. +The session view loads the snapshot and the transcript again when it sees this event. A custom client must do the same. See [When events are missing](./protocol#when-events-are-missing). + +## Keep the stream open + +Every transport sends the same frames: each event is an AG-UI wire chunk in a `harness.event` frame. A long stream stays usable through proxies and slow networks: + +- The first frame of `GET events` and of the socket is `harness.hello`, with the protocol version `v`. A client can ask for a version with `?v=2` or `v: 2`. +- When the stream is idle, the server sends a heartbeat every 20 s. Set the idle timeout of your proxy to more than that. +- A client that does not read for 30 s gets a `slow_consumer` error frame, and the stream closes. Connect again from the last cursor. +- Every answer of the handler has an `X-Request-Id` header, and every error frame has a `code` and a `requestId`. Log the id, to find the request in the server logs. + +To change the heartbeat or the limits, pass `stream` to `createHarnessHandler` and `handleHarnessSocket`. [The harness wire protocol](./protocol) lists every frame, the error codes, and the `stream` values, for a client in any language. ## Choose what clients can change @@ -113,7 +145,7 @@ export const team = defineHarness({ | `config` | Set these plugin config keys with `client.setConfig({ key, value })`. | | `commands` | Run these commands with `client.command({ name })`, or with `/name` in a [session view](./custom-ui) on a client. | | `tools` | Run these harness tools with no model turn, with `client.tool({ name, args })`. See [Run a tool without a model](./tool-runs). | -| `records` | Add records of these types to the thread with `client.append(records)`. See [Let a client add records](./session-log#let-a-client-add-records). | +| `records` | Add records of these types to the thread with `client.append({ records })`. See [Let a client add records](./session-log#let-a-client-add-records). | - Any other input of these kinds gets `{ status: 'rejected', reason: 'not_exposed' }`. - `GET describe` and `client.describe()` list only the exposed commands and config keys. A UI built from them shows only what a client can use. @@ -134,16 +166,22 @@ const client = createHarnessClient({ headers: { authorization: 'Bearer my-token' }, }) +const { cursor } = await client.snapshot() await client.send('Summarize my open tickets.') -for await (const entry of client.events()) { +for await (const entry of client.events({ from: cursor })) { 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. +`events()` reconnects after a network error and continues from the last cursor. It also reconnects when no bytes come for 60 s, so a dead connection (a laptop that slept) does not stay `open`. To change that time, pass `idleTimeoutMs` to `createHarnessClient`. Keep it larger than the [heartbeat](./protocol#stream-values) of the server, or the client reconnects again and again. A second tab, a phone, or a reload all see the same session. + +Without `from`, `events()` starts at the newest event, so it sends no old events. That is why the example reads the snapshot cursor before it sends. + +- To read every event from the start of the session, pass `from: '0'`. Do this when you rebuild a transcript from the events. +- If the snapshot answer has no `cursor` (the URL is not a harness handler), `events()` throws and does not try again. To upload files and show the media that your agents make, read [Send and show media](./media). @@ -338,6 +376,80 @@ The join has these rules: - Without `threadId`, only the instance that runs the turn finds it. Another instance answers `400`. - A turn that ended gives the events that the session still keeps, and then the stream ends. +### Stop a turn from useChat + +The `stop()` of `useChat` closes the request, but the turn continues on the server. A closed connection never cancels work, because another tab or a reload can still read the turn. To stop the turn too, send a `cancel` input for its run. + +1. On the server, serve the harness with one handler. `POST run` streams the turn to `useChat`, and `POST control` takes the cancel: + +```ts group=harness-connect-stop-server +import { createHarnessHost, defineHarness, membersOnly } from '@tanstack/ai-harness' +import { createHarnessHandler } from '@tanstack/ai-harness/server' +import { memoryPersistence } from '@tanstack/ai-persistence' +import { openaiText } from '@tanstack/ai-openai' + +export const writer = defineHarness({ + name: 'acme/writer', + adapter: openaiText('gpt-6.1-sol'), +}) + +const host = createHarnessHost({ + persistence: memoryPersistence(), + access: { + identify: ({ request }) => + request.headers.get('authorization') === `Bearer ${process.env.HARNESS_TOKEN}` + ? { id: 'user-1' } + : null, + authorize: membersOnly(), + }, +}) + +// Mount it on `/api/harness/*`. +export const handler = createHarnessHandler({ host, harness: writer }) +``` + +2. In the browser, call `client.cancel` with the `runId` of `useChat` in your stop handler. The turn of `POST run` runs as that id: + +```tsx group=harness-connect-stop-client +import { fetchServerSentEvents, useChat } from '@tanstack/ai-react' +import { createHarnessClient } from '@tanstack/ai-harness/client' +import type { writer } from './harness' + +const headers = { authorization: 'Bearer my-token' } + +const client = createHarnessClient({ + url: '/api/harness', + threadId: 'user-1-thread', + headers, +}) + +export function StoppableChat() { + const { messages, runId, isLoading, stop, sendMessage } = useChat({ + threadId: 'user-1-thread', + connection: fetchServerSentEvents('/api/harness/run', { headers }), + }) + + const stopTurn = async () => { + stop() + if (!runId) return + const receipt = await client.cancel({ operationId: runId }) + if (receipt.status === 'rejected') console.warn(receipt.reason) + } + + return ( + <> +

{messages.length} messages

+ + {isLoading ? : null} + + ) +} +``` + +- The cancel works on any server that reads the same log store. The input goes to the server that runs the turn. +- If the turn already ended, the receipt can be `rejected` with the reason `not_running`. +- The turn ends with the status `cancelled`. + ## Use a WebSocket For one connection that carries events and inputs, hand the socket and its upgrade request to `handleHarnessSocket`: @@ -351,12 +463,15 @@ export function onUpgrade(socket: WebSocketLike, request: Request) { } ``` -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. +The client sends `{ type: 'harness.subscribe', threadId }` first, and the host answers with a `harness.hello` frame. After that, each `harness.input` frame gets a `harness.receipt`, and every event arrives as a `harness.event` frame with its cursor. - The host reads the user from the upgrade request with `identify`, the same as for HTTP. Every input on the socket runs as that user. - A user that `identify` refuses gets an `unauthorized` error frame, and the socket closes. +- The socket handles the frames of one connection in the order they came. So a prompt with a large image reaches the session before the frame after it. - A user who cannot read the thread, or who loses access later, gets a `forbidden` error frame, and the socket closes. +- A frame over `media.maxBodyBytes` (default 32 MiB) gets an error frame, and the socket closes with code `1009`. Your WebSocket server reads the whole frame before this check, so the check does not cap memory. Set the frame limit of your server too, for example `maxPayload` of `ws`. - On a host with no `access`, pass `openAccess`, the same as for `createHarnessHandler`. +- Pass `socket: true` to `createHarnessHandler` too. Then `GET capabilities` says `transport.websocket: true`, so clients know that they can use the socket. The default is `false`. ## Use it as the model of another chat @@ -375,13 +490,18 @@ const stream = chat({ 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. +- Who runs the inner turn: when the outer `chat()` runs inside a harness turn, the inner session opens as the user of that turn. Else it opens as `principal`, for example `harnessText(assistant, { host, principal: { id: 'ada' } })`. +- Access: the inner thread is `/inner/`. When `host` has `access` and also runs the outer thread, the inner session is checked through the outer thread, so every user who can read a shared outer thread can use its inner session. `host.open` refuses a thread id with `/inner/` from app code. +- Usage: the outer run gets the token usage of every inner model call, so the usage totals of the outer session count them. +- Cancel: an abort of the outer call cancels the inner turn. For a remote harness, `harnessText({ url, token })` also tells the remote host to cancel the turn. + If a plugin picks the model of the harness, the harness has no adapter to read its input kinds from. Pass them with `harnessText(assistant, { host, inputModalities: ['text', 'image'] })`. ## What you have now - One handler that serves standard AG-UI and the session stream, as the user of each request. - Clients that can change only what you expose. -- A `useChat` screen on the harness, with approvals, tools that run in the browser, and a reload that joins the running turn. +- A `useChat` screen on the harness, with approvals, tools that run in the browser, a reload that joins the running turn, and a stop button that cancels the turn. - A typed web client that reconnects and resumes from a cursor. - Messages from the browser to a running agent. - A WebSocket, and a harness you can call from `chat()`. diff --git a/docs/harness/custom-ui.md b/docs/harness/custom-ui.md index fbf2e4302d..67c1875188 100644 --- a/docs/harness/custom-ui.md +++ b/docs/harness/custom-ui.md @@ -162,7 +162,7 @@ A notice has a `kind`: - `error`: a turn or an action failed. - `rejected`: the session did not accept a message. - `command`: the text result of a command. -- `ui`: a line that you added with `view.notice(text)`. +- `ui`: a line that you added with `view.notice({ text })`. ## 3. Act on the session @@ -172,18 +172,18 @@ Send work: - `view.followUp(text)`: the message waits, and runs as its own turn after the current work. - `view.continue()`: starts a turn from the saved messages, with no new message. - `view.command({ name, input })`: runs a command, for example from a button. -- `view.upload({ body, mimeType, name })`: stores a file. Send the record you get back with `view.send(text, [record])`. See [Show media](#show-media). +- `view.putMedia({ body, mimeType, name })`: stores a file. Send the record you get back with `view.send(text, [record])`. See [Show media](#show-media). Change the session: - `view.setConfig({ key, value })`: changes a config value, for example `view.setConfig({ key: 'mode', value: 'plan' })`. - `view.configure({ settings })`: changes the stored settings of the thread, for example `view.configure({ settings: { model: 'strong' } })`. - `view.reset({ note })`: the model starts fresh from the next turn, with your `note` first. See [Reset with a handoff note](./fork-and-reset#reset-with-a-handoff-note). -- `view.revert(messageId)` and `view.unrevert()`: go back to a message, and undo that. +- `view.revert({ messageId })` and `view.unrevert()`: go back to a message, and undo that. Stop work: - `view.cancel()`: cancels the running turn. `view.cancel({ operationId })` cancels one operation, for example a background agent. -- `view.cancelInput(inputId)`: drops a message from `waitingInputs` before it runs. `view.setDelivery({ inputId, delivery: 'steer' })` moves it into the running turn instead. +- `view.cancelInput({ inputId })`: drops a message from `waitingInputs` before it runs. `view.setDelivery({ inputId, delivery: 'steer' })` moves it into the running turn instead. - `view.dispose()`: stops the view when your screen closes. After that, each action throws an error. On a view with a [`HarnessClient`](#a-ui-in-the-browser), two actions need `expose` in the harness: @@ -197,7 +197,7 @@ See [Choose what clients can change](./connect#choose-what-clients-can-change). Answer what waits: -- `approval.approve()` and `approval.reject()`: answer one approval. `view.approve(id)` and `view.reject(id)` do the same by id. +- `approval.approve()` and `approval.reject()`: answer one approval. `view.approve({ id })` and `view.reject({ id })` do the same by id. - `view.approveAll()` and `view.rejectAll()`: answer all open approvals. - `call.resolve(output)` and `call.fail(message)`: answer one client tool call. - `request.resolve(value)` and `request.cancel()`: answer one item of `interrupts`. @@ -210,7 +210,7 @@ const [waiting] = view.store.get().approvals waiting?.approve({ editedArgs: { path: 'notes/draft.md' } }) ``` -`reject(payload)` sends data with a no, when the tool's approval schema takes some. +`reject({ payload })` sends data with a no, when the tool's approval schema takes some. If a turn waits for more than one approval, client tool, or interrupt, the turn continues after you answer all of them. @@ -239,9 +239,9 @@ A source must have these members: These members are optional: - `continue`, `cancelInput`, `setDelivery`, `configure`, `reset`, `revert`, and `unrevert`. Without one, its view action adds an `error` notice. -- `upload` or `putMedia`. If the source has neither, `view.upload` rejects. +- `putMedia`. Without it, `view.putMedia` rejects. - `mediaUrl` and `loadMedia`. Without `mediaUrl`, media parts have no `url`. Without `loadMedia`, `load()` rejects. -- `agentTranscript`. Without it, a reload shows no child parts. +- `agentTranscript`. Without it, `view.loadAgent` does nothing, so a saved child shows only its card. ## 4. Listen for events @@ -250,7 +250,7 @@ Some items need attention when they arrive. `view.on` calls your handler for eac ```ts group=harness-custom-ui view.on('approval', (approval) => { process.stdout.write('\u0007') - view.notice(`${approval.tool} waits for your answer.`) + view.notice({ text: `${approval.tool} waits for your answer.` }) }) ``` @@ -271,6 +271,11 @@ 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. +The view also keeps its state right on its own: + +- Each event from `client.events()` or `session.events()` has the `threadId` it belongs to. The `RUN_STARTED` event of a child run has the `parentRunId` of the run that started it. Use them to show the run tree in your own UI. +- When the server sends `harness.gap`, some events are gone. The view stops its read, loads the snapshot and the transcript again, then reads on from the new snapshot. So no turn shows twice after the gap. See [Write a client in any language](./protocol#when-events-are-missing). + ## Show a question A command, a durable tool, or an agent run can ask the user with `step.ask`. An MCP server can ask during a tool call. Each item of `state.questions` has: @@ -285,7 +290,7 @@ When a question has a `url`, show the link. The user opens it, then answers. The ```ts group=harness-custom-ui view.on('question', (question) => { - if (question.url) view.notice(`Open ${question.url}, then answer.`) + if (question.url) view.notice({ text: `Open ${question.url}, then answer.` }) }) ``` @@ -350,7 +355,7 @@ This handler shows each one, and answers the region picker right away: ```ts group=harness-custom-ui view.on('interrupt', (request) => { - view.notice(request.message ?? `The turn waits for ${request.reason}.`) + view.notice({ text: request.message ?? `The turn waits for ${request.reason}.` }) if (request.reason === 'pick_region') request.resolve({ region: 'eu-west-1' }) }) ``` @@ -420,7 +425,7 @@ const screenshot = await session.putMedia({ await view.send('What is wrong in this screenshot?', [screenshot]) ``` -In a browser, get the record from `client.upload({ body: file, mimeType: file.type, name: file.name })`. The user message then shows the file in `media`. +In a browser, get the record from `client.putMedia({ body: file, mimeType: file.type, name: file.name })`. The user message then shows the file in `media`. ## Show subagents @@ -429,13 +434,26 @@ A child agent is an `agent` part in the assistant message of the turn that start - `status`: `running`, `done`, or `failed`, and `error` when it failed. - `result`: the typed result of the agent, when it gives one. -After a reload, each child shows up again with what it wrote, not only with its final result. The view reads the saved messages of each child (see [Show what a child wrote after a reload](./subagents#show-what-a-child-wrote-after-a-reload)) and fills its part. A background child that still runs keeps getting its new text in the same part. +After a reload, each saved child shows only its card: `name`, `status`, `error`, and `result`, with no `parts`. To show what the child wrote, call `view.loadAgent({ id })`, for example when the user opens the child. The view reads the saved messages of the child (see [Show what a child wrote after a reload](./subagents#show-what-a-child-wrote-after-a-reload)) and fills its part: + +```ts group=harness-custom-ui +const last = view.store.get().messages.at(-1) +const child = + last?.role === 'assistant' + ? last.parts.find((part) => part.type === 'agent') + : undefined +if (child?.type === 'agent') await view.loadAgent({ id: child.id }) +``` + +A nested child in that part loads the same way, with its own `loadAgent` call. A background child that still runs keeps getting its new text in the same part. + +A terminal screen that prints the old turns only one time (for example Ink `Static`) must load the saved children before its first render. [`examples/harness-cli/src/tui.tsx`](https://github.com/TanStack/ai/blob/main/examples/harness-cli/src/tui.tsx) does this in `loadSavedAgents`. To show a line when a child ends, listen for `'agent'`: ```ts group=harness-custom-ui view.on('agent', (child) => { - if (child.status !== 'running') view.notice(`${child.name} is ${child.status}.`) + if (child.status !== 'running') view.notice({ text: `${child.name} is ${child.status}.` }) }) ``` @@ -450,7 +468,7 @@ import { createPluginEvent } from '@tanstack/ai-harness' const Deployed = createPluginEvent<{ env: string }>('acme/deploy:done') -view.on(Deployed, (done) => view.notice(`Deployed to ${done.env}`)) +view.on(Deployed, (done) => view.notice({ text: `Deployed to ${done.env}` })) ``` `state.plugins` is plugin data, so check its shape before you use it. diff --git a/docs/harness/deploy.md b/docs/harness/deploy.md index 711a7ba31e..b5793fcf9b 100644 --- a/docs/harness/deploy.md +++ b/docs/harness/deploy.md @@ -99,6 +99,82 @@ await buildHarness({ Bun can build for another platform with `target`, for example `bun-darwin-arm64` or `bun-windows-x64`. +## Trace runs and report operations + +A turn is slow or fails on one of your servers, and you cannot see which run or which server. Give the host a tracer and an `onOperation` listener. The host then writes a span for each run, and tells your listener when each run starts and ends. + +Pass an OpenTelemetry tracer as it is. Install `@opentelemetry/api` in your app first. It is an optional peer dependency of the harness. + +```ts +import { trace } from '@opentelemetry/api' +import { createHarnessHost } from '@tanstack/ai-harness' + +const host = createHarnessHost({ + observability: { + tracer: trace.getTracer('support-app'), + onOperation: ({ phase, operation }) => { + if (phase !== 'settled') return + console.log(operation.kind, operation.status, operation.durationMs) + console.log('tokens', operation.usage?.totalTokens) + }, + }, +}) +``` + +The host writes these spans: + +| Span | When | +|---|---| +| `harness.operation` | For each turn, agent run, and command that this host runs. A child agent run on the same host is a child span of the run that started it. | +| `harness.recovery` | This host takes the work of a host that stopped. | +| `harness.wait` | A wait for a person ends. The span starts at the time that the wait opened, also when another host opened it. | +| `harness.transfer` | `host.transfer` moves a conversation. | + +Each span has `harness.thread_id`, and `harness.tenant_id` for a tenant. A `harness.operation` span also has `harness.run_id`, `harness.parent_run_id`, `harness.operation.kind`, `harness.agent`, and `harness.operation.status`. + +`onOperation` gets one `started` report and one `settled` report for each operation: + +- Only the host that runs the operation calls it. A server that only passed the input on reports nothing. +- A `settled` report has `status`, `durationMs`, and `usage`. The usage counts the model calls of the run, of its child runs, and of plugins in the run. +- A failed operation has `error`, with its `message`, and its `code` when it has one. +- If your listener throws, the host sends a `harness.plugin.warning` event and goes on. + +Without `observability`, the host writes no spans and does no work for them. + +### Join a turn to the trace of your app + +Send the `traceparent` of your current span with the message: + +```ts +import { context, propagation } from '@opentelemetry/api' +import { createHarnessClient } from '@tanstack/ai-harness/client' + +const client = createHarnessClient({ url: '/api/harness', threadId: 'thread-1' }) + +const carrier: Record = {} +propagation.inject(context.active(), carrier) +const traceparent = carrier.traceparent +await client.send('Where is my order?', traceparent ? { traceparent } : {}) +``` + +The client sends it as the `traceparent` header. The `harness.operation` span of the turn gets a link to your span, on the server that runs the turn. To send it from another language, see [Traces](./protocol#traces). + +## Known limits + +Check these before you run many servers or long threads. + +Across servers and restarts: + +- A [`harnessAgent`](./subagents#call-a-harness-as-a-child) child waits only in memory. Its parent run waits on the same host. Another host, or a crash, does not resume it. +- A [log snapshot](./durable-sessions#keep-a-long-log-fast-to-open) holds the full fold of the thread, so it grows with the thread. To keep the log small, drop the old records with [`logRecords: 'since-snapshot'`](./lifecycle#drop-the-records-before-a-snapshot). +- `host.resumePending` opens the threads that it takes one at a time. A `GET run` by run id looks through the open sessions one at a time too. + +Users and plugins: + +- With an [OAuth MCP server](./mcp#3-sign-in-to-an-oauth-server), the MCP SDK sends its own refresh request for each session. Only the save of the new token is single-flight, so two sessions can use a rotating refresh token at the same time. +- Plugin tools do not reach background agents. The `question` tool works only in a turn of the lead agent. + ## What you have now - The same harness in your server, on background workers, in a worker process, on another machine, or as a binary. +- A span and a report for each run, linked to the trace of your app. diff --git a/docs/harness/durable-sessions.md b/docs/harness/durable-sessions.md index c4e280688a..7e0295b9fb 100644 --- a/docs/harness/durable-sessions.md +++ b/docs/harness/durable-sessions.md @@ -125,6 +125,8 @@ console.log(`Resumed ${resumed.length} threads`, resumed) It resolves with the threads that it took, as `{ threadId, harness, tenantId }` entries. Each thread opens in its own tenant, with no principal. +If a thread that the sweep took does not open, the host gives it back with a wake in 30 seconds. The next sweep or worker tries it again. + The `close` option sets what happens to a session that the sweep opened: - `'whenIdle'` (default): the session closes when its work ends. If the app opened the same thread too, the session stays open. @@ -276,8 +278,9 @@ On a serverless platform, two requests for one thread often reach two instances. - The first write of a host takes the log. From then on, this host owns the thread. - One timer renews every thread that the host owns. The host gives a thread back when it has no work. -- An input that reaches another instance waits in the log for the owner. The receipt is `queued`, and the owner runs it within a second. -- Every instance reads the thread. A client gets the events from the log, on any instance. A `POST run` on the second instance streams the turn that the owner runs. +- An input that reaches another instance waits in the log for the owner. The receipt is `queued`, and the owner runs it within a second. If the owner refuses it, `session.settled({ inputId })` on that instance rejects with `InputRejectedError`. See [Send inputs safely](./inputs). +- Every instance reads the thread, also while the owner runs a background job. A client gets the events from the log, on any instance. A `POST run` on the second instance streams the turn that the owner runs. +- The thread settings, the config values, and the state of an interrupted turn follow the changes that other hosts write. Set the timing with the `ownership` option of the host: @@ -289,7 +292,7 @@ const ownedHost = createHarnessHost({ ``` - `ttlMs`: how long the ownership lasts after each renewal. The default is 30 000. After it, another host can take the thread. -- `renewMs`: how often the timer renews. The default is 10 000. +- `renewMs`: how often the timer renews. The default is 10 000. When you set only `ttlMs`, the default is a third of `ttlMs`. It must be less than `ttlMs`, else `createHarnessHost` throws a `RangeError`. - `hostId`: a stable id for this host, for example the machine or pod name. With it, a thread whose coding tools use a folder on this host stays on this host. Other hosts refuse its tool calls. See [A folder on one server](./coding-workspaces#a-folder-on-one-server). An input that waits for the owner keeps its JSON overrides (`model`, `modelOptions`, `reasoning`, and `promptCache`). The `adapter` and `tools` overrides are objects, so the owner runs the input without them. diff --git a/docs/harness/durable-tools.md b/docs/harness/durable-tools.md index d4ebd5ff86..7c1d558aa1 100644 --- a/docs/harness/durable-tools.md +++ b/docs/harness/durable-tools.md @@ -127,7 +127,7 @@ The same tool works in plain `chat()`. There, `step.do` runs `fn` each time, and ## What you have now -- Side effects that do not repeat after a crash. +- Steps that keep their result after a crash. A step that a crash cuts can run again, so its idempotency key lets your system refuse the second side effect. - Questions to the user that any server can answer. - Your own records that land together with the tool batch, or not at all. - One tool definition for durable and plain runs. diff --git a/docs/harness/goal.md b/docs/harness/goal.md index e27cb2b2d7..2f39de279b 100644 --- a/docs/harness/goal.md +++ b/docs/harness/goal.md @@ -47,6 +47,8 @@ The first turn starts. Each new turn tells the model the reason of the last chec Keep working on the goal: all tests pass. Last check: The date parser test still fails. ``` +Each new turn runs as the user who set the goal. In a [shared thread](./shared-threads), the goal uses that user's credentials and keys, not the ones of the last person who wrote. + Run `session.command({ name: 'goal' })` to see the result: ```text diff --git a/docs/harness/handoffs.md b/docs/harness/handoffs.md index 963a3288e9..072dba0179 100644 --- a/docs/harness/handoffs.md +++ b/docs/harness/handoffs.md @@ -269,7 +269,7 @@ export async function moveToBilling(request: Request) { `envelope.note` is the first message of the target. `envelope.messages` picks what comes after it: - `'none'` (default): no message. -- `'summary'`: one summary message. The model of the source harness writes it, with the keys of the principal. +- `'summary'`: one summary message. The model of the source harness writes it, with the keys of the principal. The summary call stops after 2 minutes, and then the transfer fails. Its usage counts in the source thread under the source `transfer`. See [Usage of the router and transfers](./usage#usage-of-the-router-and-transfers). - `'all'`: the settled transcript, with new message ids and a copy of each media file. Only settled messages move. A tool call with no result yet stays behind. Plugin state, shared state, credentials, waits, settings, and plugin config never move. The source thread keeps running. diff --git a/docs/harness/inputs.md b/docs/harness/inputs.md index 0879da1f92..1eedb06604 100644 --- a/docs/harness/inputs.md +++ b/docs/harness/inputs.md @@ -63,6 +63,8 @@ console.log(receipt.operationId) Make the id once per user action (a button press), and keep it for every retry of that action. +This works for every input over HTTP and the socket. A `command` or an `agent` input that the client sends again with the same `inputId` also runs once. The `inputId` of its receipt is the id that the client sent. + ## Wait for the outcome On the server, `session.settled({ inputId })` waits until the input ends. It reads the log, so on a [durable host](./durable-sessions) it also works after a restart and on another host that opens the thread. @@ -79,6 +81,8 @@ if (settlement.outcome === 'failed') { `signal` is optional. When it aborts first, `settled` rejects with its reason. The input keeps running. +On a host that does not own the thread, the input waits for the owner. If the owner refuses it, `settled` rejects with `InputRejectedError`. The owner refuses an input when the thread is busy and the input has `busy: 'reject'`, or when the log has its id with another message. The log then has a `harness.offer.rejected` record. + - `completed`: the turn has its answer. - `failed`: the turn failed. `error.code` is set for a limit. - `aborted`: someone cancelled the turn. @@ -207,7 +211,7 @@ for (const input of session.inputs()) { console.log(input.inputId, input.delivery, input.message) } -await session.cancelInput('req-46') +await session.cancelInput({ inputId: 'req-46' }) await session.setDelivery({ inputId: 'req-47', delivery: 'steer' }) await migration ``` @@ -226,7 +230,7 @@ for (const input of waitingInputs) { console.log(input.inputId, input.delivery) } -await client.cancelInput('req-46') +await client.cancelInput({ inputId: 'req-46' }) await client.setDelivery({ inputId: 'req-47', delivery: 'steer' }) ``` @@ -246,11 +250,13 @@ await build - The tool call returns at once. The model gets a short note with the job id, and the turn goes on. - The job keeps running. When it ends, the model gets its result as a note, and an idle session starts a new turn. -- Without an argument, it moves every running call that supports it. Pass a `toolCallId` to move one call. +- Without an argument, it moves every running call that supports it. Pass `{ toolCallId }` to move one call. - `bash` and the single `subagent` tool support it. Other tools keep running as before. - With no running call that supports it, the receipt has `status: 'rejected'` and `reason: 'not_running'`. - The log keeps each job until it ends. When a crash stops a job, the next host adds the note "Background job call-1 stopped when the host restarted." +The session stays open while a job of this host runs, also when no handle holds it. A sweep with `close: 'whenIdle'` waits for these jobs too. The wake turn of a `bash` job runs as the user whose run started the job. + A client sends `{ op: 'background' }` over `POST control`. It runs no new code, so it needs no `expose` entry: ```ts group=harness-inputs-client diff --git a/docs/harness/lifecycle.md b/docs/harness/lifecycle.md index a7e1c3c446..dc6e2c24a7 100644 --- a/docs/harness/lifecycle.md +++ b/docs/harness/lifecycle.md @@ -55,10 +55,10 @@ const busyHost = createHarnessHost({ }) ``` -- `idleCloseMs`: how long an idle session with no handle stays open, in milliseconds. The default is 60 000. `0` closes it at once. +- `idleCloseMs`: how long an idle session with no handle stays open, in milliseconds. The default is 60 000. `0` closes it at once. A session is not idle while a background job of this host runs, for example a `bash` job. - `maxOpen`: how many sessions the host keeps open. The default is 500. An open past the limit first closes the least recently used sessions that are idle and that no handle holds. - If every open session is busy or held, the open goes on, and the host logs one warning. -- A close cleans up the plugins of the session. The transcript and the other data stay in the stores, so the next open continues the thread. +- A close cleans up the plugins of the session. The transcript and the other data stay in the stores, so the next open continues the thread. A workspace that `bind` gave the thread closes with `'moved'`, so an owned sandbox stays. - The idle timer does not keep the Node process alive. A plugin with `lifetime: 'host'` stays set up while another session uses it. See [Share one setup between sessions](./plugins#share-one-setup-between-sessions). @@ -115,11 +115,11 @@ What the view reads for the chat: - `transcript()`: the saved messages. - `snapshot()`: the status and the waiting approvals. When the session is open on this host, it also has the running operations. - `events({ from, signal })`: the events of the thread after a cursor. -- `agentTranscript(id)`: the messages of one child agent run. +- `agentTranscript({ subagentRunId })`: the messages of one child agent run. What the view reads for the rest: -- `getMedia(id)` and `loadMedia({ id, range })`: a media file and its bytes. +- `getMedia({ id })` and `loadMedia({ id, range })`: a media file and its bytes. - `usage()`: the token totals of the thread. - `entry()`: the index entry of the thread. - `logState()`: the `durability.reduce` fold of the log. See [One log for a session tree](./shared-logs#keep-state-for-the-whole-tree). @@ -143,7 +143,7 @@ const leanHost = createHarnessHost({ | --- | --- | --- | | `settledOperations` | 500 | `session.operation(id)` is `undefined` for the oldest settled turns, commands, and resets. | | `settledAgentRuns` | 500 | `session.agentRun(id)` and `session.agentRuns()` leave out the oldest settled agent runs. | -| `eventBytes` | 8 000 000 | `session.events()` keeps the newest events, by the JSON length of each event. A reader with an older cursor starts at the oldest event kept. | +| `eventBytes` | 8 000 000 | `session.events()` keeps the newest events, by the JSON length of each event. A reader with an older cursor reads the older events from the log. | The log keeps what memory drops: @@ -166,6 +166,7 @@ const trimmedHost = createHarnessHost({ - `'all'` (default): the store keeps every record. - `'since-snapshot'`: after each snapshot, the owner truncates the records before it. The type is `LogRecordRetention`. - A reader whose cursor is in the dropped part gets one `harness.gap` event, then the stream goes on. See [Cursors](./connect#cursors). +- `session.revert()` still works. The snapshot keeps the last 500 file steps of [`snapshots()`](./snapshots#go-back-to-an-earlier-message). CAUTION: With `'since-snapshot'`, do not change `durability.reduce` or `durability.project` often. When the host cannot use the snapshot (its `reduce` version or its `project` changed, or the snapshot has a newer version), a thread with dropped records cannot open. The open rejects with `LogTruncatedError`. Keep `'all'` if you change `reduce` often. @@ -184,6 +185,23 @@ First, `delete` closes the sessions of the thread and of its child threads, for - The uploaded media and the generated artifacts. - The log checkpoint in `stores.metadata`. +The first-party plugins also end the thread: + +- [`workspaceTools()`](./coding-workspaces) reattaches an owned sandbox of the thread and closes it with `'closed'`, so it is destroyed. +- [`snapshots()`](./snapshots#remove-the-snapshots-of-a-thread) removes the snapshot repository of the thread. +- This needs the index entry of the thread, and the harness of the thread on this host. + +A host knows a harness when it opened a session of it. After a deploy or a restart, a new host has opened nothing. Give it the harnesses that it can delete threads of: + +```ts group=harness-lifecycle +const deleteHost = createHarnessHost({ + persistence: memoryPersistence(), + harnesses: [assistant], +}) +``` + +A host that does not know the harness of a thread still deletes the data. It warns, and the plugins do not end the thread: an owned sandbox stays. + - On a host with `access`, `delete` needs the `delete` action. A child thread is checked through its root thread. - A handle that calls again after the delete opens a new, empty thread. - A host without `stores.sessions` still closes the open sessions of the thread and deletes its data. It has no index entries, so it finds no child threads. diff --git a/docs/harness/mcp.md b/docs/harness/mcp.md index 4e5fcd5f82..9ab920df7b 100644 --- a/docs/harness/mcp.md +++ b/docs/harness/mcp.md @@ -61,7 +61,7 @@ export const assistant = defineHarness({ }) ``` -The name of a server is the prefix of its tools. A `search` tool of the `docs` server is `docs_search` for the model. +The name of a server is the prefix of its tools. A `search` tool of the `docs` server is `docs_search` for the model. To choose other names, see [Rename the tools of a server](#rename-the-tools-of-a-server). - `stdio`: the server runs as a child process on this machine (Node only). The host starts it one time for each harness, and every session of that harness shares it. It stops when the last session that uses it closes, or at `host.close()`. - `http`: a server that speaks Streamable HTTP. Send an API key in `headers`. Each session connects on its own. @@ -82,9 +82,23 @@ linear: failed: Not signed in. Run /connect linear. In your own UI, read the same list from `session.snapshot().plugins['tanstack/mcp'].servers`. Each entry is an `McpServerStatus`: - `name`: the name of the server. -- `status`: `connecting`, `connected`, or `failed`. A `failed` entry has an `error`. +- `status`: `connecting`, `connected`, `reconnecting`, or `failed`. A `failed` entry has an `error`. - `toolCount`: how many tools the server gives the model. +### When a server stops + +A `stdio` server can crash, and a connection can close. Then `/mcp` shows `files: reconnecting`, and the turns run without the tools of that server. `mcp()` connects a server without `oauth: true` again: + +- It waits 0.5 s, then 1 s, then 2 s, and so on, up to 30 s between tries. A random part makes each wait shorter, so many sessions do not connect at the same moment. +- The waits start again at 0.5 s only after a connection stays open for 30 s. A server that crashes right after it connects waits longer each time. +- It tries until the server connects, or until the session closes. For a `stdio` server, it tries until the host stops the server. +- The next turn after the new connection has the tools again. All sessions that share a `stdio` server use the new process. +- A server that fails its first connect shows `failed`, and tries again with the same waits. A server that is not up yet, like a slow `npx`, gets its tools at a later turn. + +An OAuth server, also one from `mcpConnector`, connects again at the next turn after its connection closes. After `tools/list_changed`, the next turn lists its tools again. If that list fails, the old list stays. + +Some servers change their tools while they run. When a server sends `tools/list_changed`, the next turn lists the tools of that server again, and `/mcp` shows the new count. + ## 3. Sign in to an OAuth server 1. Run `/connect linear` in the CLI. @@ -117,6 +131,52 @@ The model never sees the token. To keep sign-ins after a restart, give the host - A tool of an `oauth: true` server asks for approval before it runs, unless the server marks it read-only. - The tools of the other servers run without a question. +Set `needsApproval` on a server to change this. `true` asks before every tool of the server. A function gets the MCP tool and returns `true` to ask: + +```ts group=harness-mcp +export const careful = mcp({ + servers: { + files: { + type: 'stdio', + command: 'npx', + args: ['-y', '@modelcontextprotocol/server-filesystem', '.'], + // Ask before each tool that the server does not mark read-only. + needsApproval: (tool) => tool.annotations?.readOnlyHint !== true, + }, + docs: { type: 'http', url: 'https://mcp.example.com/mcp', needsApproval: true }, + }, +}) +``` + +## Rename the tools of a server + +Two servers can give a tool with the same name. Set `toolName` on a server to choose the names that the model sees. It gets the MCP tool and returns the name: + +```ts group=harness-mcp +export const renamed = mcp({ + servers: { + docs: { + type: 'http', + url: 'https://mcp.example.com/mcp', + toolName: (tool) => `docs_v2_${tool.name}`, + }, + }, +}) +``` + +When two servers still give the same tool name, the model gets the tool of the server that comes first in `servers`. The other tool is dropped, and the console shows one warning with the names of both servers. + +To treat a clash as a config error, set `onDuplicateTool: 'error'`. Then the later server shows `failed` with an error that names both servers and the tool, and the model gets none of its tools. The other servers keep working: + +```ts group=harness-mcp +export const strict = mcp({ + onDuplicateTool: 'error', + servers: { + docs: { type: 'http', url: 'https://mcp.example.com/mcp' }, + }, +}) +``` + ## Answer a question from a server An MCP server can ask the user for input in the middle of a tool call. MCP calls this elicitation. The session asks the user, and the tool call waits for the answer. @@ -124,6 +184,7 @@ An MCP server can ask the user for input in the middle of a tool call. MCP calls - A form request becomes a question with the `message` and the requested `schema`. - A URL request becomes a question with a `url`. The user opens the page, then answers. The session does not open the page. - `mcp()` and `mcpConnector` both ask this way, for servers on MCP spec 2025 and on spec 2026-07-28. +- A spec 2025 request does not name the tool call that asks. So the session asks the user only when all the running calls of that server belong to one run. Else the server gets `decline`, so a question never goes to another user. - One tool call can ask more than once. On spec 2026-07-28, one call can ask up to 5 times. One more request fails the tool call. - The question is a live [wait](./waits) of the run that made the tool call. The host that runs the call keeps the thread until the answer comes. When that run stops, the wait ends `cancelled`. @@ -218,13 +279,13 @@ export const lateTools = definePlugin({ }) ``` -- A tool whose name is already taken is skipped. The first tool keeps the name. +- A tool whose name is already taken is skipped. The first tool keeps the name. When two owners give the same name, clients get one `harness.plugin.warning` event that names both. - If `discoverTools` throws, the turn runs without that plugin's tools, and clients get a `harness.plugin.warning` event. ## What you have now - Tools from local, remote, and OAuth MCP servers, with the server name as a prefix. -- A status for each server, in `/mcp` and in your UI. +- Servers that connect again after they stop, with a status in `/mcp` and in your UI. - One browser sign-in per service, kept in your credential store. - A question before the tools that change data. - Server requests for input, asked as session questions. diff --git a/docs/harness/media.md b/docs/harness/media.md index 09018f0106..a41d9486d2 100644 --- a/docs/harness/media.md +++ b/docs/harness/media.md @@ -47,13 +47,14 @@ console.log(turn.text) ``` - The transcript keeps a `harness-media:` URL for the file. Only the model call gets the bytes. +- The host keeps up to 32 MB of base64 text in memory, for all its sessions together, so the next turn does not encode a file again. Above that, a turn encodes the oldest files again. - If the harness refuses the file, `putMedia` throws a `MediaError`. Its `status` is `413` for a file over the size limit, and `415` for a kind that the harness does not take. The session also has: -- `getMedia(id)`: the record, or `null` when the thread has no file with this id. +- `getMedia({ id })`: the record, or `null` when the thread has no file with this id. - `loadMedia({ id, range })`: the bytes, or one range of them. `range` is optional. -- `mediaUrl(id)`: `{ url }`, with a data URL for a file up to 1 MB. A bigger file gets `{}`. Use `loadMedia` for it, and serve the bytes from your own route. +- `mediaUrl({ id })`: `{ url }`, with a data URL for a file up to 1 MB. A bigger file gets `{}`. Use `loadMedia` for it, and serve the bytes from your own route. ## Media that agents make @@ -169,6 +170,81 @@ gpt-6.1-sol cannot read audio files. Add media.transcribe, or send a text summar In a child agent, `ctx.chat` gets the bytes of the files too. There, only `accepts` limits the kinds, because the harness does not know the adapter of the child. +## Files inside a client message + +A browser can put a file in the message itself, as base64 bytes or as a URL. Base64 bytes make the transcript, the log, and each model call big. A URL can point to a server that the model provider then fetches. So the harness checks each part of an input from a client before the turn starts. + +The browser sends the bytes in the message: + +```ts group=harness-media-client +import { createHarnessClient } from '@tanstack/ai-harness/client' +import type { assistant } from './harness' + +const client = createHarnessClient({ + url: '/api/harness', + threadId: 'user-1-thread', +}) + +async function toBase64(file: File) { + const bytes = new Uint8Array(await file.arrayBuffer()) + return btoa(Array.from(bytes, (byte) => String.fromCharCode(byte)).join('')) +} + +export async function askAbout(file: File) { + const receipt = await client.send([ + { type: 'text', content: 'What is on this receipt?' }, + { + type: 'image', + source: { type: 'data', value: await toBase64(file), mimeType: file.type }, + }, + ]) + if (receipt.status === 'rejected') console.log(receipt.reason) +} +``` + +On the server, the harness does this for each part: + +- Base64 bytes (a `data` source, or a base64 `data:` URL) go to the media store. The message keeps a `harness-media:` reference, the same as a file from `putMedia`. `maxBytes`, `kinds`, and the model check apply. +- An `http:` or `https:` URL is refused. Set `externalUrls` to allow it. +- A provider file handle (`source: { type: 'file' }`) is refused. The adapter reads the file with your provider key, so a client that knows the file id of another user could make the model read it. Set `allowProviderFiles: true` to allow it. +- A URL with another scheme, for example `file:`, is always refused. + +`POST run` (the route of `useChat`) checks the parts of the last user message the same way. + +Three options control this: + +```ts group=harness-media-inline +import { defineHarness } from '@tanstack/ai-harness' +import { openaiText } from '@tanstack/ai-openai' + +export const assistant = defineHarness({ + name: 'acme/assistant', + adapter: openaiText('gpt-6.1-sol'), + media: { + externalUrls: (url) => url.hostname === 'cdn.acme.com', + allowProviderFiles: false, + maxBodyBytes: 8 * 1024 * 1024, + }, +}) +``` + +- `externalUrls`: `true` allows every `http:` and `https:` URL. A function decides for each URL. Default `false`. +- `allowProviderFiles`: `true` allows provider file handles. Default `false`. +- `maxBodyBytes`: the largest request body or socket frame, in bytes. Default 32 MiB. A bigger body answers `413`. A bigger frame closes the socket. The socket checks a frame only after your WebSocket server read it, so also set the frame limit of your server, for example `maxPayload` of `ws`. + +If the harness refuses a part, no turn starts. The receipt is `rejected`, and its `reason` tells you why: + +| `reason` | Why | +| ----------------------- | -------------------------------------------------------------------------------------------------------------- | +| `media_too_large` | The bytes are bigger than `maxBytes`. | +| `media_not_allowed` | The kind is not in `kinds`, or the type has no kind. | +| `media_invalid` | The bytes are not base64. | +| `media_url_not_allowed` | `externalUrls` refuses the URL, the scheme is bad, or it is a file handle and `allowProviderFiles` is not set. | + +`POST run` has no receipt. It answers with the `reason` as `error`, and with `413` for `media_too_large`, `415` for `media_not_allowed`, or `400` for the other two. + +Server code that calls `session.send` directly is not checked, so it can send file handles. The checks are for inputs from a client. + ## Keep media across restarts The host keeps media in three stores of its persistence: diff --git a/docs/harness/overview.md b/docs/harness/overview.md index 46868300bd..c8aeb8bdc9 100644 --- a/docs/harness/overview.md +++ b/docs/harness/overview.md @@ -205,6 +205,8 @@ Each host keeps every thread in a session log: the transcript, the inputs, the r `memoryPersistence()` has every store, but it keeps them in memory only. To keep the log in your database, implement the [`LogStore` contract](../persistence/store-reference#logstore). Then a turn continues after a crash or a deploy. See [Durable sessions](./durable-sessions). +Before you run many servers or long threads, read the [known limits](./deploy#known-limits). + ## Prompt caching Every session caches the stable start of its requests by default, with its `threadId` as the cache key. To change the default for every session, set `promptCache` on the harness: diff --git a/docs/harness/permissions.md b/docs/harness/permissions.md index 1a2dcab832..7f1d145d3a 100644 --- a/docs/harness/permissions.md +++ b/docs/harness/permissions.md @@ -98,6 +98,8 @@ await session.command({ name: 'mode', input: 'plan' }) `session.command({ name: 'mode' })` with no input shows the current mode. In `acceptEdits` mode, an edit outside `root` or of a `.env` file still asks. +A mode change applies from the next turn. A turn that runs keeps the mode that it started with, so a switch to `bypass` does not change its later tool calls. + ### Keep the mode on the server A client that can set `mode` can set `bypass`. Then every call runs with no question. So by default, a client cannot change the mode: @@ -161,6 +163,8 @@ Two more rules apply to each tool: - A tool with no rule runs without a question. To ask for each tool that has no rule, set `default: 'ask'`. - When the last rule of a tool is a `deny` with no `resource`, the model does not get the tool. +The plugin decides each call in the guard hook, [`onGuardToolCall`](../advanced/middleware#onguardtoolcall). It runs after the input check and after every middleware that changes the arguments. So the rules see the arguments that the tool gets. + ### Paths - Inside `root`, a rule matches the path from `root`, for example `src/app.ts`. Outside `root`, it matches the full path, for example `/etc/hosts`. @@ -224,7 +228,7 @@ The rules decide the formatter as a part of the call that writes, before that ca | `plan` | Neither runs. No file changes in this mode. | - A `reject` stops the call. Nothing is written, and the formatter does not run. -- An `always` answer also saves `formatter:prettier` for the project, so the user answers once. +- An `always` answer also saves the formatter for the project, tied to a hash of its config files (`formatter:prettier@`). The user answers once, until those files change. - A rule that denies the formatter skips it. The write still runs, and the tool result says that the formatter did not run. - Without `permissions()`, the formatter runs without a question. @@ -299,7 +303,7 @@ The same check runs in every agent run of the session: - A background agent run that needs an approval stops with status `waiting`. Its approval is in `pendingWaits` of the same session. Answer it with `resolve`, and the run continues. See [When a background agent waits](./subagents#when-a-background-agent-waits). - An agent that the model calls as a tool inside a turn cannot stop for an approval. A call of it that asks gets the tool error `This call needs the approval of the user, and this run cannot ask for it.` - In `plan` mode, a child cannot change a file. Its call gets the tool error `This tool is not allowed in plan mode.` -- The `permissions` of an [agent profile](./agents) count only while that profile is the primary agent. They do not apply when the profile runs as a subagent. +- The `permissions` of an [agent profile](./agents) count while that profile is the primary agent, in its turns and in each child run. A child run does not use the rules of its own profile. - [Coding agents](./coding-agents#plan-mode) such as Claude Code start read-only in `plan` mode. ## Rules for your own tool diff --git a/docs/harness/plugins.md b/docs/harness/plugins.md index 6676d50630..4497c6e0cd 100644 --- a/docs/harness/plugins.md +++ b/docs/harness/plugins.md @@ -31,7 +31,7 @@ Add it with `plugins: () => [today]` in `defineHarness`. `setup` runs once per s For the model: - `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. +- `prompts`: text for the system prompt. A function runs for each turn, so it can show current state. For text that changes often, see [Put changing text last](#put-changing-text-last). - `middleware`: chat middleware for the lead turn, the same type as `chat({ middleware })`. `agentMiddleware` is for agent runs, see [Middleware in every agent](#middleware-in-every-agent). - `generationMiddleware`: middleware for the activities agents call. - `adapter`: a picker for the model of each turn, see [Pick the model of a turn](#pick-the-model-of-a-turn). @@ -48,6 +48,59 @@ If two plugins add the same tool, command, setting, or prompt id, `host.open` fa Return these from `setup`, not from the top level of `definePlugin`. A key like `adapter` or `tools` at the top level is a type error, and `definePlugin` throws for it. Other keys at the top level stay on the plugin object, so a plugin can give your code its own functions. +### Put changing text last + +Set `placement: 'context'` on a prompt whose text changes often: + +```ts group=harness-plugins +let openTickets = 0 + +export const tickets = definePlugin({ + name: 'acme/tickets', + setup: () => ({ + prompts: [ + { + id: 'acme/tickets:count', + text: () => `Open tickets: ${openTickets}.`, + placement: 'context', + }, + ], + }), +}) +``` + +A provider caches the start of each request: the system prompts, then the history. A change of a system prompt makes the cache miss from that prompt on. A `'context'` prompt is the last system prompt, so the system prompts before it stay in the cache. + +- The model gets all `'context'` text as one system prompt, after the other system prompts. The harness makes it again at each model call. +- The text has the authority of a system prompt, not of a user message. +- The default is `'system'`: the text goes in the system prompts, in plugin order. +- The log does not keep `'context'` text. +- The first-party todos, goal, and plan mode prompts use `'context'`. + +### Refresh state at the start of a turn + +A prompt function must return at once, so it cannot read a file or call a server. Do that work in `onTurnStart`. The harness calls it once at the start of each turn, before it reads the prompts: + +```ts group=harness-plugins +import { readFile } from 'node:fs/promises' + +let rules = '' + +export const projectRules = definePlugin({ + name: 'acme/project-rules', + setup: () => ({ + onTurnStart: async () => { + rules = await readFile('RULES.md', 'utf8') + }, + prompts: [() => rules], + }), +}) +``` + +- `onTurnStart` gets the turn. `turn.run` has the user of the turn and the credentials of that user. +- It runs one time for each turn, also when the turn makes many model calls. +- If it throws, the turn continues, and clients get a `harness.plugin.warning` event. + ## Add commands and settings Commands are actions for the user, not the model: slash commands and buttons in your app. Settings are typed, saved per session, and apply at the next turn. @@ -196,6 +249,7 @@ export const counter = definePlugin({ - A key or a token of the session in the state goes into the log as `[redacted]`. `update` resolves with that same value. See [Keep tokens out of the log and the model](./auth#keep-tokens-out-of-the-log-and-the-model). - The updates of one plugin run one at a time. Each `update` function gets the value that the update before it wrote. - [`host.sessions.delete`](./lifecycle#delete-a-thread-for-good) deletes the state with the log. A fork does not copy it. +- For a long list of records, such as one record for each step, append records to the log and set `keepRecords`. Read them with `ctx.session.keptRecords(type)`. See [Keep records after a truncate](./session-log#keep-records-after-a-truncate). - Agents get a state with the same `get` and `update` from `ctx.harness.state({ key, initial })`. See [Share state between agents](./subagents#share-state-between-agents). ## Middleware in every agent @@ -271,6 +325,10 @@ Where each part of a plugin reads the run: - A session plugin (the default lifetime): `ctx.credentials` and `ctx.keys` have no user, so each call throws. `ctx.session.principal` is the owner of the session. - A plugin with `lifetime: 'host'`: its setup has no user and no session. Its tools read the user of each call from `ctx.context.run`. +`prepareTools` gets `run` too: the run of the turn, with its id and its user. `onTurnStart` gets the turn, with its run in `turn.run`. + +A plugin that starts a turn picks its user. `ctx.session.send(text, { principal })` and `ctx.session.note(text, { wake: true, principal })` run that turn as `principal`. Without it, the turn runs as `ctx.session.principal`. To continue the work of a run, pass the `principal` of that run. + The run has these fields: | Field | What it is | @@ -284,6 +342,44 @@ The run has these fields: | `signal` | Aborts when the run stops. | | `credentials`, `keys` | The credentials and the model provider keys of the user of the run. The session replaces their values with `[redacted]` in what it keeps and sends. To refresh an OAuth token, use `credentials.refresh`. See [Auth and connectors](./auth#refresh-a-token). | +## Call a model from a plugin + +Your plugin needs a model call of its own: a judge, a title, or a summary. A plain `chat()` call does not count in `session.usage()`, and it keeps running after the user cancels the turn. Call the model with `ctx.model.chat` instead: + +```ts group=harness-plugins +export const recap = definePlugin({ + name: 'acme/recap', + setup: (ctx) => ({ + commands: { + recap: defineCommand({ + description: 'Sum up the conversation in three lines', + run: async (_input, { run }) => { + const transcript = await ctx.session.transcript() + const { text } = await ctx.model.chat({ + run, + label: 'recap', + options: { + messages: [...transcript, { role: 'user', content: 'Sum up this conversation in three lines.' }], + }, + }) + return text + }, + }), + }, + }), +}) +``` + +The options: + +- `run`: the run that the call is for. The call uses the keys of its user, and the budget of the run applies. A plugin with `lifetime: 'turn'` can leave it out. Then the call uses the run of its turn. +- `adapter`: the model. Without it, the call uses the model of the session: the thread setting `model`, else the harness `adapter`. A `keyedAdapter(...)` gets the key of the user of the run. +- `options`: `messages`, `systemPrompts`, and `outputSchema`. With `outputSchema`, the call resolves to the checked value. Without it, the call resolves to `{ text, chunks }`. +- `label`: a name for the call in its usage record. +- `timeoutMs`: the time limit of the call. The default is 2 minutes. + +The call stops when its run stops, for example on a cancel, and when the time limit ends. Then it rejects. Its usage counts in the session under `plugin:`, here `plugin:acme/recap`. See [Usage of plugins](./usage#usage-of-plugins). The `goal()`, `title()`, and `compact()` plugins call their models this way. + ## Pick the model of a turn Some turns need another model: a short question does not need your strongest model, and an organization can allow only some models. Return an `adapter` picker from `setup`. When each turn starts, the session calls it with the message, context, and sender of that turn: @@ -521,6 +617,7 @@ If a `setup` throws, the session keeps working with no plugins. The promise reje - A plugin that adds tools, prompts, commands, settings, and state to any harness. - Middleware that sees every model call, in the lead turn and in every agent. - Tools and commands that act as the user of their own run. +- Model calls of your plugin that count in the session usage and stop with their run. - A picker that chooses the model of each turn from its message, context, and sender. - Plugins that share services, lists, and events without knowing each other. - Commands, agents, and whole plugins that change while the session runs. diff --git a/docs/harness/protocol.md b/docs/harness/protocol.md new file mode 100644 index 0000000000..770697b02b --- /dev/null +++ b/docs/harness/protocol.md @@ -0,0 +1,480 @@ +--- +title: The harness wire protocol +id: harness-protocol +order: 3 +description: "Every route, frame, and harness.* event of the harness session stream. Use it to write your own client or transport in any language." +keywords: + - tanstack ai + - harness + - protocol + - AG-UI + - websocket + - server-sent events +--- + +You want a client for a harness that is not `useChat` or `createHarnessClient`: a Python service, a mobile app, or a proxy of your own. This page lists every route, frame, and `harness.*` event, so you can write that client in any language. + +Send a prompt to a thread with `curl`: + +```bash +curl -X POST https://example.com/api/harness/control \ + -H 'authorization: Bearer my-token' \ + -H 'content-type: application/json' \ + -d '{"threadId":"thread-1","input":{"op":"prompt","message":"Hello"}}' +``` + +The command prints the receipt, for example `{"inputId":"in-1","status":"accepted","operationId":"op-1"}`. The first prompt also makes the thread. Then open the session stream of the thread: + +```bash +curl -N -H 'authorization: Bearer my-token' \ + 'https://example.com/api/harness/events?threadId=thread-1&v=2' +``` + +With no cursor, the stream starts at the first event of the thread. It shows the hello frame, then one SSE event per session event, and a ping comment when the stream is idle: + +```text +data: {"type":"harness.hello","v":2,"threadId":"thread-1"} + +id: 4 +data: {"type":"harness.event","cursor":"4","threadId":"thread-1","operationId":"op-1","event":{"type":"TEXT_MESSAGE_CONTENT","messageId":"m-1","delta":"Hi"}} + +: ping +``` + +The rest of this page is the reference for each part. + +## Transports + +The server is the handler of [Connect clients to a harness](./connect). One handler and one socket serve every transport on this page: + +```ts group=harness-protocol-server +import { createHarnessHost, defineHarness, membersOnly } from '@tanstack/ai-harness' +import { createHarnessHandler, handleHarnessSocket } from '@tanstack/ai-harness/server' +import { memoryPersistence } from '@tanstack/ai-persistence' +import { openaiText } from '@tanstack/ai-openai' +import type { WebSocketLike } from '@tanstack/ai' + +export const assistant = defineHarness({ + name: 'acme/assistant', + adapter: openaiText('gpt-6.1-sol'), +}) + +const host = createHarnessHost({ + persistence: memoryPersistence(), + access: { + identify: ({ request }) => + request.headers.get('authorization') === `Bearer ${process.env.HARNESS_TOKEN}` + ? { id: 'user-1' } + : null, + authorize: membersOnly(), + }, +}) + +// Mount it on `/api/harness/*`. +export const handler = createHarnessHandler({ host, harness: assistant }) + +// Call it from the upgrade handler of your server. +export function onUpgrade(socket: WebSocketLike, request: Request) { + handleHarnessSocket({ host, harness: assistant, socket, request }) +} +``` + +Pick the transport that your client can use: + +| Transport | Events | Inputs | Use it when | +| --- | --- | --- | --- | +| HTTP | `GET events` (SSE) | `POST control` (JSON) | Your client can read SSE. This works through most proxies. | +| WebSocket | `harness.event` frames | `harness.input` frames | You want one connection for events and inputs. | +| Stdio worker | NDJSON on stdout | NDJSON on stdin | You start the harness as a child process. | + +The socket and the stdio worker use the same frames. The HTTP routes send the same `harness.*` frames as SSE `data`. + +### HTTP routes + +These routes are enough for a full client. [Connect clients to a harness](./connect#serve-the-session-over-http) lists the others (media, sessions, and more). + +- `GET events?threadId=&v=&from=`: the session stream as SSE. Each SSE `data` is one frame. On a host with `access`, a thread that has no session yet answers `404`, so send the first input before you open it. +- `POST control`: the body is `{ threadId, input }`. The answer is the receipt as JSON. +- `GET snapshot?threadId=` and `GET transcript?threadId=`: the state and the saved messages of the thread. When the client starts, load both. +- `GET capabilities`: the AG-UI capabilities. `custom.tanstack.protocol` is the protocol version of the server. +- `GET host-events?v=`: the status changes and session list changes that the user may read, as SSE. Its `harness.hello` frame has no `threadId`. +- `POST run` and `GET run`: standard AG-UI for `useChat`. These routes send plain AG-UI events, with no `harness.*` frames around them. + +`POST control` answers with these statuses: + +- `200`: the receipt, `{ inputId, status, operationId?, reason? }`. +- `400`: `{ error }`. The input has a bad shape, or `threadId` is missing. +- `401`: `identify` returned `null`. +- `403`: `{ error: 'forbidden' }`. The host refused the action. + +A `rejected` receipt is still `200`. Its `reason` says why, for example `not_exposed`. + +### The WebSocket + +The first frame from the client must be `harness.subscribe`. The host reads the user from the upgrade request, with `identify`. A browser cannot set headers on a WebSocket, so `identify` must read a cookie or a token in the URL. + +```ts group=harness-protocol-socket-client +import type { ControlFrame, HostFrame } from '@tanstack/ai-harness/server' + +function isFrame(value: unknown): value is HostFrame { + return ( + typeof value === 'object' && + value !== null && + 'type' in value && + typeof value.type === 'string' + ) +} + +const socket = new WebSocket('wss://example.com/api/harness/socket') +const send = (frame: ControlFrame) => socket.send(JSON.stringify(frame)) + +socket.addEventListener('open', () => { + send({ type: 'harness.subscribe', threadId: 'thread-1', v: 2 }) +}) + +socket.addEventListener('message', (message) => { + const frame: unknown = JSON.parse(String(message.data)) + if (!isFrame(frame)) return + if (frame.type === 'harness.hello') { + send({ + type: 'harness.input', + requestId: 'req-1', + input: { op: 'prompt', message: 'Hello' }, + }) + } + if (frame.type === 'harness.receipt') console.log(frame.requestId, frame.status) + if (frame.type === 'harness.event') console.log(frame.cursor, frame.event.type) + if (frame.type === 'harness.error') console.warn(frame.code, frame.message) +}) +``` + +The socket closes with these codes: + +| Code | Why | What the client does | +| --- | --- | --- | +| `4401` | `identify` refused the user. | Sign in again. | +| `4403` | The user cannot read the thread, or lost access. | Stop. | +| `4400` | The server does not speak the `v` of the subscribe. | Stop. Use a client of the version of the server. | +| `4408` | The client read too slowly. See [Stream values](#stream-values). | Subscribe again from the last cursor. | +| `1011` | The event feed of the thread ended. | Subscribe again from the last cursor. | + +### The stdio worker + +`runHarnessWorker` reads frames from stdin and writes frames to stdout, one JSON object per line. Put it in its own script: + +```ts group=harness-protocol-worker +import { defineHarness } from '@tanstack/ai-harness' +import { runHarnessWorker } from '@tanstack/ai-harness/worker' +import { openaiText } from '@tanstack/ai-openai' + +const assistant = defineHarness({ + name: 'acme/assistant', + adapter: openaiText('gpt-6.1-sol'), +}) + +await runHarnessWorker(assistant) +``` + +Start the script as a child process, then write frames to it: + +```ts group=harness-protocol-worker-client +import { spawn } from 'node:child_process' +import { createInterface } from 'node:readline' +import type { ControlFrame } from '@tanstack/ai-harness/server' + +const worker = spawn('node', ['worker.js'], { stdio: ['pipe', 'pipe', 'inherit'] }) +const send = (frame: ControlFrame) => worker.stdin.write(`${JSON.stringify(frame)}\n`) + +send({ type: 'harness.subscribe', threadId: 'thread-1', v: 2 }) +send({ type: 'harness.input', requestId: 'req-1', input: { op: 'prompt', message: 'Hello' } }) + +for await (const line of createInterface({ input: worker.stdout })) { + console.log(line) +} +``` + +- The worker runs the frames in order, so an input right after the subscribe is safe. +- The worker has no login. Every input runs with no user. +- The worker sends no heartbeat. +- When stdin closes, the worker exits. It also exits after an `unsupported_version` error. + +To ship a harness as a worker with a checked bundle, see [Deploy a harness](./deploy). + +## Frames + +Each frame is one JSON object with a `type`. Ignore a frame with a `type` that you do not know. A later version of the server can add frames. + +From the client to the host (socket and stdio worker): + +| `type` | Fields | What it does | +| --- | --- | --- | +| `harness.subscribe` | `threadId`, `from?`, `v?` | The first frame. Starts the event stream after the cursor `from`. | +| `harness.input` | `requestId`, `input`, `traceparent?` | Sends one input. `input` is the same object as the `input` of `POST control`. | +| `harness.snapshot` | none | Asks for the snapshot of the session. | + +`input` is one of the ops of [Send inputs to a session](./inputs), for example `{ op: 'prompt', message }` or `{ op: 'cancel', operationId }`. Pick a new `requestId` for each input. The receipt and any error frame repeat it. + +From the host to the client: + +| `type` | Fields | When | +| --- | --- | --- | +| `harness.hello` | `v`, `threadId` | First, after a subscribe that the host accepts. On SSE, the first event of `GET events`. | +| `harness.receipt` | `requestId`, `status`, `inputId?`, `operationId?`, `reason?` | The answer to one `harness.input`. | +| `harness.event` | `cursor`, `threadId`, `operationId`, `event` | One event of the session. | +| `harness.snapshot` | `snapshot` | The answer to a `harness.snapshot` frame. | +| `harness.error` | `message`, `code`, `requestId` | Something failed. See [Errors](#errors). | +| `harness.ping` | none | A heartbeat on the socket. Ignore it. | + +### The event frame + +- `event` is an AG-UI event in its wire form. It has only the keys of the AG-UI spec. +- `operationId` is the run that made the event. It is also the AG-UI `runId`. +- `threadId` is the thread that keeps the event. On the stream of a thread, it is that thread. +- The AG-UI `RUN_STARTED` event of a child run has `parentRunId`: the run that started it. A child run is an agent run that agent code started, or that a tool of a turn started. A top-level run has no `parentRunId`. +- An event of a subagent also has `subagentRunId` on the `event` itself. + +The stream of a thread has the events of every run of its tree, in one order. To rebuild the tree, keep the `parentRunId` of the `RUN_STARTED` of each `operationId`. For example, a lead agent that runs two helpers in a group gives these frames (the `event` is cut): + +```text +{"type":"harness.event","cursor":"9","threadId":"thread-1","operationId":"op-lead","event":{"type":"RUN_STARTED","runId":"op-lead",...}} +{"type":"harness.event","cursor":"12","threadId":"thread-1","operationId":"op-a","event":{"type":"RUN_STARTED","runId":"op-a","parentRunId":"op-lead",...}} +{"type":"harness.event","cursor":"13","threadId":"thread-1","operationId":"op-b","event":{"type":"RUN_STARTED","runId":"op-b","parentRunId":"op-lead",...}} +``` + +`op-a` and `op-b` are children of `op-lead`. + +The status of a receipt is `accepted`, `queued`, or `rejected`. A `rejected` receipt has a `reason`, for example `not_exposed`, `not_running`, or `unknown_agent`. + +Give each input an `inputId`, and send the same id again after a lost answer. A `prompt`, `command`, `agent`, or `append` input with the same `inputId` runs once. The `inputId` of its receipt is the id that you sent. + +The socket handles the frames of one connection in the order they came. + +## Cursors and reconnects + +Each `harness.event` has a `cursor`. On SSE, the `id` of each event is the same cursor. A cursor is an opaque string. Store the last one, and do not parse it. + +To continue after a disconnect, connect again with the last cursor: + +- On SSE, send it as the `Last-Event-ID` header, or as `?from=`. When you send both, the header wins. +- On the socket and the stdio worker, send it as `from` in `harness.subscribe`. + +```bash +curl -N -H 'authorization: Bearer my-token' -H 'Last-Event-ID: 4' \ + 'https://example.com/api/harness/events?threadId=thread-1&v=2' +``` + +The stream starts with the first event after that cursor. A disconnect does not stop a running turn. To stop a turn, send a `cancel` input. + +A cursor works on every server and after a restart only on a durable host (`stores.log`). See [Cursors](./connect#cursors). + +### When events are missing + +Sometimes the server no longer has the events after your cursor. A truncate of the log removes old records, for example. Then the stream sends one `CUSTOM` event with the name `harness.gap`, and continues from the first event that it has: + +```json +{ + "type": "harness.event", + "cursor": "120", + "threadId": "thread-1", + "operationId": "session", + "event": { + "type": "CUSTOM", + "name": "harness.gap", + "value": { "from": "12", "resumeAt": "120", "reason": "truncated" } + } +} +``` + +- `from` is the cursor that you sent. +- `resumeAt` is the cursor of this event. The stream continues after it. +- `reason` is `truncated` (a truncate of the log) or `evicted` (the server memory dropped old events). + +When you get `harness.gap`, do these steps: + +1. Delete the state that you built from events. +2. Load `GET snapshot` and `GET transcript` again. On the socket, send `harness.snapshot`. +3. Store `resumeAt` as your cursor, and continue to read the stream. + +A [session view](./custom-ui) does these steps for you. + +## Protocol version + +The protocol version is `2`. + +- On SSE, send `?v=2` on `GET events`. +- On the socket and the stdio worker, send `v: 2` in `harness.subscribe`. +- With no `v`, the server uses its current version. + +The `harness.hello` frame gives the version that the server uses. If the server does not speak your `v`, it sends one error frame and closes the stream: + +```json +{ + "type": "harness.error", + "message": "Protocol version 3 is not supported. This host speaks version 2.", + "code": "unsupported_version", + "requestId": "req-1" +} +``` + +The server speaks one version. Use a client of that version. + +## Errors + +A `harness.error` frame has a `code` and a `requestId`. `message` is for people, so read `code` in your code. + +| `code` | Why | Does the stream close? | +| --- | --- | --- | +| `unauthorized` | `identify` refused the user. | Yes. | +| `forbidden` | The user cannot read the thread, or lost access. On the socket, also when the host refuses one input. | Yes. When the host refuses one input, no. | +| `invalid_frame` | The frame is not valid JSON, has an unknown `type` or a bad `input`, comes before `harness.subscribe`, or is a second `harness.subscribe`. | No. | +| `unsupported_version` | The server does not speak the `v` that you sent. | Yes. | +| `slow_consumer` | The client read too slowly. | Yes. | +| `internal` | The server failed to read the events, or failed to apply one input. | When the read failed, yes: connect again from the last cursor. When one input failed, no. | + +On HTTP, a refused login or action before the stream starts is the `401` or `403` status of the request, with no frame. + +`createHarnessClient` reads these frames for you. After `slow_consumer`, `events()` and `hostEvents()` connect again. Any other code ends the read with a `HarnessProtocolError`, which has the `code` and the `message`. + +The `requestId` connects the error to its request: + +- On the socket and the stdio worker, it is the `requestId` of the frame that failed. A frame with no `requestId` gets an id from the server. On the socket, that id is the id of the upgrade request. +- On HTTP, the id is your `X-Request-Id` header when it is a short safe id: 1 to 128 letters, digits, `-`, `_`, `.`, or `:`. Else the server makes a new UUID. +- Every HTTP answer of the handler sends the id back in the `X-Request-Id` header. + +Log the `requestId` with your errors. Then you can find the request in the server logs. + +## Stream values + +A long stream passes through proxies and slow networks. The server keeps it open and protects itself with these values: + +| Value | Option | Default | What happens | +| --- | --- | --- | --- | +| Heartbeat | `heartbeatMs` | 20 s | SSE gets a `: ping` comment. The socket gets `{"type":"harness.ping"}`. | +| SSE queue | `sseQueueBytes` | 1 MiB | The server stops reading the log while this much data waits for the client. | +| Socket buffer | `socketBufferBytes` | 8 MiB | Above this `bufferedAmount`, the server stops sending. It checks again every 50 ms. | +| Slow close | `slowConsumerMs` | 30 s | If the client does not read for this long, the server sends `slow_consumer` and closes. | + +- `POST run` and `GET run` also get the `: ping` comment. SSE parsers skip comments. +- A slow client loses no events. The server reads them from the log when the client reads again. +- Set the idle timeout of your proxy to more than the heartbeat. +- A client can use the heartbeat to find a dead connection. `createHarnessClient` connects again when no bytes come for 60 s (its `idleTimeoutMs`). If you set `heartbeatMs` to 60 s or more, set a larger `idleTimeoutMs` on the client. +- The stdio worker has no heartbeat and does not use these values. + +To change a value, pass `stream` to the handler and to the socket. A value that you do not set keeps its default: + +```ts group=harness-protocol-server +// A proxy that closes idle streams after 15 s, and mobile clients that read slowly. +const stream = { heartbeatMs: 10_000, slowConsumerMs: 60_000 } + +export const slowNetworkHandler = createHarnessHandler({ host, harness: assistant, stream }) + +export function onSlowNetworkUpgrade(socket: WebSocketLike, request: Request) { + handleHarnessSocket({ host, harness: assistant, socket, request, stream }) +} +``` + +Each value must be a finite number above 0. If a value is not, `createHarnessHandler` and `handleHarnessSocket` throw at the start, and the error names the field. + +## Traces + +To join a turn to the trace of your app, send a W3C `traceparent`: + +- On HTTP, send the `traceparent` header on `POST control` or `POST run`. +- On the socket and the stdio worker, set `traceparent` on the `harness.input` frame. + +```ts +const response = await fetch('/api/harness/control', { + method: 'POST', + headers: { + 'content-type': 'application/json', + traceparent: '00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01', + }, + body: JSON.stringify({ + threadId: 'thread-1', + input: { op: 'prompt', message: 'Hello' }, + }), +}) +console.log(response.headers.get('X-Request-Id')) +``` + +What the harness does with it: + +- It reads the `traceparent` of a message input only: `prompt`, `steer`, `followUp`, and the turn of `POST run`. Other inputs ignore it. +- It drops a value that is not a valid W3C `traceparent`. +- It stores the value on the `harness.input` record of the log. When another server owns the thread, the value goes to that server with the input. +- When the input starts a turn, the server that runs the turn gives its `harness.operation` span a link to the value. A link is not a parent, so your trace UI shows the turn as a linked trace. + +To give the host a tracer, see [Trace runs and report operations](./deploy#trace-runs-and-report-operations). + +## Event catalog + +A harness adds its own events to the stream as AG-UI `CUSTOM` events. `name` is the event name, and `value` has the fields below. Ignore a name that you do not know. A later version can add names. + +In TypeScript, `HarnessEventValues` from `@tanstack/ai-harness` maps each name to its value type, and `HARNESS_EVENTS` has the names. `harness.plugin.warning` is the only name that is in neither. + +### Runs + +| `name` | `value` | +| --- | --- | +| `harness.operation.started` | `{ operationId, kind, agent? }`. `kind` is `chat`, `agent`, `command`, or `compact`. | +| `harness.operation.finished` | `{ operationId, status }`. `status` is `completed`, `failed`, `cancelled`, or `interrupted`. | +| `harness.operation.resumed` | `{ operationId, resumedFrom? }`. A host continued a run after another host stopped. | +| `harness.turn.retry` | `{ operationId, retries, error: { message, code?, retryAfterMs? }, continued }`. `continued` is true when the model continues a partial answer. | +| `harness.handoff` | `{ from, to, note?, runId, by, waitId? }`. `from` and `to` are an agent name, `null` for the main model, or `{ human: true }`. `by` is `tool`, `code`, `router`, or `operator`. A handoff to a person has the `waitId` of its wait. | +| `harness.usage` | `{ model, sender?, source?, label?, usage, total }`. `usage` is this model call, and `total` is the new total of the thread. `source` is `plugin:`, `router`, or `transfer`. A turn or an agent run has no `source`. | + +### Inputs + +| `name` | `value` | +| --- | --- | +| `harness.input.accepted` | `{ inputId, op }` | +| `harness.input.queued` | `{ inputId, delivery }`. The owner of the log stored the input, and it waits behind the running turn or the turns before it. `delivery` is `steer` or `queue`. | +| `harness.input.applied` | `{ inputId, operationId }` | +| `harness.input.rejected` | `{ inputId, reason }` | +| `harness.input.settled` | `{ inputId, operationId?, outcome, error? }`. `outcome` is `completed`, `aborted`, `interrupted`, or `failed`. Only `failed` has `error: { message, code? }`. | +| `harness.input.delivery` | `{ inputId, delivery }`. A `setDelivery` input moved a waiting input. | + +### Waits for a person + +| `name` | `value` | +| --- | --- | +| `harness.wait.opened` | `{ waitId, kind, runId, message?, expiresAt? }`. A wait opened in the log: a question, a `step.ask`, a handoff to a person, or a wait of a background agent run. `kind` is `approval`, `question`, `sign-in`, `client-tool`, or `human`. `expiresAt` is in epoch milliseconds. The interrupts of a turn come in its `RUN_FINISHED`. | +| `harness.wait.settled` | `{ waitId, outcome }`. `outcome` is `answered`, `cancelled`, or `expired`. | +| `harness.question` | `{ questionId, message, schema?, secret?, url? }`. Answer it with the `answer` input. | +| `harness.question.answered` | `{ questionId }` | +| `harness.auth_required` | `{ connector, url?, userCode? }`. A tool needs a sign-in. | + +### Session changes + +When you get one of these events, load the part of the session that it names again. + +| `name` | `value` | +| --- | --- | +| `harness.config.changed` | `{ key, value }` | +| `harness.settings.changed` | `{ settings }` | +| `harness.commands.changed` | `{}`. Load `GET describe` again. | +| `harness.reloaded` | `{ error? }`. The plugins started again. Load `GET describe` again. | +| `harness.reset` | `{ inputId, note? }`. The next turn starts a fresh model context. | +| `harness.revert` | `{ messageId }`. `messageId` is `null` when no revert stands. Load `GET transcript` again. | + +### Files, plugins, and hosts + +| `name` | `value` | +| --- | --- | +| `harness.media` | A media record: `{ id, threadId, kind, mimeType, name, size, createdAt, source }`. See [Send and show media](./media). | +| `harness.plugin.event` | `{ plugin, name, value }`. An event that a plugin sent. | +| `harness.plugin.warning` | `{ plugin, message }`. A plugin hook, a listener, or a write that no one waits for failed. The work goes on. | +| `harness.ownership.changed` | `{ epoch, reason }`. A new host owns the log, from its first write. `reason` is `taken_over` when the ownership of the old owner expired with no release (it stopped), else `acquired`. | +| `harness.stranded` | `{ operationId?, hostId? }`. The workspace that the log binds the thread to is gone, so the workspace tools fail. `operationId` is the run that tried it. `hostId` is the host of a folder. | +| `harness.transferred` | `{ transferId, from?, to? }`. `from` and `to` are `{ threadId }`. The conversation came from, or went to, another thread. | +| `harness.gap` | `{ from, resumeAt, reason }`. See [When events are missing](#when-events-are-missing). | + +## What you have now + +- A client that reads the session stream over SSE, a WebSocket, or a stdio worker. +- Inputs with a receipt for each one, and errors that name their request. +- Reconnects that continue from the last cursor, and a reload when events are missing. +- A stream that stays open through proxies. +- The name and the value of every `harness.*` event. + +Next: send your first prompt with the two `curl` commands at the top of this page, and watch the events arrive. diff --git a/docs/harness/session-log.md b/docs/harness/session-log.md index 302b9656b0..c64f9dd31d 100644 --- a/docs/harness/session-log.md +++ b/docs/harness/session-log.md @@ -32,15 +32,42 @@ const assistant = defineHarness({ const host = createHarnessHost({ persistence }) const session = await host.open(assistant, { threadId: 'thread-1' }) -await session.append([{ type: 'app.signal', text: 'The build failed.' }]) +await session.append({ records: [{ type: 'app.signal', text: 'The build failed.' }] }) ``` - A record is JSON with a `type`. Types that start with `harness.` belong to the harness, so `append` refuses them. -- `append` resolves when the records are in the log. +- `append` resolves with a receipt when the records are in the log. When another host owns the thread, the receipt is `queued`, and the owner writes the records. See [One host drives a thread](./durable-sessions#one-host-drives-a-thread). +- Pass `inputId` to make the append safe to send again: `session.append({ records, inputId: 'signal-7' })`. A second append with the same id writes nothing. Without `inputId`, the append gets a new id. The receipt has the id, as the log keeps it. `principal` sets who appends. The default is the user who opened the session. - A record can belong to a child thread of the same log. See [One log for a session tree](./shared-logs). - In a tool, use `append` from [`durableTool`](./durable-tools), so the records land with the tool batch. - A key or a token of the session in a record goes into the log as `[redacted]`. See [Keep tokens out of the log and the model](./auth#keep-tokens-out-of-the-log-and-the-model). +### Keep records after a truncate + +With [`logRecords: 'since-snapshot'`](./lifecycle#drop-the-records-before-a-snapshot), the store drops the records before each snapshot. A plugin can make the snapshot keep the last records of its own types with `keepRecords`, and read them back with `ctx.session.keptRecords(type)`: + +```ts group=harness-session-log +import { defineCommand, definePlugin } from '@tanstack/ai-harness' + +export const steps = definePlugin({ + name: 'acme/steps', + keepRecords: { 'acme/steps:step': 100 }, + setup: (ctx) => ({ + commands: { + steps: defineCommand({ + description: 'Count the kept steps', + run: async () => (await ctx.session.keptRecords('acme/steps:step')).length, + }), + }, + }), +}) +``` + +- The snapshot keeps the last 100 `acme/steps:step` records of each thread. +- `keptRecords` gives the kept records of that type of this thread, oldest first. It works before and after a truncate. +- [`snapshots()`](./snapshots#go-back-to-an-earlier-message) keeps its last 500 file steps this way, so `session.revert()` works after a truncate. +- The plugins of the harness of the root thread decide, because that harness folds the log. A type in two plugins keeps the larger count. + ## Fold records into the model context By default, the model does not see your records. Set `durability.project` on the harness to fold a record into the messages. Its `record` function gets one object, `{ messages, record }`, and returns the new list, or `undefined` for no change: @@ -122,7 +149,7 @@ const client = createHarnessClient({ headers: { authorization: 'Bearer my-token' }, }) -await client.append([{ type: 'app.signal', text: 'The build failed.' }]) +await client.append({ records: [{ type: 'app.signal', text: 'The build failed.' }] }) await client.continue() ``` diff --git a/docs/harness/sessions.md b/docs/harness/sessions.md index 4cd6e5d547..d7c31d01f8 100644 --- a/docs/harness/sessions.md +++ b/docs/harness/sessions.md @@ -143,7 +143,7 @@ A fork keeps the old chat. To go back in the same chat, revert it to a message: const [question] = await session.transcript() if (!question?.id) throw new Error('The chat has no messages.') -await session.revert(question.id) +await session.revert({ messageId: question.id }) // The transcript now ends at `question`. To bring the later messages back: await session.unrevert() ``` @@ -257,7 +257,9 @@ void watch() - `session-deleted`: the host removed an entry. `entry` is the removed entry. - `revoked`: `host.revoke` or `unshare` took the access of the user `principalId` to `threadId`. See [Use your own access tables](./access#use-your-own-access-tables). -Each host reports only the sessions that are open in it. A session with no status event is not open, so show it as idle. +Each host reports only the sessions that are open in it. For those sessions, the status also counts the turns and the questions that another host writes to the same log. A session with no status event is not open, so show it as idle. + +A reader that does not keep up stops. When more than 1 000 changes wait for it, its loop ends. The first statuses do not count. Call `host.events()` again: the current statuses come first again. `host.close()` does not end the readers of `events()`. Abort the signal of each reader before you close the host. If you do not abort it, the loop waits forever: diff --git a/docs/harness/shared-logs.md b/docs/harness/shared-logs.md index dd811a1bf8..50ebc329db 100644 --- a/docs/harness/shared-logs.md +++ b/docs/harness/shared-logs.md @@ -40,7 +40,7 @@ const root = await host.open(assistant, { threadId: 'agent-1' }) - The log has the id of the root thread, in its tenant: `logIdOf({ tenantId, logId: 'agent-1' })`. - Each thread folds only its own transcript, inputs, and events from the log. A child never sees the messages of the root. -- A child thread does not open. Open its root, and read the child with `session.agentTranscript(subagentRunId)`. `host.open` throws for the thread id of a `harnessAgent` child. +- A child thread does not open. Open its root, and read the child with `session.agentTranscript({ subagentRunId })`. `host.open` throws for the thread id of a `harnessAgent` child. - `host.observe` also reads a child thread by its own thread id. It finds the root through the session index, so it needs `stores.sessions`. `memoryPersistence()` has it. ## Write to several threads in one append @@ -48,10 +48,12 @@ const root = await host.open(assistant, { threadId: 'agent-1' }) `session.append` adds records to the root thread. A record with `thread: ''` belongs to that thread of the same log. The whole append is all or nothing: ```ts group=harness-shared-logs -await root.append([ - { type: 'app.child_created', thread: 'task-1', task: 'research' }, - { type: 'app.child_linked', child: 'task-1' }, -]) +await root.append({ + records: [ + { type: 'app.child_created', thread: 'task-1', task: 'research' }, + { type: 'app.child_linked', child: 'task-1' }, + ], +}) ``` Here the child record and the link in the root land together, or not at all. @@ -85,9 +87,9 @@ const tracker = defineHarness({ durability: { reduce }, }) const tracked = await host.open(tracker, { threadId: 'agent-2' }) -await tracked.append([ - { type: 'app.child_created', thread: 'task-2', task: 'research' }, -]) +await tracked.append({ + records: [{ type: 'app.child_created', thread: 'task-2', task: 'research' }], +}) console.log((await tracked.logState())?.tasks) ``` diff --git a/docs/harness/shared-threads.md b/docs/harness/shared-threads.md index 2fdaed0a9a..5280e8f63b 100644 --- a/docs/harness/shared-threads.md +++ b/docs/harness/shared-threads.md @@ -98,9 +98,9 @@ export async function invite(userId: string, role: 'member' | 'viewer') { When `unshare` removes a user: -- The queued and running work that the user started on this host stops. The log keeps who cancelled it. -- The streams of the user on the thread close. A socket gets a `forbidden` error frame. -- The open handles of the user throw `AccessDeniedError`. +- The queued and running work that the user started stops, on every host. The log keeps who cancelled it. The schedules of the user on the thread are removed. +- The streams of the user on the thread close. A socket gets a `forbidden` error frame. A running `events()` read of a handle or a view ends with `AccessDeniedError`. +- The open handles and views of the user throw `AccessDeniedError`. - The signed media URLs of the user stop working. ## 3. Connect each person diff --git a/docs/harness/skills.md b/docs/harness/skills.md index 89a399d56e..7fb8b7254a 100644 --- a/docs/harness/skills.md +++ b/docs/harness/skills.md @@ -73,7 +73,7 @@ export const cliSkills = skills({ ## Skills that change while the session runs -The plugin watches the folders, and it lists them again at the start of each turn: +The plugin watches the folders, and the commands follow them: - A new `SKILL.md` adds its command. - A removed skill takes its command away. @@ -82,6 +82,10 @@ The plugin watches the folders, and it lists them again at the start of each tur Each change sends a `harness.commands.changed` event, so a UI can show the new command list at once. +The plugin watches each folder once for each host. All sessions of the harness on that host share the watchers, so many open sessions still use one watcher for each folder. A change only tells the sessions to update. The plugin reads the folders again at the start of each turn. + +If the plugin cannot watch a folder, it logs one warning with the folder and the reason. An example is a Linux system with no free file watchers. Changes in that folder still show at the next turn. + ## Many skills The model's list has one line for each skill: the name and the first sentence of the description. When the list is longer than `maxCatalogTokens` (default 4000), the skills of the earlier folders stay, and the rest leave the model's list. A skill that the user calls with its command always stays in it. diff --git a/docs/harness/snapshots.md b/docs/harness/snapshots.md index 985bcdfc5d..bfa89ee4d4 100644 --- a/docs/harness/snapshots.md +++ b/docs/harness/snapshots.md @@ -45,6 +45,7 @@ const session = await host.open(coder, { threadId: 'fix-login' }) - `dataDir` is a folder for the snapshot repositories, an absolute path. Put it outside `root`. Each thread of each tenant gets its own repository in it. - The plugin needs the `git` command. `/undo` and `/redo` need git 2.26 or later. - Your project does not need to be a git repository, and the plugin never uses the `.git` folder of your project. +- With [`workspaceTools({ workspace: { bind } })`](./coding-workspaces) before it, each thread keeps its snapshots in its own workspace. Leave out `root`. ## 2. Undo and redo a turn @@ -99,13 +100,13 @@ for (const { record } of await log.read({ logId: logIdOf({ logId: 'fix-login' }) ```ts group=harness-snapshots const messages = await session.transcript() const firstAnswer = messages.find((message) => message.role === 'assistant') -if (firstAnswer?.id) await session.revert(firstAnswer.id) +if (firstAnswer?.id) await session.revert({ messageId: firstAnswer.id }) // Bring back the files and the messages. await session.unrevert() ``` -- A revert uses the last 500 steps of the thread. +- A revert uses the last 500 steps of the thread. The log snapshot keeps them, so a revert also works with [`logRecords: 'since-snapshot'`](./lifecycle#drop-the-records-before-a-snapshot). - A revert touches only the files that the tool calls after that message changed. Each file goes back to its state before the first of these calls. A file that did not exist then is removed. - Other files stay as they are. Your own edits between the hidden turns stay too. - `unrevert()` touches only the same files. It puts them back as they were when the revert started, with your edits. @@ -114,7 +115,7 @@ await session.unrevert() - The next prompt drops the hidden messages for good. Then `unrevert()` does not change the files. - The revert is saved, so `unrevert()` still works after a restart. -In your own UI, call `client.revert(messageId)` and `client.unrevert()`. The client needs `'undo'` in `expose.commands`. More about reverts: [Go back to an earlier message](./sessions#go-back-to-an-earlier-message). +In your own UI, call `client.revert({ messageId })` and `client.unrevert()`. The client needs `'undo'` in `expose.commands`. More about reverts: [Go back to an earlier message](./sessions#go-back-to-an-earlier-message). ## In a sandbox @@ -145,7 +146,12 @@ The undo point and the steps are in the session log. So when another server take ## Remove the snapshots of a thread -When you delete a thread, remove its snapshot repository too. Call `purgeSnapshots` with the options of the plugin and the thread: +[`host.sessions.delete`](./lifecycle#delete-a-thread-for-good) removes the snapshot repository of the thread for you. This needs two things: + +- The thread has an index entry. +- This host knows the harness: it opened a session of it before, or it has the harness in its [`harnesses`](./lifecycle#delete-a-thread-for-good) option. + +In other cases, call `purgeSnapshots` with the options of the plugin and the thread: ```ts group=harness-snapshots import { purgeSnapshots } from '@tanstack/ai-harness/plugins/coding' diff --git a/docs/harness/subagents.md b/docs/harness/subagents.md index 0a5c57648c..ec82d23add 100644 --- a/docs/harness/subagents.md +++ b/docs/harness/subagents.md @@ -249,6 +249,12 @@ For a follow-up, `receipt.operationId` is the id of the new run. While the run b A steer joins only a run of the same sender. A steer from another sender runs as a follow-up for that sender. As a result, no message runs with the credentials of another person. +Who sends: + +- `run.send` and `run.cancel` act as the user of the run, not as the user who opened the session. +- To send as another user, pass `principal`: `run.send('Add a title.', { principal })`. +- `session.agents.drafter.run(input, { principal })` and `start` run the agent as that user. + An agent written as plain code, with no `ctx.chat` call, makes no model call. It gets a steer only as a follow-up. To message a run from the browser, see [Send a message to a running agent](./connect#send-a-message-to-a-running-agent). @@ -411,7 +417,7 @@ Known limits: ## Show what a child wrote after a reload -Your UI shows each child's text and tool calls while they stream. Then the user reloads the page. The thread's transcript has only the agent tool call and its result, because each child keeps its own messages. `session.agentTranscript(subagentRunId)` gives you those messages back. +Your UI shows each child's text and tool calls while they stream. Then the user reloads the page. The thread's transcript has only the agent tool call and its result, because each child keeps its own messages. `session.agentTranscript({ subagentRunId })` gives you those messages back. Each agent tool result names its child with `subagentRunId`. Read the result, then ask for the child: @@ -444,7 +450,7 @@ for (const message of await session.transcript()) { if (!('subagentRunId' in result)) continue if (typeof result.subagentRunId !== 'string') continue - const child = await session.agentTranscript(result.subagentRunId) + const child = await session.agentTranscript({ subagentRunId: result.subagentRunId }) console.log(child.map((item) => item.content)) } ``` @@ -469,7 +475,7 @@ const client = createHarnessClient({ }) export async function childMessages(subagentRunId: string) { - return client.agentTranscript(subagentRunId) + return client.agentTranscript({ subagentRunId }) } ``` @@ -537,7 +543,35 @@ export const studio = defineHarness({ The router can also return a pick of its own: `'main'`, a name, a list of names, `{ names, order }`, or `{ steps }`. Anywhere a name goes, `{ name, input }` can go too, for example `{ name: 'pricer', input: { vendor: 'Acme' } }`. If a picked agent has an `inputSchema` and no input, the turn fails with `Agent "pricer" needs input. Return { name: 'pricer', input } from the router.` See [Route, or let the model pick](../chat/subagents#route-or-let-the-model-pick) for each pick. -The router also gets `session`, the `input` of the turn as the user sent it, `operationId`, `inputId`, `abortSignal`, and `activeAgent`. Next to `router`, `routing` takes `order`, `limits`, `sandbox`, `after`, and `humans`. The first three work the same as in `subagents`. +The router also gets `session`, the `input` of the turn as the user sent it, `operationId`, `inputId`, `abortSignal`, `activeAgent`, and `model`. Next to `router`, `routing` takes `order`, `limits`, `sandbox`, `after`, and `humans`. The first three work the same as in `subagents`. + +### Count the model calls of the router + +Your router asks a model which agent fits. A plain `chat()` call does not count in `session.usage()`, and it keeps running after the user cancels the turn. Call the model with `model.chat` instead: + +```ts group=harness-routing +export const helpdesk = defineHarness({ + name: 'acme/helpdesk', + adapter: openaiText('gpt-6.1-sol'), + agents: [writer, pricer], + routing: { + router: async ({ model, messages }) => { + const { text } = await model.chat({ + label: 'pick', + options: { + messages, + systemPrompts: ['Answer "writer" for a blog post, else "main". Answer one word.'], + }, + }) + return text.trim() === 'writer' ? 'writer' : 'main' + }, + }, +}) +``` + +- `model.chat` takes the same options as [`ctx.model.chat`](./plugins#call-a-model-from-a-plugin) of a plugin. Without `adapter`, it uses the thread setting `model`, else the harness `adapter`. +- The call uses the keys of the sender of the turn. A cancel of the turn stops it. +- Its usage counts in the turn, and in the session under `router`. See [Usage of the router and transfers](./usage#usage-of-the-router-and-transfers). What the turn does with the pick: diff --git a/docs/harness/thread-settings.md b/docs/harness/thread-settings.md index c8dafd4864..aa05330a2a 100644 --- a/docs/harness/thread-settings.md +++ b/docs/harness/thread-settings.md @@ -51,7 +51,7 @@ console.log(receipt.status) // 'accepted' console.log(session.settings()) // { model: 'strong', reasoning: 'high', instructions: 'Answer in French.' } ``` -The settings apply from the next turn. A field that you leave out stays as it is. Set a field to `null` to clear it: +The settings apply from the next turn. A field that you leave out stays as it is. `configure` merges your fields into the latest stored settings, also when another host changed them. Set a field to `null` to clear it: ```ts group=harness-thread-settings await session.configure({ settings: { instructions: null } }) @@ -70,7 +70,7 @@ The six fields: An unknown field, model, or plugin, or a bad value, gets a receipt with `status: 'rejected'` and the reason. `stores.metadata` keeps the settings, so they stay after a restart. Without a metadata store, they stay in memory. A [fork](./fork-and-reset) copies them. -`configure` is an input, so the log keeps who changed what. Each change sends a `harness.settings.changed` event with the new `settings`. +`configure` is an input, so the log keeps who changed what. Each change sends a `harness.settings.changed` event with the new `settings`. A session on another host reads the change too, so its next turn uses the new settings. A client can change only the fields in `expose.settings` of the harness. See [Choose what clients can change](./connect#choose-what-clients-can-change). @@ -101,8 +101,9 @@ For `reasoning`, the `overrides` of the prompt win, then the stored setting, the - A `tools` list keeps only the tools with those names. `{ remove }` drops the tools with those names. Either way, the tools of the prompt `overrides` stay. - The harness applies `tools` before the plugins prepare the tool list (`prepareTools`). So a plugin that wraps tools sees only the tools that a thread keeps. -- `plugins: { remove }` leaves the tools, prompts, middleware, and model pick of those plugins out of the turn. Plugins that run once per session stay mounted. -- The agents and subagents that a removed plugin gives stay available. +- `plugins: { remove }` leaves everything that those plugins give out of the thread: tools, prompts, middleware, the model pick, agents, subagents, commands, and the items that they add for other plugins. For example, a thread without `agents()` gets no subagent tools, no `/agent` command, and no permission rules of its agent profiles. The thread also cannot run an agent of a removed plugin by its name. +- Plugins that run once per session stay mounted, but they give the thread nothing. +- If a plugin that stays (of any lifetime) or a harness `middleware` entry needs a capability of a removed plugin, the setting is rejected. ## The working folder diff --git a/docs/harness/turn-control.md b/docs/harness/turn-control.md index 68b09b00e1..c486d3cedb 100644 --- a/docs/harness/turn-control.md +++ b/docs/harness/turn-control.md @@ -191,6 +191,10 @@ The join stops at the first message that: That message and the ones after it wait. `canJoin` runs again before each model call and at each final answer, so it can run more than once for one message. Keep it a quick check with no side effects. What still waits when the turn ends runs as its own turn, before other queued turns. A cancelled one settles `aborted` and does not run. +A `session.tool()` turn takes no messages. It runs no model, so nothing could answer them. A message that arrives while it runs waits, and then runs as its own turn. See [Run a tool without a model](./tool-runs). + +A turn that `routing` sends to root agents takes messages, the same as the main model. The root agents get them at each model call. + A joined message runs with the credentials of the sender of the turn. So without `turn.canJoin`, a waiting message joins the running turn only when the same person sent both. A message from another person waits and runs as its own turn after the running one, as its sender. Joins are a prefix, so it also holds back the messages that wait after it. Set `canJoin` to choose your own rule. Compare `principal` with `turnPrincipal` to keep one person's message out of another person's turn. See [Keep other people out of a running turn](./shared-threads#5-keep-other-people-out-of-a-running-turn). `onJoin` runs before the model call that the joined messages reach. It can return messages, which come after the joined ones, and records, which land in the same append as the join. diff --git a/docs/harness/usage.md b/docs/harness/usage.md index 0bc97a9363..7646cfb45f 100644 --- a/docs/harness/usage.md +++ b/docs/harness/usage.md @@ -2,7 +2,7 @@ title: Track usage and cost id: harness-usage order: 7 -description: "Read the tokens and the cost of a thread, in total, by model, by sender, and by run. Show a live total, stop a person who is over budget, and add /usage with your own prices." +description: "Read the tokens and the cost of a thread, in total, by model, by sender, by run, and by plugin. Show a live total, stop a person who is over budget, and add /usage with your own prices." keywords: - tanstack ai - harness @@ -39,11 +39,12 @@ console.log(total.calls, total.totalTokens, total.cost) console.log(byModel['openai/gpt-6.1-sol']?.totalTokens, bySender.ada?.cost) ``` -The totals have three groups. `snapshot().usage` has the same data: +The totals have four groups. `snapshot().usage` has the same data: - `total`: every model call of the thread. - `byModel`: the calls of each model, keyed by `provider/model`, for example `openai/gpt-6.1-sol`. - `bySender`: the calls of each person, keyed by the principal id. It has only the calls of a known sender, so its sum can be less than `total`. +- `bySource`: the model calls that are not a turn or an agent run, keyed by their source: `plugin:`, `router`, or `transfer`. See [Usage of plugins](#usage-of-plugins). Each entry has `calls`, `promptTokens`, `completionTokens`, `totalTokens`, `cachedTokens`, and `cacheWriteTokens`. If a provider reported a cost, the entry also has `cost`. The harness adds the costs that the providers report, in their currency. It uses no price list of its own. @@ -55,6 +56,9 @@ At the end of each call, the provider reports its usage. These calls count: - The subagents that a turn starts, for the same sender. - Background agents, for the person who started the agent. - A harness that runs as a child with [`harnessAgent`](./subagents#call-a-harness-as-a-child), for the same person. +- The model calls of plugins with `ctx.model.chat`, for the user of their run. +- The model calls of the router with `model.chat`, for the sender of the turn. +- The summary of a transfer with `messages: 'summary'`, for the person who made the transfer, in the source thread. A call that fails before it ends does not count. Each message runs as its sender, so in a [shared thread](./shared-threads), `bySender` shows what each person used. @@ -94,6 +98,43 @@ export async function tokensByRun(log: LogStore, threadId: string) { - When an agent run ends, its run record gets `usage`: the tokens of the run and of all its children. It is in the `harness.run.update` record of the end. - To stop a run that uses too much, give it a `budget`. See [Give one run a budget and a time limit](./subagents#give-one-run-a-budget-and-a-time-limit). +## Usage of plugins + +A goal judge or a title writer calls a model next to the turn, and you want to see what these calls cost. A plugin that calls its model with [`ctx.model.chat`](./plugins#call-a-model-from-a-plugin) counts each call under its own name: + +```ts group=harness-usage +import { title } from '@tanstack/ai-harness/plugins' + +export const titled = defineHarness({ + name: 'acme/titled', + adapter: openaiText('gpt-6.1-sol'), + plugins: () => [title({ adapter: openaiText('gpt-6-luna') })], +}) + +const titledSession = await host.open(titled, { threadId: 'titled-thread' }) +await titledSession.send('Plan a trip to Rome.') + +const { bySource } = titledSession.usage() +console.log(bySource['plugin:tanstack/title']?.totalTokens) +``` + +- The title call runs next to the turn, so its usage can arrive a moment after the turn ends. +- The calls of a plugin also count in `total`, `byModel`, and `bySender`, and in the use of their run. +- The calls of turns and agent runs have no source, so they are not in `bySource`. +- Each `harness.usage` record of a plugin call has `source`, and the `label` of the call when it has one. The `goal()` judge has the label `judge`, `title()` has `title`, and `/compact` has `compact`. + +## Usage of the router and transfers + +Two more kinds of calls run outside a turn model call. Each one has its own source in `bySource`: + +- `router`: the calls of [`model.chat` in the router](./subagents#count-the-model-calls-of-the-router). They also count in the use of the turn. +- `transfer`: the summary call of a [transfer](./handoffs) with `messages: 'summary'`. Its record has the label `summary`. It counts in the source thread, also when the transfer fails later. + +```ts group=harness-usage +const { router, transfer } = session.usage().bySource +console.log(router?.totalTokens, transfer?.totalTokens) +``` + ## Show a live total Want a token counter that ticks up while the chat runs? `session.onUsage` calls you after each model call, with the new thread totals: @@ -107,7 +148,7 @@ const stopCounting = session.onUsage(({ total }) => { stopCounting() ``` -You also get the `model`, the `sender` when it is known, and the `usage` of that one call. Reading raw events? Each call is a `harness.usage` event (`HARNESS_EVENTS.usage`). +You also get the `model`, the `sender` when it is known, and the `usage` of that one call. A call of a plugin or of the router also has its `source`, and its `label` when it has one. Reading raw events? Each call is a `harness.usage` event (`HARNESS_EVENTS.usage`). ## Stop a person who is over budget @@ -193,7 +234,7 @@ Limits of the plugin: ## What you have now -- The tokens and the cost of a thread, in total, by model, and by sender. +- The tokens and the cost of a thread, in total, by model, by sender, and by plugin. - The tokens of each run, with the use of its children in the run above them. - Totals that stay after a restart. - A live total from the `harness.usage` event. diff --git a/docs/persistence/store-reference.md b/docs/persistence/store-reference.md index eb78ec7bfe..e017a2fa66 100644 --- a/docs/persistence/store-reference.md +++ b/docs/persistence/store-reference.md @@ -672,7 +672,7 @@ interface LogStore { ownerId: string ttlMs: number affinity?: string - }): Promise<{ epoch: number } | null> + }): Promise<{ epoch: number; expired?: boolean } | null> renew(claim: { logId: string ownerId: string @@ -803,6 +803,9 @@ or a `wakeInMs` less than 0, rejects with `RangeError`. - For the owner of a live record, it extends the expiry and returns the same epoch. - Else it adds 1 to the epoch, stores `harness`, clears the wake time, and returns `{ epoch }`. +- When the record had another owner that did not release, it returns + `{ epoch, expired: true }`. That owner stopped. The harness then sends + `harness.ownership.changed` with `reason: taken_over`. For the other ownership methods: diff --git a/docs/tools/server-tools.md b/docs/tools/server-tools.md index 9106cb0e00..e3d290d0c3 100644 --- a/docs/tools/server-tools.md +++ b/docs/tools/server-tools.md @@ -359,7 +359,7 @@ Malformed JSON arguments and input validation failures (Standard Schema or JSON The `inputSchema` checks the input just before the tool runs: -- The check runs after approval edits and after middleware `onBeforeToolCall`. An approval edit that does not fit the schema stops the resume with an `invalid-edited-args` error. A middleware change that does not fit becomes a tool error. +- The check runs after approval edits and after middleware `onBeforeToolCall`. An approval edit or a middleware change that does not fit the schema becomes a tool error. The resume goes on, the tool does not run, and the model gets the input error as the result of that call. - An async validator is awaited. - The tool gets the schema's output, so a transform such as `z.string().transform(Number)` applies once. diff --git a/examples/harness-cli/src/tui.tsx b/examples/harness-cli/src/tui.tsx index 53b8b89694..3e942440cd 100644 --- a/examples/harness-cli/src/tui.tsx +++ b/examples/harness-cli/src/tui.tsx @@ -104,7 +104,9 @@ function Screen({ view }: { view: SessionView }) { .answer(answerOf(text, question.schema)) .then((receipt) => { if (receipt.status === 'rejected') { - view.notice(`Answer again: ${receipt.reason ?? 'not accepted'}`) + view.notice({ + text: `Answer again: ${receipt.reason ?? 'not accepted'}`, + }) } }) } @@ -168,7 +170,35 @@ function Screen({ view }: { view: SessionView }) { ) } +/** The ids of the child agent parts in `parts`, at any depth. */ +function agentIds(parts: ReadonlyArray): Array { + return parts.flatMap((part) => + part.type === 'agent' ? [part.id, ...agentIds(part.parts)] : [], + ) +} + +/** + * Fill the saved child agents of a reopened thread. Ink prints the old turns + * once, so they load before the first render. A loaded child can have saved + * children of its own, so it loads again until no new id shows up. + */ +async function loadSavedAgents(view: SessionView) { + const loaded = new Set() + for (;;) { + const ids = view.store + .get() + .messages.flatMap((message) => + message.role === 'assistant' ? agentIds(message.parts) : [], + ) + .filter((id) => !loaded.has(id)) + if (ids.length === 0) return + for (const id of ids) loaded.add(id) + await Promise.all(ids.map((id) => view.loadAgent({ id }))) + } +} + /** The `ui` of runCli: show the screen until the user quits. */ export async function runTui(view: SessionView) { + await loadSavedAgents(view) await render().waitUntilExit() } diff --git a/packages/ai-acp/src/agent/index.ts b/packages/ai-acp/src/agent/index.ts index 75c504f503..8bb30e55d6 100644 --- a/packages/ai-acp/src/agent/index.ts +++ b/packages/ai-acp/src/agent/index.ts @@ -16,7 +16,6 @@ import type { AgentApp, AgentContext, AgentConnection, - ClientCapabilities, ElicitationSchema, PermissionOption, SessionUpdate, @@ -230,6 +229,19 @@ function parseArgs(raw: string | undefined): unknown { } } +/** + * The tool call of an interrupt: its own `toolCallId`, or for an approval of + * `permissions()`, the `toolCallId` in its payload. Core gives such an + * approval no `toolCallId` of its own. + */ +function toolCallIdOf(interrupt: Interrupt) { + if (interrupt.toolCallId) return interrupt.toolCallId + const payload = interrupt.metadata?.[INTERRUPT_PAYLOAD_METADATA_KEY] + return isRecord(payload) && typeof payload.toolCallId === 'string' + ? payload.toolCallId + : undefined +} + /** The sign-in page of a sign-in interrupt, when it has one. */ function signInUrl(interrupt: Interrupt) { const payload = interrupt.metadata?.[INTERRUPT_PAYLOAD_METADATA_KEY] @@ -338,10 +350,17 @@ type WaitAnswer = * checks each action, and a session that the principal may not read is not * found. Without `principal`, a host with access throws at the start. * - * The agent holds one handle per session. `session/close` releases it, and - * the close of the connection releases them all. The host then closes each - * session when it is idle. `session/resume` only checks the session index: - * the session opens at its first prompt. + * The agent holds a handle on a session only while the session has work: a + * new session, or prompts that run or wait in its queue. When the queue is + * empty, the agent releases the handle, and the host closes the session when + * it is idle. The next prompt opens it again. `session/close` and the close + * of the connection release the handles too. `session/resume` only checks the + * session index: the session opens at its first prompt. + * + * One app serves one connection: it keeps the session handles and prompt + * queues of that connection, and a second connection to it throws. For more + * connections, make one app for each, for example + * `new AcpServer({ createAgent: () => createAcpAgent(options) })`. * * ACP v2 is a draft, so this API is experimental too. */ @@ -355,13 +374,23 @@ export function createAcpAgent

( ) } /** - * The handle of each session that this agent opened, by session id. The - * agent holds one per session. `session/close` releases it, and so does - * the close of the connection. + * The handle of each session that this agent holds, by session id. The + * agent holds one while the session has work, and releases it when the + * prompt queue of the session is empty. `session/close` releases it too, + * and so does the close of the connection. */ const handles = new Map>() - /** The last pump of each session. The pumps of one session run in order. */ + /** + * The last prompt of each session. The prompts of one session run in the + * order that they came in. + */ const tails = new Map>() + /** + * The sessions that this agent opened in this connection. They are known + * without the session index, so a host with no `stores.sessions` still + * takes their prompts after the agent released their handles. + */ + const opened = new Set() /** The handle of a session: the one that the agent holds, or a new one. */ const live = (sessionId: string) => { @@ -370,9 +399,12 @@ export function createAcpAgent

( const opening = host.open(harness, { threadId: sessionId, principal }) handles.set(sessionId, opening) // A failed open holds nothing, so the next request opens again. - opening.catch(() => { - if (handles.get(sessionId) === opening) handles.delete(sessionId) - }) + opening.then( + () => opened.add(sessionId), + () => { + if (handles.get(sessionId) === opening) handles.delete(sessionId) + }, + ) return opening } @@ -388,13 +420,23 @@ export function createAcpAgent

( } /** - * Throw unless this agent or the session index knows the id. A subagent + * Release the handle of a session when no prompt of it runs or waits. An + * idle session holds nothing, and the next prompt opens it again. + */ + const releaseIfIdle = (sessionId: string) => { + if (tails.has(sessionId)) return + void release(sessionId) + } + + /** + * Throw unless this agent or the session index knows the id. The agent + * knows the sessions that it opened in this connection. A subagent * child is not a session of its own: ACP reaches it through its parent. A * session that the principal may not read gets the same answer as an * unknown id. It reads the index only: no session opens. */ const checkKnown = async (sessionId: string) => { - if (handles.has(sessionId)) return + if (handles.has(sessionId) || opened.has(sessionId)) return try { const entry = await host.sessions.get({ threadId: sessionId, principal }) const isOwnSession = @@ -427,12 +469,17 @@ export function createAcpAgent

( update: SessionUpdate, ) => client.notify('session/update', { sessionId, update }).catch(() => {}) - /** What the editor can show, from its `initialize`. */ - let capabilities: ClientCapabilities = {} - - /** True when the editor can answer `wait` (see `askEditor`). */ - const canAsk = (wait: Wait, url: string | undefined) => { - const elicitation = capabilities.elicitation + /** + * True when the editor can answer `wait` (see `askEditor`). What the editor + * can show comes from the `initialize` of its own connection. + */ + const canAsk = async ( + client: AgentContext, + wait: Wait, + url: string | undefined, + ) => { + const { request } = await client.initialized + const elicitation = request.capabilities?.elicitation switch (wait.kind) { case 'approval': return true @@ -468,7 +515,7 @@ export function createAcpAgent

( toolCallId?: string url?: string }): Promise => { - if (!canAsk(wait, url)) return { status: 'cancelled' } + if (!(await canAsk(client, wait, url))) return { status: 'cancelled' } if (wait.kind === 'approval') { const { outcome } = await client.request('session/request_permission', { sessionId, @@ -530,7 +577,10 @@ export function createAcpAgent

( wait: waitOfInterrupt({ interrupt, runId: sessionId }), url: signInUrl(interrupt), })) - if (items.some(({ wait, url }) => canAsk(wait, url))) { + const askable = await Promise.all( + items.map(({ wait, url }) => canAsk(client, wait, url)), + ) + if (askable.includes(true)) { await notify(client, sessionId, { sessionUpdate: 'state_update', state: 'requires_action', @@ -538,11 +588,12 @@ export function createAcpAgent

( } const resume: Array = [] for (const { interrupt, wait, url } of items) { + const toolCallId = toolCallIdOf(interrupt) const reply = await askEditor({ client, sessionId, wait, - ...(interrupt.toolCallId ? { toolCallId: interrupt.toolCallId } : {}), + ...(toolCallId ? { toolCallId } : {}), ...(url ? { url } : {}), }) if (reply.status === 'stop') return { cancelled: true } @@ -727,16 +778,30 @@ export function createAcpAgent

( ) } + /** + * True after the first connection. The SDK gives a request handler no + * public id of its connection, so the state above is for one connection + * only. The SDK way to serve many is one app for each connection. + */ + let connected = false + return ( agent() .onConnect((connection) => { + // A second connection would share the handles and queues of the + // first, so the SDK closes it with this error. + if (connected) { + throw new Error( + 'An ACP agent app serves one connection. Call createAcpAgent once for each connection.', + ) + } + connected = true // The handles of the sessions go with the connection. void connection.closed.then(() => Promise.all([...handles.keys()].map(release)), ) }) - .onRequest('initialize', ({ params }) => { - capabilities = params.capabilities ?? {} + .onRequest('initialize', () => { return { protocolVersion: PROTOCOL_VERSION, info: { name: harness.name, version: options.version ?? '0.0.0' }, @@ -749,7 +814,10 @@ export function createAcpAgent

( }) .onRequest('session/new', async () => { const sessionId = `acp-${Date.now().toString(36)}-${Math.random().toString(36).slice(2, 9)}` + // The open saves the session in the session index. The first prompt + // opens it again. await live(sessionId) + releaseIfIdle(sessionId) return { sessionId } }) .onRequest('session/list', async ({ params }) => { @@ -781,36 +849,53 @@ export function createAcpAgent

( }) .onRequest('session/prompt', async ({ params, client }) => { const { sessionId } = params - const session = await known(sessionId) - const input = await promptInput(session, params.prompt).catch( - (error: unknown) => { - // The editor sent a file the harness refuses, so the prompt is - // invalid: answer with the reason, not an internal error. - if (error instanceof MediaError) { - throw RequestError.invalidParams( - { status: error.status }, - error.message, - ) - } - throw error - }, - ) const messageId = `user-${crypto.randomUUID()}` - await notify(client, sessionId, { - sessionUpdate: 'user_message', - messageId, - content: params.prompt, - }) - // `pump` catches its own errors and always reports idle. The prompts of - // one session go to the harness in order, each after the turn before it - // ended, so a later prompt cannot overtake a permission request. - const tail = (tails.get(sessionId) ?? Promise.resolve()).then(() => - pump(client, sessionId, session, input), - ) + /** + * Open the session, store the files of the prompt, and echo the user + * message. A file upload can take long, so this runs at once, also + * while an earlier prompt still runs. + */ + const accepted = (async () => { + const session = await known(sessionId) + const input = await promptInput(session, params.prompt).catch( + (error: unknown) => { + // The editor sent a file the harness refuses, so the prompt is + // invalid: answer with the reason, not an internal error. + if (error instanceof MediaError) { + throw RequestError.invalidParams( + { status: error.status }, + error.message, + ) + } + throw error + }, + ) + await notify(client, sessionId, { + sessionUpdate: 'user_message', + messageId, + content: params.prompt, + }) + return { session, input } + })() + // Take the place in the queue before the first await, so a prompt + // keeps its place while its files upload. `pump` catches its own + // errors and always reports idle. The prompts of one session go to the + // harness in order, each after the turn before it ended, so a later + // prompt cannot overtake a permission request. A prompt that fails + // here runs no turn: the request answers with the error. + const tail = (tails.get(sessionId) ?? Promise.resolve()) + .then(() => accepted) + .then( + ({ session, input }) => pump(client, sessionId, session, input), + () => {}, + ) tails.set(sessionId, tail) void tail.then(() => { - if (tails.get(sessionId) === tail) tails.delete(sessionId) + if (tails.get(sessionId) !== tail) return + tails.delete(sessionId) + releaseIfIdle(sessionId) }) + await accepted return { messageId } }) // Only the handle of this agent goes. Other handles keep the session. diff --git a/packages/ai-acp/tests/agent-handlers.test.ts b/packages/ai-acp/tests/agent-handlers.test.ts index 2275da7594..1198e52af8 100644 --- a/packages/ai-acp/tests/agent-handlers.test.ts +++ b/packages/ai-acp/tests/agent-handlers.test.ts @@ -338,32 +338,133 @@ describe('ACP agent handlers', () => { close = undefined }) - it('releases its handle at session/close, and every handle when the connection closes', async () => { + it('holds no handle on an idle session, and opens it again at the next prompt', async () => { // A session that no handle holds closes at once when it is idle. - const { harness, host, editor } = setup({ lifecycle: { idleCloseMs: 0 } }) + const { harness, host, editor, running, states, textOf } = setup({ + lifecycle: { idleCloseMs: 0 }, + }) close = () => host.close() const isOpen = (threadId: string) => host.isOpen(harness, { threadId }) - let kept = '' await editor.connectWith(createAcpAgent({ host, harness }), async (ctx) => { await ctx.request('initialize', { protocolVersion: 2, info: { name: 'editor', version: '1.0.0' }, }) - const closed = await ctx.request('session/new', { + const { sessionId } = await ctx.request('session/new', { cwd: '/tmp', mcpServers: [], }) - kept = (await ctx.request('session/new', { cwd: '/tmp', mcpServers: [] })) - .sessionId - expect(isOpen(closed.sessionId)).toBe(true) + await vi.waitFor(() => expect(isOpen(sessionId)).toBe(false)) - await ctx.request('session/close', { sessionId: closed.sessionId }) + await ctx.request('session/prompt', { + sessionId, + prompt: [{ type: 'text', text: 'wait' }], + }) + await vi.waitFor(() => expect(running(sessionId)).toBe(true)) + // A running prompt holds the session. + expect(isOpen(sessionId)).toBe(true) + await ctx.notify('session/cancel', { sessionId }) + await vi.waitFor(() => expect(isOpen(sessionId)).toBe(false)) - expect(isOpen(closed.sessionId)).toBe(false) - expect(isOpen(kept)).toBe(true) + await ctx.request('session/prompt', { + sessionId, + prompt: [{ type: 'text', text: 'hello' }], + }) + await vi.waitFor(() => + expect(states(sessionId)).toContain('idle:end_turn'), + ) + // The session opened again and ran the prompt. + expect(textOf(sessionId)).toBe('echo: hello') + await vi.waitFor(() => expect(isOpen(sessionId)).toBe(false)) + }) + }) + + it('takes the prompts of its new sessions on a host with no session index', async () => { + const { harness, editor, states, textOf } = setup() + // Only a log store: `host.sessions.get` finds no session. + const host = createHarnessHost({ + persistence: { stores: { log: memoryPersistence().stores.log } }, + lifecycle: { idleCloseMs: 0 }, + }) + close = () => host.close() + await editor.connectWith(createAcpAgent({ host, harness }), async (ctx) => { + await ctx.request('initialize', { + protocolVersion: 2, + info: { name: 'editor', version: '1.0.0' }, + }) + const { sessionId } = await ctx.request('session/new', { + cwd: '/tmp', + mcpServers: [], + }) + await vi.waitFor(() => + expect(host.isOpen(harness, { threadId: sessionId })).toBe(false), + ) + + await ctx.request('session/prompt', { + sessionId, + prompt: [{ type: 'text', text: 'hello' }], + }) + await vi.waitFor(() => + expect(states(sessionId)).toContain('idle:end_turn'), + ) + expect(textOf(sessionId)).toBe('echo: hello') + }) + }) + + it('keeps the order of two prompts when the first one uploads a file', async () => { + const { harness, host, persistence, editor, updates, states, textOf } = + setup() + close = () => host.close() + // The store of the first file answers only when the test lets it. + let upload = () => {} + const uploaded = new Promise((resolve) => { + upload = resolve + }) + const put = persistence.stores.blobs.put.bind(persistence.stores.blobs) + const slowPut = vi + .spyOn(persistence.stores.blobs, 'put') + .mockImplementation(async (...args) => { + await uploaded + return put(...args) + }) + await editor.connectWith(createAcpAgent({ host, harness }), async (ctx) => { + await ctx.request('initialize', { + protocolVersion: 2, + info: { name: 'editor', version: '1.0.0' }, + }) + const { sessionId } = await ctx.request('session/new', { + cwd: '/tmp', + mcpServers: [], + }) + const first = ctx.request('session/prompt', { + sessionId, + prompt: [ + { type: 'text', text: 'first' }, + { type: 'image', data: 'AQIDBA==', mimeType: 'image/png' }, + ], + }) + await vi.waitFor(() => expect(slowPut).toHaveBeenCalled()) + const second = ctx.request('session/prompt', { + sessionId, + prompt: [{ type: 'text', text: 'second' }], + }) + // The second prompt has no file, so the agent takes it at once. + await vi.waitFor(() => + expect( + updates.filter( + (entry) => entry.update.sessionUpdate === 'user_message', + ), + ).toHaveLength(1), + ) + upload() + await Promise.all([first, second]) + await vi.waitFor(() => + expect( + states(sessionId).filter((state) => state === 'idle:end_turn'), + ).toHaveLength(2), + ) + expect(textOf(sessionId)).toBe('echo: firstecho: second') }) - // The connection closed when the callback ended. - await vi.waitFor(() => expect(isOpen(kept)).toBe(false)) }) it('keeps the tool from running when the editor rejects it', async () => { @@ -823,7 +924,6 @@ describe('ACP agent handlers', () => { }) await vi.waitFor(() => expect(idle('kept-session')).toBeDefined()) expect(textOf('kept-session')).toBe('echo: hello') - expect(isOpen()).toBe(true) }) }) diff --git a/packages/ai-acp/tests/agent.test.ts b/packages/ai-acp/tests/agent.test.ts index ebb7984f94..89c8d95cc5 100644 --- a/packages/ai-acp/tests/agent.test.ts +++ b/packages/ai-acp/tests/agent.test.ts @@ -249,12 +249,14 @@ describe('ACP v2 agent', () => { const host = createHarnessHost({ persistence: memoryPersistence() }) const updates: Array = [] const offered: Array> = [] + const subjects: Array = [] const editor = client() .onNotification('session/update', ({ params }) => { updates.push(params.update) }) .onRequest('session/request_permission', ({ params }) => { offered.push(params.options.map((option) => option.kind)) + subjects.push(params.subject) return { outcome: { outcome: 'selected', optionId: 'always' } } }) @@ -275,6 +277,10 @@ describe('ACP v2 agent', () => { }) expect(offered).toEqual([['allow_once', 'allow_always', 'reject_once']]) + // The request names the tool call that the editor already shows. + expect(subjects).toEqual([ + { type: 'tool_call', toolCall: { toolCallId: 'call_1' } }, + ]) expect(ran).toEqual(['bash']) expect(agentText(updates)).toContain('Listed.') await host.close() @@ -350,4 +356,80 @@ describe('ACP v2 agent', () => { expect(result?.content).toContain('Ada') await host.close() }) + + it('refuses a second connection to one app, and keeps the first', async () => { + const greet = durableTool( + toolDefinition({ + name: 'greet', + description: 'Ask the user for a name', + inputSchema: { type: 'object', properties: {} }, + }), + async (_args, { step }) => { + const answer = await step.ask({ + name: 'name', + message: 'Name?', + schema: { + type: 'object', + properties: { name: { type: 'string' } }, + required: ['name'], + }, + }) + return { greeted: answer } + }, + ) + const harness = defineHarness({ + name: 'test/acp-connections', + adapter: scripted([toolTurn('greet', '{}'), textTurn('Greeted.')]), + tools: [greet], + }) + const host = createHarnessHost({ persistence: memoryPersistence() }) + const app = createAcpAgent({ host, harness }) + const updates: Array = [] + const asked: Array = [] + const first = client() + .onNotification('session/update', ({ params }) => { + updates.push(params.update) + }) + .onRequest('elicitation/create', ({ params }) => { + asked.push(params) + return { action: 'accept', content: { name: 'Ada' } } + }) + // An editor that shows no elicitation, on its own connection. + const second = client() + + await first.connectWith(app, async (ctx) => { + await ctx.request('initialize', { + protocolVersion: 2, + info: { name: 'editor-a', version: '1.0.0' }, + capabilities: { elicitation: { form: {} } }, + }) + const { sessionId } = await ctx.request('session/new', { + cwd: '/tmp', + mcpServers: [], + }) + // A second editor on the same app would share its sessions, so the + // connection fails at once. + expect(() => + second.connectWith(app, async (other) => { + await other.request('initialize', { + protocolVersion: 2, + info: { name: 'editor-b', version: '1.0.0' }, + }) + }), + ).toThrow( + 'An ACP agent app serves one connection. Call createAcpAgent once for each connection.', + ) + await ctx.request('session/prompt', { + sessionId, + prompt: [{ type: 'text', text: 'greet me' }], + }) + await waitFor(() => expect(isIdle(updates)).toBe(true)) + }) + + // The first editor still gets its question: the second one did not + // change what the first one can show. + expect(asked).toHaveLength(1) + expect(agentText(updates)).toContain('Greeted.') + await host.close() + }) }) diff --git a/packages/ai-code-mode/src/harness.ts b/packages/ai-code-mode/src/harness.ts index 973c711d5f..8e4fc4497c 100644 --- a/packages/ai-code-mode/src/harness.ts +++ b/packages/ai-code-mode/src/harness.ts @@ -7,8 +7,17 @@ import { } from '@tanstack/ai-harness' import { createCodeMode } from './create-code-mode' import type { AnyTool } from '@tanstack/ai' +import type { PermissionMode } from '@tanstack/ai-harness' import type { CodeModeTool, CodeModeToolConfig } from './types' +/** The permission modes of `permissions()`, for the `mode` setting. */ +const MODES: ReadonlyArray = [ + 'default', + 'plan', + 'acceptEdits', + 'bypass', +] + export interface CodeModePluginOptions extends Omit< CodeModeToolConfig, 'tools' @@ -20,6 +29,12 @@ export interface CodeModePluginOptions extends Omit< * the rules of tool plugins decide. A safe tool with * `metadata.codeMode: true`, for example a tool of an `mcp()` server with * `codeMode: true`, moves even when `include` does not pick it. + * + * `include` cannot move a tool with `needsApproval`, or a tool that a rule + * asks for or denies: such a tool stays a normal tool call, with its + * approval and its checks. Code mode warns once for each such tool that `include` + * picks. The rules decide in the permission mode of the session at each + * turn, so in plan mode a tool that edits stays a normal tool call. */ include?: (tool: CodeModeTool) => boolean /** @@ -58,8 +73,10 @@ function isServerTool(tool: AnyTool): tool is CodeModeTool { * tools that are safe to run without a question move into code mode. The * other tools stay normal tool calls. A tool with `metadata.codeMode: true` * moves even when `include` does not pick it, but only when it is safe in - * the same way. Calls inside the isolate also skip the permission checks, - * so a tool that declares `PermissionResources` never moves. With + * the same way. `include` can pick more tools, but never a tool with + * `needsApproval` or a tool that a rule asks for or denies. Calls inside the + * isolate also skip the permission checks, so a tool that declares + * `PermissionResources` never moves. With * `lazy: true`, the model gets only the names of the moved tools and asks * `discover_tools` for the signatures. * @@ -77,6 +94,15 @@ function isServerTool(tool: AnyTool): tool is CodeModeTool { */ export function codeMode(options: CodeModePluginOptions) { const { include, lazy = false, ...config } = options + /** The tools that `include` picked but cannot move, warned once each. */ + const warned = new Set() + const warnKept = (name: string) => { + if (warned.has(name)) return + warned.add(name) + console.warn( + `[@tanstack/ai-code-mode] include picked the tool ${JSON.stringify(name)}, but it needs approval, or a rule asks for it or denies it. A call inside code mode skips those checks, so the tool stays a normal tool call.`, + ) + } return definePlugin({ name: 'tanstack/code-mode', optionalRequires: [PermissionDecisionCapability], @@ -104,10 +130,20 @@ export function codeMode(options: CodeModePluginOptions) { // Ask permissions() when it is mounted: it knows all rules in its // order and its default. Read here, so it can come after code mode. const decide = ctx.getOptional(PermissionDecisionCapability) + const decisionOf = (tool: CodeModeTool, mode: PermissionMode) => + decide?.(tool.name, mode) ?? + decidePermission(rules, tool.name, mode) const safe = (tool: CodeModeTool) => - !tool.needsApproval && - (decide?.(tool.name, 'plan') ?? - decidePermission(rules, tool.name, 'plan')) === 'allow' + !tool.needsApproval && decisionOf(tool, 'plan') === 'allow' + // The permission mode of the session now, read at each turn. In + // plan mode, a tool that edits must stay a tool call, where the + // rules deny it. + const value = ctx.config.get('mode') + const mode = MODES.find((each) => each === value) ?? 'default' + // `include` widens the pick past plan mode, but a call inside the + // isolate must not skip an approval, an `ask` rule, or a deny rule. + const canInclude = (tool: CodeModeTool) => + !tool.needsApproval && decisionOf(tool, mode) === 'allow' const moves = (tool: CodeModeTool) => { // The permission checks run on tool calls of the model. A call // inside the isolate skips them, so such a tool stays a tool call. @@ -118,7 +154,11 @@ export function codeMode(options: CodeModePluginOptions) { // A marked tool still needs the checks of the default pick. const isMarked = tool.metadata?.codeMode === true if (isMarked && safe(tool)) return true - return include ? include(tool) : safe(tool) + if (!include) return safe(tool) + if (!include(tool)) return false + if (canInclude(tool)) return true + warnKept(tool.name) + return false } // Two tools that map to the same identifier: the first one moves. const byIdentifier = new Map() diff --git a/packages/ai-code-mode/tests/harness.test.ts b/packages/ai-code-mode/tests/harness.test.ts index f5556e3513..a8b9ad4df6 100644 --- a/packages/ai-code-mode/tests/harness.test.ts +++ b/packages/ai-code-mode/tests/harness.test.ts @@ -267,9 +267,9 @@ describe('codeMode', () => { defineHarness({ name: 'test/code-mode-include', adapter, - tools: [tools.lookup, tools.remove], + tools: [tools.lookup, tools.notionSearch], plugins: () => [ - codeMode({ driver, include: (tool) => tool.name === 'remove' }), + codeMode({ driver, include: (tool) => tool.name === 'lookup' }), ], }), { @@ -278,7 +278,7 @@ describe('codeMode', () => { ) await picked.send('hi') expect(calls[0].tools.map((tool: { name: string }) => tool.name)).toEqual([ - 'lookup', + 'notion_search', 'execute_typescript', ]) @@ -387,6 +387,65 @@ describe('codeMode', () => { await host.close() }) + it('keeps a tool that needs approval, or that a rule asks for or denies, when include picks it, and warns once', async () => { + const warn = vi.spyOn(console, 'warn').mockImplementation(() => {}) + const denyDocs = definePlugin({ + name: 'test/deny-docs', + setup: () => ({ + contribute: [ + PermissionRules.item({ tool: 'docs_deny', decision: 'deny' }), + PermissionRules.item({ tool: 'docs_ask', decision: 'ask' }), + ], + }), + }) + const { adapter, calls } = scripted([textTurn('one'), textTurn('two')]) + const host = createHarnessHost({ persistence: memoryPersistence() }) + const session = await host.open( + defineHarness({ + name: 'test/code-mode-include-guard', + adapter, + tools: [ + tools.lookup, + tools.remove, + toolDefinition({ + name: 'docs_deny', + description: 'A tool that a rule denies', + inputSchema: z.object({}), + }).server(async () => 'denied'), + toolDefinition({ + name: 'docs_ask', + description: 'A tool that a rule asks for', + inputSchema: z.object({}), + }).server(async () => 'asked'), + ], + plugins: () => [ + denyDocs, + codeMode({ driver: fakeDriver().driver, include: () => true }), + ], + }), + { threadId: 't' }, + ) + await session.send('first') + await session.send('second') + await host.close() + + for (const call of calls) { + expect(call.tools.map((tool: { name: string }) => tool.name)).toEqual([ + 'remove', + 'docs_deny', + 'docs_ask', + 'execute_typescript', + ]) + } + // One warning for each kept tool, not one for each turn. + expect(warn.mock.calls.map(([message]) => message)).toEqual([ + '[@tanstack/ai-code-mode] include picked the tool "remove", but it needs approval, or a rule asks for it or denies it. A call inside code mode skips those checks, so the tool stays a normal tool call.', + '[@tanstack/ai-code-mode] include picked the tool "docs_deny", but it needs approval, or a rule asks for it or denies it. A call inside code mode skips those checks, so the tool stays a normal tool call.', + '[@tanstack/ai-code-mode] include picked the tool "docs_ask", but it needs approval, or a rule asks for it or denies it. A call inside code mode skips those checks, so the tool stays a normal tool call.', + ]) + warn.mockRestore() + }) + it('moves a marked tool that include does not pick, unless a rule or an approval stops it', async () => { const rules = definePlugin({ name: 'test/marked-rules', @@ -418,6 +477,40 @@ describe('codeMode', () => { ]) }) + it('keeps a tool that include picks and that edits as a tool call in plan mode', async () => { + const { adapter, calls } = scripted([textTurn('ok'), textTurn('ok')]) + const host = createHarnessHost({ persistence: memoryPersistence() }) + const session = await host.open( + defineHarness({ + name: 'test/code-mode-include-plan', + adapter, + tools: [tools.lookup], + plugins: () => [ + permissions({ + rules: [{ tool: 'lookup', decision: 'allow', kind: 'edit' }], + }), + codeMode({ + driver: fakeDriver().driver, + include: (tool) => tool.name === 'lookup', + }), + ], + }), + { threadId: 't' }, + ) + const namesOf = (call: number) => + calls[call].tools.map((tool: { name: string }) => tool.name) + await session.send('hi') + // In the default mode, the rule allows it, so it moves. + expect(namesOf(0)).toEqual(['execute_typescript']) + + await session.command({ name: 'mode', input: 'plan' }) + await session.send('plan it') + // Plan mode denies an edit, so a call inside the isolate must not run + // it. No tool moves, so code mode adds no tool. + expect(namesOf(1)).toEqual(['lookup']) + await host.close() + }) + it('keeps a tool that a rule of permissions() asks for as a tool call', async () => { const names = await firstToolNames({ tools: [ diff --git a/packages/ai-harness-cli/src/args.ts b/packages/ai-harness-cli/src/args.ts index c5b89c797f..0fc59a389c 100644 --- a/packages/ai-harness-cli/src/args.ts +++ b/packages/ai-harness-cli/src/args.ts @@ -11,8 +11,8 @@ export interface CliArgs { port: number hostname: string token?: string + mediaKey?: string thread: string - dashboard?: string mediaDir?: string } @@ -32,8 +32,8 @@ Options: --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) + --media-key Key that signs media URLs of --serve, 32+ bytes (default: HARNESS_MEDIA_KEY, or a new random key at each start) --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 @@ -54,8 +54,8 @@ export function parseCliArgs(argv: ReadonlyArray): CliArgs { port: { type: 'string' }, host: { type: 'string' }, token: { type: 'string' }, + 'media-key': { type: 'string' }, thread: { type: 'string' }, - dashboard: { type: 'string' }, 'media-dir': { type: 'string' }, }, strict: true, @@ -79,8 +79,10 @@ export function parseCliArgs(argv: ReadonlyArray): CliArgs { port, hostname: values.host ?? '127.0.0.1', ...(values.token !== undefined ? { token: values.token } : {}), + ...(values['media-key'] !== undefined + ? { mediaKey: values['media-key'] } + : {}), 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/commands.ts b/packages/ai-harness-cli/src/commands.ts index 8f980f3691..8ab550724e 100644 --- a/packages/ai-harness-cli/src/commands.ts +++ b/packages/ai-harness-cli/src/commands.ts @@ -123,9 +123,14 @@ export async function handleLine( text: error instanceof Error ? error.message : String(error), } } - if (session.snapshot().status === 'running') - await session.send(input, { mode: 'steer' }).receipt - else void session.send(input) + // While a turn runs, the message joins it. Else it starts a turn. The + // session decides this when it stores the message, so there is no race. + const operation = session.send(input, { busy: 'steer' }) + // The view shows a failed turn from its events. + operation.then(undefined, () => {}) + const receipt = await operation.receipt + if (receipt.status === 'rejected') + return { type: 'notice', text: `Not sent: ${receipt.reason}` } 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 61459a8dde..81ed988690 100644 --- a/packages/ai-harness-cli/src/index.ts +++ b/packages/ai-harness-cli/src/index.ts @@ -24,7 +24,6 @@ import type { } from '@tanstack/ai-harness/view' import type { AcpAgentModule, - DashboardConnectModule, McpHarnessModule, McpStdioModule, } from './optional' @@ -54,9 +53,9 @@ export interface RunCliOptions

{ */ host?: HarnessHost

/** - * Who uses the CLI. Sessions open as this principal in every mode except - * `--dashboard`, so credentials and provider keys are theirs. On a host - * without access, `--serve` runs each request with the token as them. + * Who uses the CLI. Sessions open as this principal in every mode, so + * credentials and provider keys are theirs. On a host without access, + * `--serve` runs each request with the token as them. */ principal?: P /** Default: `process.argv.slice(2)`. */ @@ -68,8 +67,10 @@ 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. + * the other modes (`--print`, `--acp`, `--mcp`, `--serve`) do not use it. * `view.send(text)` sends each `@path` file in `text`, as line mode does. + * A saved child agent of a reopened thread shows only its card until the + * screen calls `view.loadAgent({ id })`. */ ui?: (view: SessionView) => Promise | void /** @@ -138,7 +139,7 @@ function attachingSource(session: HarnessSession) { snapshot: () => session.snapshot(), transcript: () => session.transcript(), describe: () => session.describe(), - mediaUrl: (id) => session.mediaUrl(id), + mediaUrl: (options) => session.mediaUrl(options), loadMedia: (options) => session.loadMedia(options), } return source @@ -260,44 +261,16 @@ export async function runCli

( return code } - if (args.dashboard) { - const dashboard = await loadOptional( - '@tanstack/ai-dashboard/connect', - 'connectDashboard', - ) - if (!dashboard) throw needs('--dashboard', '@tanstack/ai-dashboard') - const savedToken = env.HARNESS_DASHBOARD_TOKEN - // ponytail: the dashboard takes no principal, so its sessions open - // without one. - const connection = await dashboard.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`) - const code = await stop.stopped - connection.close() - return code - } - if (args.serve) { const token = args.token ?? env.HARNESS_TOKEN ?? createToken() + const mediaKey = args.mediaKey ?? env.HARNESS_MEDIA_KEY const server = await serve({ host, harness, port: args.port, hostname: args.hostname, token, + ...(mediaKey !== undefined ? { mediaKey } : {}), threadId: args.thread, approvals, ...(principal ? { principal } : {}), @@ -327,7 +300,7 @@ export async function runCli

( : stdin.isTTY && options.ui ? showUi(session, options.ui, stop.stopped) : runLines(session, stdin, stdout, { - openSignIns: Boolean(stdin.isTTY), + interactive: Boolean(stdin.isTTY), mediaDir, signal: stop.signal, onCtrlC: () => stop.receive('SIGINT'), diff --git a/packages/ai-harness-cli/src/lines.ts b/packages/ai-harness-cli/src/lines.ts index b8308d7823..944706961b 100644 --- a/packages/ai-harness-cli/src/lines.ts +++ b/packages/ai-harness-cli/src/lines.ts @@ -69,15 +69,17 @@ function secretTyping(input: NodeJS.ReadableStream, onCtrlC: () => void) { } /** - * 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 its waits, - * each next line answers one of them, by its kind and schema (see - * `answerOf`: `y` approves an approval). After the last one, the turn goes - * on. A question of a command or a durable tool takes the next line. In a - * terminal, the answer to a secret question (a key) is not shown as the - * user types. Line mode never prints the answer. + * The line mode: one message or command per line. When a turn stops for + * its waits, each next line answers one of them, by its kind and schema + * (see `answerOf`: `y` approves an approval). After the last one, the turn + * goes on. A question of a command or a durable tool takes the next line. + * In a terminal, the answer to a secret question (a key) is not shown as + * the user types. Line mode never prints the answer. * - * `openSignIns`: open sign-in links in the browser (for an interactive terminal). + * `interactive`: the user types the lines (a terminal). Each line acts at + * once: `/cancel` cancels the running turn, and a message steers it. Sign-in + * links open in the browser. Without it (piped input), each line waits for + * the turn that the line before it started, so a script runs in order. * `mediaDir`: the folder for the media a turn makes. Each new file is saved * there, with one `[image saved: ]` line. * `signal`: stop reading lines when it aborts. @@ -89,7 +91,7 @@ export async function runLines( input: NodeJS.ReadableStream, stdout: Output, options: { - openSignIns?: boolean + interactive?: boolean mediaDir: string signal: AbortSignal onCtrlC: () => void @@ -109,7 +111,7 @@ export async function runLines( // Line mode prints only what happens from now on. printer.mark(view.store.get()) const printing = view.store.subscribe((state) => printer.print(state)) - const stopOpening = options.openSignIns + const stopOpening = options.interactive ? view.on('signIn', (signIn) => { if (signIn.url) openUrl(signIn.url) }) @@ -124,6 +126,23 @@ export async function runLines( const pendingLater: Array> = [] /** The answers to the waits of the stopped turn so far, by wait id. */ const decided = new Map() + /** + * Wait for the turn to end. Then ask for its next wait that has no + * answer, and hide the typing for a secret question. In a terminal, this + * runs while the next lines are read. Only one runs at a time, so each + * wait is asked once. + */ + const afterTurn = async () => { + await waitIdle(session) + const after = session.snapshot() + const ask = + after.status === 'requires_action' + ? turnWaits(after).find(({ waitId }) => !decided.has(waitId)) + : undefined + if (ask) line(waitQuestion(ask)) + if (after.pendingQuestions[0]?.secret) typing.hide() + } + let waiting: Promise | undefined try { for await (const raw of lines) { // Lines that came before the stop do not run. @@ -151,19 +170,21 @@ export async function runLines( pendingLater.push(result.later.then((text) => line(text))) } } - await waitIdle(session) - const after = session.snapshot() - const ask = - after.status === 'requires_action' - ? turnWaits(after).find(({ waitId }) => !decided.has(waitId)) - : undefined - if (ask) line(waitQuestion(ask)) - if (after.pendingQuestions[0]?.secret) typing.hide() + if (!options.interactive) { + await afterTurn() + continue + } + // A secret question that the line asked hides the next line at once. + if (session.snapshot().pendingQuestions[0]?.secret) typing.hide() + waiting ??= afterTurn().finally(() => { + waiting = undefined + }) } } finally { typing.show() } lines.close() + await waiting await Promise.allSettled(pendingLater) await waitIdle(session) stopOpening() diff --git a/packages/ai-harness-cli/src/optional.ts b/packages/ai-harness-cli/src/optional.ts index 003219fd37..57bbbd46d2 100644 --- a/packages/ai-harness-cli/src/optional.ts +++ b/packages/ai-harness-cli/src/optional.ts @@ -1,25 +1,54 @@ import type { AnyHarness, HarnessHost, Principal } from '@tanstack/ai-harness' import type { OpenAccess } from '@tanstack/ai-harness/server' -// The packages behind `--acp`, `--mcp`, `/mcp`, and `--dashboard` are -// optional. The specifier is a variable, so a build or a type check never -// needs them, and a version without harness support is the same as none. +// The packages behind `--acp`, `--mcp`, and `/mcp` are optional. The +// specifier is a variable, so a build or a type check never needs them, and +// a version without harness support is the same as none. /** * Import `specifier` and check that it has the functions in `names`. - * Resolves to `undefined` when the package is not installed, or when it has - * no such functions. + * Resolves to `undefined` when the package is not installed, when its + * version does not export `specifier`, or when it has no such functions. + * Any other error (a syntax error, or a failing import inside the module) + * throws, with `specifier` in the message. */ export async function loadOptional( specifier: string, ...names: Array -): Promise { - const loaded: unknown = await import(/* @vite-ignore */ specifier).catch( - () => undefined, - ) +) { + // Resolve first: a resolve error is about `specifier` itself, so its code + // alone says "not installed". A missing package that the module imports + // fails the import below, and throws. + try { + import.meta.resolve(specifier) + } catch (error) { + if (isMissing(error)) return undefined + throw loadError(specifier, error) + } + let loaded: unknown + try { + loaded = await import(/* @vite-ignore */ specifier) + } catch (error) { + throw loadError(specifier, error) + } return hasFunctions(loaded, names) ? loaded : undefined } +/** True when a resolve error says that the package or its export is not there. */ +function isMissing(error: unknown) { + if (!(error instanceof Error) || !('code' in error)) return false + return ( + error.code === 'ERR_MODULE_NOT_FOUND' || + error.code === 'MODULE_NOT_FOUND' || + error.code === 'ERR_PACKAGE_PATH_NOT_EXPORTED' + ) +} + +function loadError(specifier: string, error: unknown) { + const reason = error instanceof Error ? error.message : String(error) + return new Error(`Could not load ${specifier}: ${reason}`, { cause: error }) +} + function hasFunctions( value: unknown, names: ReadonlyArray, @@ -66,19 +95,3 @@ export interface AcpAgentModule { close: () => void } } - -/** - * `@tanstack/ai-dashboard/connect`. It takes no principal: its sessions - * open without one. - */ -export interface DashboardConnectModule { - connectDashboard:

(options: { - host: HarnessHost

- harness: AnyHarness - url: string - threads: Array - token?: string - onPairingCode?: (code: string) => void - onToken?: (token: string) => void - }) => Promise<{ close: () => void }> -} diff --git a/packages/ai-harness-cli/src/printer.ts b/packages/ai-harness-cli/src/printer.ts index f39bb7aaf0..7c008e35a3 100644 --- a/packages/ai-harness-cli/src/printer.ts +++ b/packages/ai-harness-cli/src/printer.ts @@ -23,6 +23,11 @@ export function createPrinter( ) { const done = new Set() const written = new Map() + /** + * The messages and parts that printed. The view keeps a message or a part + * that did not change, so a print walks only the new ones. + */ + const printed = new WeakSet() // True while streamed text has no line break at its end yet. let midLine = false let lastStream = '' @@ -87,19 +92,25 @@ export function createPrinter( } const print = (state: SessionViewState) => { for (const message of state.messages) { + if (printed.has(message)) continue + printed.add(message) if (message.role === 'notice') { if (message.kind === 'error' || message.kind === 'info') once(message.id, message.text) continue } if (message.role !== 'assistant') continue + const fresh: Array = [] message.parts.forEach((part, index) => { + if (printed.has(part)) return + printed.add(part) + fresh.push(part) const key = `${message.id}:${index}` if (part.type === 'text') stream(key, part.text) if (part.type === 'tool-call') once(key, `tool ${part.name}`) if (part.type === 'agent') agent(key, part) }) - media(message.parts) + media(fresh) } for (const question of state.questions) { const isHidden = question.secret === true && options.hidesSecrets === true diff --git a/packages/ai-harness-cli/src/serve.ts b/packages/ai-harness-cli/src/serve.ts index a695efc495..ba6ea3f932 100644 --- a/packages/ai-harness-cli/src/serve.ts +++ b/packages/ai-harness-cli/src/serve.ts @@ -23,6 +23,13 @@ export interface ServeOptions

{ * media URL needs no token: its signature is its login. */ token: string + /** + * The key that signs media URLs, at least 32 bytes (UTF-8). Default: a new + * random key at each start, so a signed URL stops working after a restart. + * It is never made from `token`: a signed URL is given out without the + * token, and must not let anyone guess it. + */ + mediaKey?: string /** The conversation of an MCP tool call that names none. Default `'main'`. */ threadId?: string /** How `/mcp` handles tool calls that need approval. Default `'ask'`. */ @@ -134,16 +141,20 @@ export async function serve

( const isAllowed = (request: Request) => sameToken(request.headers.get('authorization'), options.token) const withAccess = host.hasAccess - // ponytail: the token signs media URLs, so they work as long as it does. - const media = { secret: options.token } + // Without `mediaKey`, the handler signs with a new random key, so signed + // media URLs stop working after a restart. + const media = + options.mediaKey !== undefined + ? { media: { secret: options.mediaKey } } + : {} const handler = withAccess - ? createHarnessHandler({ host, harness, media }) + ? createHarnessHandler({ host, harness, ...media }) : createHarnessHandler({ host, harness, // One principal holds the token, so it can use every thread. openAccess: principal ? { principal } : true, - media, + ...media, }) const mcp = await mcpServer(options, withAccess) const route = (request: Request) => { diff --git a/packages/ai-harness-cli/tests/attach.test.ts b/packages/ai-harness-cli/tests/attach.test.ts index 26db146163..6d4fd9c72a 100644 --- a/packages/ai-harness-cli/tests/attach.test.ts +++ b/packages/ai-harness-cli/tests/attach.test.ts @@ -96,7 +96,7 @@ async function storedFiles(input: UserInput) { const ids = parts.flatMap((part) => mediaIdOf(part) ?? []) return Promise.all( ids.map(async (id) => ({ - name: (await session.getMedia(id))?.name, + name: (await session.getMedia({ id }))?.name, text: new TextDecoder().decode(await session.loadMedia({ id })), })), ) diff --git a/packages/ai-harness-cli/tests/cli.test.ts b/packages/ai-harness-cli/tests/cli.test.ts index 93b25a6537..09a7fe64b2 100644 --- a/packages/ai-harness-cli/tests/cli.test.ts +++ b/packages/ai-harness-cli/tests/cli.test.ts @@ -1,3 +1,4 @@ +import { PassThrough } from 'node:stream' import { memoryCredentialStore } from '@tanstack/ai/byok/server' import { describe, expect, it, vi } from 'vitest' import { defineAgent, toolDefinition } from '@tanstack/ai' @@ -14,9 +15,19 @@ import { import { question as questionTool } from '@tanstack/ai-harness/plugins' import { createMCPClient } from '@tanstack/ai-mcp' import { memoryPersistence } from '@tanstack/ai-persistence' -import { EXIT, parseCliArgs, runCli } from '../src' +import { EXIT, USAGE, parseCliArgs, runCli } from '../src' import { serve } from '../src/serve' -import { capture, scripted, stdinFrom, textTurn, toolTurn } from './helpers' +import { + asStdin, + capture, + gate, + heldTurn, + scripted, + stdinFrom, + textTurn, + toolTurn, +} from './helpers' +import type { AnyHarness } from '@tanstack/ai-harness' /** `/connect acme` asks for a key as a secret question, then saves it. */ const acmeKeys = definePlugin({ @@ -77,6 +88,13 @@ describe('parseCliArgs', () => { }) expect(parseCliArgs([])).toMatchObject({ mcp: false, yes: false }) }) + + it('refuses --dashboard as an unknown flag', () => { + expect(() => parseCliArgs(['--dashboard', 'http://127.0.0.1:1'])).toThrow( + "Unknown option '--dashboard'", + ) + expect(USAGE).not.toContain('--dashboard') + }) }) describe('print mode', () => { @@ -216,6 +234,106 @@ describe('line mode', () => { }) }) +describe('line mode in a terminal', () => { + /** Line mode on a terminal that the test types into. */ + function terminal(harness: AnyHarness) { + const input = new PassThrough() + const stdout = capture() + const exit = runCli(harness, { + argv: [], + stdin: asStdin(input, { isTTY: true }), + stdout, + stderr: capture(), + persistence: memoryPersistence(), + }) + return { + stdout, + type: (text: string) => input.write(`${text}\n`), + shows: (text: string) => + vi.waitFor(() => expect(stdout.text).toContain(text)), + quit: () => { + input.end('/exit\n') + return exit + }, + } + } + + it('cancels the running turn at once with /cancel', async () => { + // Never opened: only the cancel can end this turn. + const release = gate() + const { adapter } = scripted([heldTurn('Working on it.', release.opened)]) + const cli = terminal(defineHarness({ name: 'test/cancel', adapter })) + cli.type('write a long story') + await cli.shows('Working on it.') + cli.type('/cancel') + await cli.shows('Cancelled.') + expect(await cli.quit()).toBe(EXIT.ok) + }) + + it('sends a line typed while a turn runs as a steer that the turn sees', async () => { + const release = gate() + const { adapter, seen } = scripted([ + heldTurn('Working on it.', release.opened), + textTurn('Checked the logs too.'), + ]) + const joined: Array = [] + const cli = terminal( + defineHarness({ + name: 'test/steer', + adapter, + turn: { + canJoin: (candidate) => { + joined.push(candidate.message) + return true + }, + }, + }), + ) + cli.type('fix the bug') + await cli.shows('Working on it.') + cli.type('also check the logs') + // The steer waits for the next model call of the running turn. + cli.type('/status') + await cli.shows('Queued turns: 1.') + release.open() + await cli.shows('Checked the logs too.') + expect(await cli.quit()).toBe(EXIT.ok) + expect(joined).toEqual(['also check the logs']) + expect(seen[1]?.at(-1)).toMatchObject({ + role: 'user', + content: 'also check the logs', + }) + }) + + it('asks for an approval once, and the next line answers it', async () => { + const remove = vi.fn(async () => 'removed') + const { adapter } = scripted([ + toolTurn('remove', '{}'), + textTurn('Done removing.'), + ]) + const cli = terminal( + defineHarness({ + name: 'test/approve', + adapter, + tools: [ + toolDefinition({ + name: 'remove', + description: 'Remove', + needsApproval: true, + }).server(remove), + ], + }), + ) + cli.type('remove the file') + await cli.shows('[y/n]') + cli.type('y') + await cli.shows('Done removing.') + expect(await cli.quit()).toBe(EXIT.ok) + expect(remove).toHaveBeenCalledTimes(1) + expect(cli.stdout.text.split('[y/n]')).toHaveLength(2) + }) +}) + describe('waits in line mode', () => { it('answers a question of the question tool by its schema', async () => { const { adapter, seen } = scripted([ @@ -631,6 +749,32 @@ describe('serve mode', () => { } }) + it('signs media URLs with --media-key, never with the token', async () => { + expect(parseCliArgs(['--serve', '--media-key', 'k'])).toMatchObject({ + mediaKey: 'k', + }) + const host = createHarnessHost({ persistence: memoryPersistence() }) + const start = (mediaKey?: string) => + serve({ + host, + harness: defineHarness({ + name: 'test/serve-key', + adapter: scripted([]).adapter, + }), + port: 0, + hostname: '127.0.0.1', + token: 'secret', + ...(mediaKey !== undefined ? { mediaKey } : {}), + }) + + // The key is checked: a short one is refused at the start. + await expect(start('short')).rejects.toThrow('Media signing key 0 is 5') + // A short token is fine: it does not sign anything. + const server = await start() + await server.close() + await host.close() + }) + it('reads the user of each request with identify on a host with access, behind the token', async () => { const host = accessHost() const server = await serve({ diff --git a/packages/ai-harness-cli/tests/helpers.ts b/packages/ai-harness-cli/tests/helpers.ts index 0cb846d311..5938d78092 100644 --- a/packages/ai-harness-cli/tests/helpers.ts +++ b/packages/ai-harness-cli/tests/helpers.ts @@ -54,11 +54,44 @@ export const toolTurn = (name: string, args: string): Array => [ }, ] +/** A model call that the test drives. It gets the abort signal of the call. */ +type LiveTurn = (signal: AbortSignal | undefined) => AsyncIterable + +/** A promise that the test resolves with `open`. */ +export function gate() { + let open = () => {} + const opened = new Promise((resolve) => { + open = resolve + }) + return { open, opened } +} + +/** Resolves when `signal` aborts. Never resolves without a signal. */ +function whenAborted(signal: AbortSignal | undefined) { + return new Promise((resolve) => { + if (signal?.aborted) resolve() + signal?.addEventListener('abort', () => resolve(), { once: true }) + }) +} + +/** + * A model call that streams `text`, then holds until `release` resolves or + * the call is aborted (a cancel). Then it ends. + */ +export function heldTurn(text: string, release: Promise): LiveTurn { + const chunks = textTurn(text) + return async function* (signal) { + yield* chunks.slice(0, 3) + await Promise.race([release, whenAborted(signal)]) + yield* chunks.slice(3) + } +} + /** * A model that answers each call with the next of `turns`, and keeps the * messages of each call in `seen`. */ -export function scripted(turns: Array>) { +export function scripted(turns: Array | LiveTurn>) { let call = 0 const seen: Array> = [] const adapter: AnyTextAdapter = { @@ -82,10 +115,17 @@ export function scripted(turns: Array>) { }, chatStream: (options) => { seen.push(options.messages) - const chunks = turns[call] ?? textTurn('') + const turn = turns[call] ?? textTurn('') call += 1 + if (typeof turn === 'function') { + return turn( + options.abortController?.signal ?? + options.request?.signal ?? + undefined, + ) + } return (async function* () { - yield* chunks + yield* turn })() }, structuredOutput: async () => ({ data: {}, rawText: '{}' }), diff --git a/packages/ai-harness-cli/tests/optional.test.ts b/packages/ai-harness-cli/tests/optional.test.ts index a448cfeffe..022ea3f5f5 100644 --- a/packages/ai-harness-cli/tests/optional.test.ts +++ b/packages/ai-harness-cli/tests/optional.test.ts @@ -1,10 +1,28 @@ -import { describe, expect, it } from 'vitest' +import { mkdtempSync, rmSync, writeFileSync } from 'node:fs' +import { tmpdir } from 'node:os' +import { join } from 'node:path' +import { pathToFileURL } from 'node:url' +import { afterAll, describe, expect, it, vi } from 'vitest' import { fakeText } from '@tanstack/ai/testing' import { defineHarness } from '@tanstack/ai-harness' import { memoryPersistence } from '@tanstack/ai-persistence' import { EXIT, runCli } from '../src' import { loadOptional } from '../src/optional' +// An installed version of @tanstack/ai-acp without harness support: it has +// no serveAcp function. +vi.mock('@tanstack/ai-acp/agent', () => ({ serveAcp: undefined })) + +const folder = mkdtempSync(join(tmpdir(), 'cli-optional-')) +afterAll(() => rmSync(folder, { recursive: true, force: true })) + +/** The URL of a temp module with `source`. */ +function tempModule(name: string, source: string) { + const path = join(folder, name) + writeFileSync(path, source) + return pathToFileURL(path).href +} + describe('optional packages', () => { it('loads a module that has the functions', async () => { const path = await loadOptional<{ join: () => string }>('node:path', 'join') @@ -18,27 +36,49 @@ describe('optional packages', () => { 'run', ), ).toBeUndefined() + expect( + await loadOptional<{ run: () => void }>( + '@tanstack/no-such-package/harness', + 'run', + ), + ).toBeUndefined() // An installed package with no harness support. expect( await loadOptional<{ serveAcp: () => void }>('node:path', 'serveAcp'), ).toBeUndefined() }) - it('fails a flag whose package is missing, with a clear message', async () => { - let text = '' - const harness = defineHarness({ - name: 'test/optional', - adapter: fakeText(), - }) - const code = await runCli(harness, { - argv: ['--dashboard', 'http://127.0.0.1:1'], - stdout: { write: () => {} }, - stderr: { write: (chunk) => (text += chunk) }, - persistence: memoryPersistence(), - }) + it('throws, with the specifier, for a module that fails on import', async () => { + const broken = tempModule('broken.mjs', 'throw new Error("boom")\n') + await expect( + loadOptional<{ run: () => void }>(broken, 'run'), + ).rejects.toThrow(`Could not load ${broken}: boom`) + }) + + it('throws for a missing package that the module imports', async () => { + const importer = tempModule( + 'importer.mjs', + 'import "@tanstack/no-such-dependency"\n', + ) + await expect( + loadOptional<{ run: () => void }>(importer, 'run'), + ).rejects.toThrow(`Could not load ${importer}`) + }) + + it('fails a flag whose package has no harness support, with a clear message', async () => { + let stderr = '' + const code = await runCli( + defineHarness({ name: 'test/optional', adapter: fakeText() }), + { + argv: ['--acp'], + stdout: { write: () => {} }, + stderr: { write: (chunk) => (stderr += chunk) }, + persistence: memoryPersistence(), + }, + ) expect(code).toBe(EXIT.failed) - expect(text).toContain( - '--dashboard needs @tanstack/ai-dashboard with harness support.', + expect(stderr).toContain( + '--acp needs @tanstack/ai-acp with harness support.', ) }) }) diff --git a/packages/ai-harness-cli/tests/printer.test.ts b/packages/ai-harness-cli/tests/printer.test.ts index a4f93c23c4..dc8bb14d4a 100644 --- a/packages/ai-harness-cli/tests/printer.test.ts +++ b/packages/ai-harness-cli/tests/printer.test.ts @@ -243,4 +243,39 @@ describe('line printer', () => { expect(out).toBe('') expect(onMedia).not.toHaveBeenCalled() }) + + it('walks only new messages and parts, and writes only the new text', () => { + const writes: Array = [] + const printer = createPrinter((text) => writes.push(text)) + let reads = 0 + const oldParts: Array = [{ type: 'text', text: 'Old.' }] + const old = { + id: 'turn-0', + role: 'assistant' as const, + get parts() { + reads += 1 + return oldParts + }, + } + const tool: ViewPart = { + type: 'tool-call', + id: 'c', + name: 'read', + argsText: '{}', + args: {}, + status: 'done', + } + const turn = (text: string) => ({ + id: 'turn-1', + role: 'assistant' as const, + parts: [tool, { type: 'text' as const, text }], + }) + + printer.print({ ...base, messages: [old] }) + printer.print({ ...base, messages: [old, turn('Hi')] }) + printer.print({ ...base, messages: [old, turn('Hi there')] }) + + expect(writes).toEqual(['Old.', '\n', '[tool read]\n', 'Hi', ' there']) + expect(reads).toBe(1) + }) }) diff --git a/packages/ai-harness/package.json b/packages/ai-harness/package.json index 1ca73466a6..19f5f00cf2 100644 --- a/packages/ai-harness/package.json +++ b/packages/ai-harness/package.json @@ -88,6 +88,7 @@ "@tanstack/store": "^0.11.1" }, "devDependencies": { + "@opentelemetry/api": "^1.9.0", "@tanstack/ai": "workspace:*", "@tanstack/ai-persistence": "workspace:*", "@vitest/coverage-v8": "4.1.10", @@ -96,12 +97,16 @@ "zod": "^4.2.0" }, "peerDependencies": { + "@opentelemetry/api": ">=1.9.0", "@tanstack/ai": "workspace:^", "@tanstack/ai-persistence": "workspace:^", "@vscode/ripgrep": "^1.15.0", "turndown": "^7.2.0" }, "peerDependenciesMeta": { + "@opentelemetry/api": { + "optional": true + }, "@vscode/ripgrep": { "optional": true }, diff --git a/packages/ai-harness/src/agents.ts b/packages/ai-harness/src/agents.ts index b3984575c6..a8d245f929 100644 --- a/packages/ai-harness/src/agents.ts +++ b/packages/ai-harness/src/agents.ts @@ -16,7 +16,7 @@ import type { AgentStartOptions, RunArgs, } from './session/types' -import type { AgentRun } from './types' +import type { AgentRun, RunContext } from './types' import type { PluginState } from './plugins' import type { AnyHarness } from './define' @@ -118,6 +118,11 @@ export class AgentRegistry implements AgentRegistryView { return this.agents.get(name)?.agent } + /** Who added the agent `name`: `'the harness'` or a plugin name. */ + ownerOf(name: string): string | undefined { + return this.agents.get(name)?.owner + } + find(query: { produces: AgentProduces }): AnyAgent | undefined { return this.list().find((agent) => agent.produces === query.produces) } @@ -304,15 +309,32 @@ export type HarnessChildRun = ( */ const childRuns = new WeakMap() -/** @internal Give `context` the child hook `run`. Returns `context`. */ +/** The run context of the run of each `ctx.harness` that a session made. */ +const runContexts = new WeakMap() + +/** + * @internal Give `context` the child hook `run`, and the run context of the + * run that it belongs to. Returns `context`. + */ export function withHarnessChild( context: HarnessAgentContext, run: HarnessChildRun, + runContext: RunContext, ): HarnessAgentContext { childRuns.set(context, run) + runContexts.set(context, runContext) return context } +/** + * @internal The run context of the run of `context`: for an agent that runs + * as a tool of a turn, the run of that turn. `undefined` outside a run that + * a harness session started. + */ +export function runContextOf(context: HarnessAgentContext | undefined) { + return context ? runContexts.get(context) : undefined +} + /** * @internal The child hook of `context`, or `undefined` outside a run that * a harness session started. diff --git a/packages/ai-harness/src/build.ts b/packages/ai-harness/src/build.ts index 7210fb7453..41ee2decc8 100644 --- a/packages/ai-harness/src/build.ts +++ b/packages/ai-harness/src/build.ts @@ -245,6 +245,8 @@ const FORWARDED = new Set([ interface Worker { send: (frame: object) => void frames: (listener: (frame: HostFrame) => void) => () => void + /** True for the `requestId` of a frame that was sent to this worker. */ + sent: (requestId: string) => boolean kill: () => void } @@ -257,10 +259,16 @@ function startWorker( const child = spawn(node, [bundle], { stdio: ['pipe', 'pipe', 'inherit'] }) const listeners = new Set<(frame: HostFrame) => void>() // A worker that fails to start or stops fails the streams that wait on it. + const requestIds = new Set() const fail = (message: string) => { onExit() for (const listener of listeners) - listener({ type: 'harness.error', message }) + listener({ + type: 'harness.error', + message, + code: 'internal', + requestId: 'worker', + }) } child.on('error', (error) => fail(`The worker failed: ${error.message}`)) child.on('exit', (code, signal) => @@ -275,8 +283,11 @@ function startWorker( // A line that is not a frame (for example a stray log) is skipped. } }) - const send = (frame: object) => + const send = (frame: object) => { + if ('requestId' in frame && typeof frame.requestId === 'string') + requestIds.add(frame.requestId) child.stdin.write(`${JSON.stringify(frame)}\n`) + } send({ type: 'harness.subscribe', threadId }) return { send, @@ -284,6 +295,7 @@ function startWorker( listeners.add(listener) return () => listeners.delete(listener) }, + sent: (requestId) => requestIds.has(requestId), kill: () => { child.stdin.end() child.kill() @@ -394,7 +406,15 @@ export async function artifactText( } const frame = queue.shift() if (!frame) continue - if (frame.type === 'harness.error') throw new Error(frame.message) + if (frame.type === 'harness.error') { + // The error of another input of this worker, for example the + // cancel of an older turn, is not about this turn. + if (frame.requestId !== requestId && worker.sent(frame.requestId)) + continue + throw new Error(frame.message, { + cause: { code: frame.code, requestId: frame.requestId }, + }) + } if ( frame.type === 'harness.receipt' && frame.requestId === requestId diff --git a/packages/ai-harness/src/client.ts b/packages/ai-harness/src/client.ts index d986f07eb6..74704b67b4 100644 --- a/packages/ai-harness/src/client.ts +++ b/packages/ai-harness/src/client.ts @@ -1,7 +1,8 @@ // Browser-safe: this module has type-only imports from the rest of the -// package, plus `media-ref`, which has type-only imports itself. So no server -// code (adapters, stores, secrets) reaches a client. +// package, plus `media-ref` and `utils`, which have no runtime imports +// themselves. So no server code (adapters, stores, secrets) reaches a client. import { isMediaRecord } from './media-ref' +import { isRecord } from './utils' import type { ModelMessage } from '@tanstack/ai' import type { LogRecord, @@ -17,6 +18,7 @@ import type { SessionMetadataOf, } from './define' import type { HostEvent } from './host' +import type { HarnessErrorCode } from './protocol' import type { SessionDescription, SessionSnapshot } from './session/index' import type { SessionEntryChange } from './session-entry' import type { @@ -52,6 +54,36 @@ export type { HostStatusEvent, } from './host' +/** + * A `harness.error` frame that ended a stream of the client: `code` says + * what failed. `events()` and `hostEvents()` throw it. A `slow_consumer` + * error is not thrown: they connect again. + */ +export class HarnessProtocolError extends Error { + readonly code: HarnessErrorCode + constructor(code: HarnessErrorCode, message: string) { + super(message) + this.name = 'HarnessProtocolError' + this.code = code + } +} + +/** + * The error of a `harness.error` frame. Undefined for another frame and + * for `slow_consumer`: the stream ends, and the client connects again. + */ +function protocolErrorOf(frame: unknown): HarnessProtocolError | undefined { + if (!isRecord(frame) || frame.type !== 'harness.error') return undefined + if (frame.code === 'slow_consumer') return undefined + const code = + typeof frame.code === 'string' + ? (frame.code as HarnessErrorCode) + : 'internal' + const message = + typeof frame.message === 'string' ? frame.message : 'Harness stream failed' + return new HarnessProtocolError(code, message) +} + export interface HarnessClientOptions { /** The base URL of `createHarnessHandler`, for example `/api/harness`. */ url: string @@ -59,8 +91,20 @@ export interface HarnessClientOptions { /** Extra headers, for example `Authorization`. */ headers?: Record | (() => Record) fetch?: typeof fetch - /** Wait before a reconnect of `events()`. Default 1000 ms. */ + /** + * The first wait before a reconnect of `events()`. Default 1000 ms. Each + * failed try after it doubles the wait, up to 30 s, with a random part. + * A read from the server sets the wait back to this value. + */ reconnectDelayMs?: number + /** + * A connection of `events()` that gets no bytes for this long is dead, + * for example after a laptop sleeps or a NAT drops the mapping. The + * client closes it and connects again from its last cursor. Default + * 60 000 ms. Keep it larger than the `stream.heartbeatMs` of the server + * (default 20 000 ms), or the client reconnects in a loop. + */ + idleTimeoutMs?: number } type StartArgs = @@ -76,7 +120,14 @@ export type ClientAgentHandles = { } /** The options of `HarnessClient.send`. */ -export type ClientSendOptions = { inputId?: string } & ( +export type ClientSendOptions = { + inputId?: string + /** + * A W3C `traceparent` of the caller, sent as the `traceparent` header. + * The span of the turn on the server links it. + */ + traceparent?: string +} & ( | { to?: undefined mode?: 'prompt' @@ -147,8 +198,13 @@ export interface HarnessClient { /** * Cancel an input that waits (see `snapshot().waitingInputs`). An input * that started is refused with `not_waiting`: use `cancel`. + * + * @example + * ```ts + * await client.cancelInput({ inputId }) + * ``` */ - cancelInput: (inputId: string) => Promise + cancelInput: (options: { inputId: string }) => Promise /** * Move an input that waits: `'steer'` joins the running turn, `'queue'` * runs it as its own turn later. @@ -193,9 +249,18 @@ export interface HarnessClient { }) => Promise /** * Add host records to the session log, as `session.append` does. The - * harness must list each record type in `expose.records`. + * harness must list each record type in `expose.records`. With `inputId`, + * a second append with the same id is a duplicate and writes nothing. + * + * @example + * ```ts + * await client.append({ records: [{ type: 'app.signal', text: 'Saved.' }] }) + * ``` */ - append: (records: ReadonlyArray) => Promise + append: (options: { + records: ReadonlyArray + inputId?: string + }) => Promise /** Change a session setting. */ setConfig: (change: { key: string; value: unknown }) => Promise /** @@ -220,12 +285,19 @@ export interface HarnessClient { /** * Go back to the message `messageId`, as `session.revert` does. The * harness must list `'undo'` in `expose.commands`. + * + * @example + * ```ts + * await client.revert({ messageId }) + * ``` */ - revert: (messageId: string) => Promise + revert: (options: { messageId: string }) => Promise /** End the revert that stands, as `session.unrevert` does. */ unrevert: () => Promise /** - * The session events from `from` (exclusive). Reconnects after a network + * The session events from `from` (exclusive). Without `from`, the events + * start at the cursor of the snapshot: only new events come. Pass + * `from: '0'` to read the events from the start. Reconnects after a network * error and resumes from the last cursor. Ends when `signal` aborts. * `onConnection` reports `'open'` for each connection and `'reconnecting'` * before each new try. When the log dropped the events after `from`, the @@ -251,8 +323,15 @@ export interface HarnessClient { * The saved messages of a child agent run of the thread, by its * `subagentRunId`. Empty for an unknown id. A child with no messages has * one message with no content, with its card. See `session.agentTranscript`. + * + * @example + * ```ts + * const messages = await client.agentTranscript({ subagentRunId }) + * ``` */ - agentTranscript: (subagentRunId: string) => Promise> + agentTranscript: (options: { + subagentRunId: string + }) => Promise> /** * The commands, settings, and tools of the session. Only the commands in * `expose.commands` and the config keys in `expose.config`. @@ -265,17 +344,24 @@ export interface HarnessClient { * * @example * ```ts - * const record = await client.upload({ body: file, mimeType: file.type, name: file.name }) + * const record = await client.putMedia({ body: file, mimeType: file.type, name: file.name }) * ``` */ - upload: ( + putMedia: ( file: MediaUpload, ) => Promise /** * A signed URL for a media file, for ``, `

{ ownership?: OwnershipOptions /** How the sessions of this host live. Leave it out for the defaults. */ lifecycle?: HostLifecycleOptions + /** + * Spans and operation reports. `tracer`: an OpenTelemetry tracer, or any + * object with the same `startSpan`. The host writes a `harness.operation` + * span for each turn, agent run, and command that it runs, and + * `harness.recovery`, `harness.wait`, and `harness.transfer` spans. A + * `traceparent` of the caller becomes a link of the span of its turn. + * `onOperation`: called when an operation that this host runs starts and + * when it settles, with its status, duration, and usage. Without it, the + * host writes no spans and reports nothing. + * + * @example + * ```ts + * import { trace } from '@opentelemetry/api' + * + * const host = createHarnessHost({ + * observability: { + * tracer: trace.getTracer('my-app'), + * onOperation: ({ phase, operation }) => { + * if (phase === 'settled') metrics.record(operation) + * }, + * }, + * }) + * ``` + */ + observability?: ObservabilityOptions + /** + * The harnesses whose threads this host may delete without an open first. + * `host.sessions.delete` ends a thread in the plugins of its harness, so + * they free what lives as long as the thread (an owned sandbox, a shadow + * repository). It finds the harness by the name in the index entry, among + * these and the harnesses that this host opened. A thread whose harness + * it does not know skips that step, with a warning. + * + * @example + * ```ts + * const host = createHarnessHost({ harnesses: [assistant] }) + * ``` + */ + harnesses?: ReadonlyArray } /** How the sessions of a host live. */ @@ -217,9 +261,8 @@ export interface RetentionLimits { settledAgentRuns?: number /** * How many bytes of events a session keeps in memory for `events()`, as - * the JSON length of each event. A reader whose cursor is older starts - * at the oldest event kept. Older events come from the log. Default - * 8 000 000. + * the JSON length of each event. A reader whose cursor is older reads + * the older events from the log. Default 8 000 000. */ eventBytes?: number /** @@ -412,10 +455,12 @@ export interface HostSessions

{ }) => Promise /** * Remove the user `userId` from the members of a thread. Needs - * `membersOnly()`, and the `share` action. The work that the user started - * on the session in this host stops, and `host.revoke` runs for the user. - * The owner cannot be removed. Resolves to the changed entry, or - * `undefined` for a thread that has no entry. + * `membersOnly()`, and the `share` action. The running and queued work + * that the user started on the thread stops, on any host, their schedules + * do not run, and `host.revoke` runs for the user. When that work does + * not stop, it rejects after the member is removed: call `host.revoke` + * for the user to stop it. The owner cannot be removed. Resolves to the + * changed entry, or `undefined` for a thread that has no entry. * * @example * ```ts @@ -772,9 +817,14 @@ export interface HarnessHost

{ * Take the access of the user `principal` to a thread, after your app * changed its own access tables. `host.events()` yields a `revoked` * event, so the user's streams close. On a host with access, the open - * handles of that user on the thread throw `AccessDeniedError` from now - * on, and a new `host.open` asks `authorize` again. The `tenantId` of - * `principal` picks the tenant. Only `id` and `tenantId` are read. + * handles and `host.observe` views of that user on the thread throw + * `AccessDeniedError` from now on, their running `events()` reads end + * with it, and a new `host.open` asks `authorize` again. The running and + * queued work of the user on the thread is cancelled on any host, and + * their schedules do not run. Resolves when the cancels are sent, and + * rejects when they could not be sent: the access is gone already, so + * call it again to stop the work. The `tenantId` of `principal` picks + * the tenant. Only `id` and `tenantId` are read. * * @example * ```ts @@ -784,7 +834,7 @@ export interface HarnessHost

{ revoke: (options: { threadId: string principal: { id: string; tenantId?: string } - }) => void + }) => Promise /** * The user of `request`, from `access.identify`, or `null` when * `identify` refuses it. Network code other than the HTTP handler (an MCP @@ -797,7 +847,9 @@ export interface HarnessHost

{ * The status changes and the session index changes of this host, for a * list of sessions that stays current. First the current status of each * open session, then each change as it occurs. Each call reads on its own. - * The read stops when `signal` aborts or the loop ends. + * The read stops when `signal` aborts or the loop ends. It also stops when + * more than 1 000 events wait for a slow reader: read again to get the + * current statuses. * * @example * ```ts @@ -931,6 +983,19 @@ function checkSessionStore(store: SessionIndexStore | undefined) { /** A thread id that `harnessAgent` derives for a child: `.../agent/...`. */ const isHarnessChildThread = (threadId: string) => threadId.includes('/agent/') +/** + * @internal The thread of the inner session of `harnessText` in the outer + * thread `outerThread`. Only an open marked with `throughOuterThread` for + * that outer thread opens it. + */ +export const innerThreadOf = (outerThread: string, harnessName: string) => + `${outerThread}/inner/${harnessName}` + +const isInnerThread = (threadId: string) => threadId.includes('/inner/') + +/** The index `metadata` key that links an inner thread to its outer thread. */ +const OUTER_THREAD = 'harness:outerThread' + /** * The key of an open session in the host: the harness, the tenant, and the * thread. One thread id in two tenants is two sessions. @@ -957,26 +1022,50 @@ export type SessionIndexWriter = ReturnType * question waits. When an operation ends or a question gets its answer, the * session waits while a question or an interrupt still waits, runs while * another operation runs, and else is idle. + * + * The questions of the snapshot come from the folded log, so the questions + * of every host count. Its operations are only the work of this host: + * `seen.running` keeps the operations that run from the events, so the + * events that another host wrote count too. */ -function statusAfter(session: HarnessSession, event: StreamChunk) { +function statusAfter( + session: HarnessSession, + event: StreamChunk, + seen: { running: Set }, +) { if (event.type !== EventType.CUSTOM) return undefined + const value = isRecord(event.value) ? event.value : {} + const { operationId } = value switch (event.name) { case HARNESS_EVENTS.operationStarted: + if (typeof operationId === 'string') seen.running.add(operationId) return 'running' + case HARNESS_EVENTS.ownershipChanged: + // The old owner stopped: its operations never finish. + if (value.reason === 'taken_over') seen.running.clear() + return undefined + case HARNESS_EVENTS.operationResumed: + // The input of a stopped host runs again as a new operation. The old + // operation never finishes. + if (typeof value.resumedFrom === 'string') { + seen.running.delete(value.resumedFrom) + } + return undefined case HARNESS_EVENTS.question: return 'waiting' case HARNESS_EVENTS.operationFinished: case HARNESS_EVENTS.questionAnswered: { + if (typeof operationId === 'string') seen.running.delete(operationId) const { pendingQuestions, pendingInterrupts, activeOperations } = session.snapshot() - const isInterrupted = - isRecord(event.value) && event.value.status === 'interrupted' + const isInterrupted = value.status === 'interrupted' const stillWaits = isInterrupted || pendingQuestions.length > 0 || pendingInterrupts.length > 0 if (stillWaits) return 'waiting' - return activeOperations.length > 0 ? 'running' : 'idle' + const isRunning = activeOperations.length > 0 || seen.running.size > 0 + return isRunning ? 'running' : 'idle' } default: return undefined @@ -1010,6 +1099,30 @@ export function sessionOfTurn

( return turnFinders.get(host)?.(lookup) ?? Promise.resolve(undefined) } +/** A read of the host events of one tenant. */ +export interface TenantEventsRead { + /** The tenant to read. `undefined` is the default tenant. */ + tenantId: string | undefined + /** Runs before the queue: an event it rejects takes no memory. */ + filter: (event: HostEvent) => boolean + signal?: AbortSignal +} + +/** How each host reads its events for one tenant. Not public API. */ +const tenantEventReaders = new WeakMap< + object, + (read: TenantEventsRead) => AsyncIterable +>() + +/** + * @internal The `host.events()` of one tenant, with the statuses of that + * tenant first. `filter` runs before the queue. `undefined` for a host that + * `createHarnessHost` did not make. + */ +export function tenantEventsOf

(host: HarnessHost

) { + return tenantEventReaders.get(host) +} + /** A log that a sweep took over, and the end of what the sweep opened. */ export interface TakenLog { thread: ResumedThread @@ -1031,7 +1144,8 @@ export interface HostSweep { * Acquire the log of `orphan`, open its thread with `harness`, and take * its work over, as `host.resumePending` does for each log. Resolves * `null` when another host got the log. Rejects when the thread does not - * open: the log goes back, so a later sweep tries again. + * open: the log goes back with a wake of `OPEN_RETRY_MS`, so a sweep + * after that time tries again. */ takeOver: ( orphan: OrphanedLog, @@ -1144,6 +1258,20 @@ function forkedMembers( return { ownerId: forker.id, members: [owner, ...old, ...others] } } +/** + * The host records that the log fold of `harness` keeps: the `keepRecords` + * of its plugins, with the larger count for a type in two plugins. + */ +function keptRecordsOf(harness: AnyHarness) { + const keep = new Map() + for (const plugin of harness.plugins?.() ?? []) { + for (const [type, count] of Object.entries(plugin.keepRecords ?? {})) { + keep.set(type, Math.max(keep.get(type) ?? 0, count)) + } + } + return keep +} + /** * The messages of `messages` up to the last point where each tool call has * its result. A running turn can have a call with no result yet: it does @@ -1203,6 +1331,25 @@ function checkDurableStores(persistence: HarnessPersistence) { } } +const accessThreads = new WeakMap() + +/** + * @internal Mark `openOptions` for the inner session of `harnessText`, with + * the thread id from `innerThreadOf(outerThread, ...)`. On a host with access + * that has an entry for `outerThread`, `host.open` with these options checks + * the session through that thread, as a child thread is checked through its + * root, so every user of the outer thread can use it. The first open links + * the inner entry to `outerThread`. An inner entry with another link gets + * the usual checks. + */ +export function throughOuterThread( + openOptions: OpenSessionOptions, + outerThread: string, +) { + accessThreads.set(openOptions, outerThread) + return openOptions +} + /** * Create a host for harness sessions. With `access`, the principal type of * the host is what `access.identify` returns, and each call is checked (see @@ -1267,21 +1414,9 @@ export function createHarnessHost

( } const logStore = persistence.stores.log const { sessions: sessionStore } = persistence.stores - /** One per running `host.events()` read. */ - const listeners = new Set<(event: HostEvent) => void>() - const emit = (event: HostEvent) => { - for (const listener of [...listeners]) { - // One reader that throws must not stop the others. - try { - listener(event) - } catch (error) { - console.error( - '[@tanstack/ai-harness] A host events reader threw.', - error, - ) - } - } - } + /** The readers of `host.events()`. */ + const bus = new HostEventBus() + const emit = (event: HostEvent) => bus.emit(event) /** The latest status of each open session, by session key. */ const statuses = new Map() const setStatus = ( @@ -1339,12 +1474,14 @@ export function createHarnessHost

( if (open) return open const project = harness.durability?.project const reduce = harness.durability?.reduce + const keep = keptRecordsOf(harness) const loaded = loadLogState({ store, logId, rootThread, ...(project ? { project } : {}), ...(reduce ? { reduce } : {}), + keep, }).then((state) => { const log = new SharedLog({ store, @@ -1357,6 +1494,7 @@ export function createHarnessHost

( eventBytes: retention.eventBytes, ...(project ? { project } : {}), ...(reduce ? { reduce } : {}), + keep, ...(options.log?.snapshot ? { snapshot: options.log.snapshot } : {}), logRecords: retention.logRecords, // The next open folds the log again. @@ -1373,6 +1511,15 @@ export function createHarnessHost

( return loaded } const sessions = new Map>() + /** + * The last harness definition that this host opened under each name, + * first from the `harnesses` option. A thread purge reads the harness + * name in the index entry, and mounts the plugins of that definition to + * end the thread (see `endThread`). + */ + const harnessesByName = new Map( + (options.harnesses ?? []).map((harness) => [harness.name, harness]), + ) /** The plugins with `lifetime: 'host'`, shared by the sessions of a harness. */ const hostPlugins = new HostPlugins() /** The sessions of `sessions` that finished opening, by key. */ @@ -1609,6 +1756,7 @@ export function createHarnessHost

( }, ...(access ? { access: { envKeys: access.envKeys === true } } : {}), openChild: openChildIn(ids, harness, tenantId), + ...(options.observability ? { observability: options.observability } : {}), }) /** @@ -1624,6 +1772,7 @@ export function createHarnessHost

( ): OpenChildSession => async (harness, { threadId, principal, child }) => { if (isClosed) throw hostClosed() + harnessesByName.set(harness.name, harness) const created: SessionImpl = new SessionImpl({ ...sessionDeps(ids, rootHarness, tenantId), harness, @@ -1662,6 +1811,7 @@ export function createHarnessHost

( tenantId: string | undefined, ) => { if (isClosed) return Promise.reject(hostClosed()) + harnessesByName.set(harness.name, harness) const key = sessionKey(harness.name, tenantId, threadId) const of = { tenantId, threadId } touch(key) @@ -1680,6 +1830,8 @@ export function createHarnessHost

( // The log of a session tree has the id of its root thread, in its // tenant. Only a root opens: its child threads write to this log. const storeLogId = logIdOf({ tenantId, logId: threadId }) + // For the status rule. See `statusAfter`. + const seen = { running: new Set() } const created = new SessionImpl({ ...sessionDeps( { logId: storeLogId, rootThread: threadId }, @@ -1693,7 +1845,7 @@ export function createHarnessHost

( onEvent: (event) => { // A session that left the cache has no status any more. if (sessions.get(key) !== session) return - const status = statusAfter(created, event) + const status = statusAfter(created, event, seen) if (status === 'running') cancelIdle(key) if (status) setStatus(key, of, status) }, @@ -1748,10 +1900,11 @@ export function createHarnessHost

( // Of two hosts that sweep at once, one gets the log. const epoch = held ?? (await ownership.acquire(storeLogId, name)) if (epoch === null) return null - // Only an ownership that this sweep acquired goes back here. - const giveBack = () => + // Only an ownership that this sweep acquired goes back here. `wakeIn` + // lists the log again after that time, so a later sweep tries again. + const giveBack = (wakeIn?: number) => held === undefined - ? ownership.release(storeLogId, epoch) + ? ownership.release(storeLogId, epoch, wakeIn) : Promise.resolve() // The store id of a log is its tenant and its root thread. const { tenantId, logId: threadId } = parseLogId(storeLogId) @@ -1777,9 +1930,10 @@ export function createHarnessHost

( return null } } catch (error) { - // The log goes back, so the next sweep tries again. + // The log goes back with a wake, so a sweep after the wait tries + // again. With no wake, no sweep lists the log again. if (isOwn) endSweep(key) - await giveBack().catch(() => {}) + await giveBack(OPEN_RETRY_MS).catch(() => {}) throw error } const thread: ResumedThread = { @@ -1797,8 +1951,29 @@ export function createHarnessHost

( // ponytail: one count per revoked user and thread, for the life of the // host. Drop the old ones if a host revokes millions. const revocations = new Map() + /** + * The reads that run now (`handle.events()`, a view of `host.observe`), + * by revoked key: a revoke ends them. + */ + const revokeListeners = new Map void>>() + /** Call `listener` at the next revoke of `key`. Returns the stop. */ + const onRevoke = (key: string, listener: () => void) => { + let listeners = revokeListeners.get(key) + if (!listeners) { + listeners = new Set() + revokeListeners.set(key, listeners) + } + const set = listeners + set.add(listener) + return () => { + set.delete(listener) + if (set.size === 0 && revokeListeners.get(key) === set) { + revokeListeners.delete(key) + } + } + } - /** Count a revoke, and tell the readers of `host.events()`. */ + /** Count a revoke, end the reads of the user, and tell `host.events()`. */ const revokeUser = ({ tenantId, threadId, @@ -1810,6 +1985,7 @@ export function createHarnessHost

( }) => { const key = revokedKey(tenantId, threadId, principalId) revocations.set(key, (revocations.get(key) ?? 0) + 1) + for (const listener of [...(revokeListeners.get(key) ?? [])]) listener() emit({ type: 'revoked', ...(tenantId !== undefined ? { tenantId } : {}), @@ -1936,6 +2112,7 @@ export function createHarnessHost

( const count = revocations.get(key) ?? 0 return { isRevoked: () => (revocations.get(key) ?? 0) !== count, + onRevoke: (listener) => onRevoke(key, listener), authorize: async (action, run, wait) => allow({ principal, @@ -1952,10 +2129,15 @@ export function createHarnessHost

( * `openSession` with the checks of `host.open`: with access, a thread that * exists needs the `read` action, and a new thread is checked unless its * open made the principal its owner. Hold the key first. + * + * With `accessThread` (the inner session of `harnessText`), a thread of + * that id on this host is checked instead, as a child thread is checked + * through its root. Without an entry for it, the checks are the usual ones. */ const openChecked = async ( harness: THarness, openOptions: OpenSessionOptions

, + accessThread?: string, ) => { if (!access) { return openSession(harness, openOptions, openOptions.principal?.tenantId) @@ -1963,6 +2145,22 @@ export function createHarnessHost

( const principal = need(openOptions.principal, 'host.open') const { threadId } = openOptions const ids = { tenantId: principal.tenantId, threadId } + const outer = + accessThread === undefined + ? undefined + : await index.get({ + tenantId: principal.tenantId, + threadId: accessThread, + }) + if (outer && accessThread !== undefined) { + await allow({ + principal, + action: 'read', + threadId: accessThread, + entry: outer, + }) + return openSession(harness, openOptions, principal.tenantId) + } const before = await index.get(ids) // The open refuses a child thread, and says why. const isRoot = before !== undefined && before.parentThreadId === undefined @@ -1992,18 +2190,22 @@ export function createHarnessHost

( session, openOptions, key, + accessThread, }: { harness: THarness session: SessionImpl openOptions: OpenSessionOptions

key: string + /** See `openChecked`. The handle checks each action on this thread. */ + accessThread?: string }) => { const { principal } = openOptions const { threadId } = session + const guarded = accessThread ?? threadId return sessionHandle({ session, principal, - ...(access && principal ? { guard: guardOf(principal, threadId) } : {}), + ...(access && principal ? { guard: guardOf(principal, guarded) } : {}), lease: { current: () => { // After `host.close()`, a call refuses at once. @@ -2016,7 +2218,7 @@ export function createHarnessHost

( }, reopen: () => { if (isClosed) throw hostClosed() - return openChecked(harness, openOptions) + return openChecked(harness, openOptions, accessThread) }, hold: () => hold(key), release: (releaseOptions) => { @@ -2046,6 +2248,83 @@ export function createHarnessHost

( ) } + /** + * Stop the work of the user `userId` on a thread, on any host, after + * `unshare` or `revoke` took their access. The sessions open here cancel + * it at once. Then each input of the user that did not end gets a cancel + * through the input lane, so the owner of the log stops it, on any host. + * Each schedule of the user takes its input id in the lane, so it never + * runs: the store keeps each input id once. `principal` is who stops the + * work, and the log keeps it. A failure rejects: the access is gone + * already, so the caller can stop the work again. + */ + const stopWorkOf = async ({ + tenantId, + threadId, + userId, + principal, + }: { + tenantId: string | undefined + threadId: string + userId: string + principal: P | undefined + }) => { + try { + const open = await openSessionsOf(tenantId, threadId) + for (const session of open) { + await session.cancelWorkOf({ userId, principal }) + } + const entry = await index.get({ tenantId, threadId }) + // The lane needs the harness of the log. + if (entry?.harness === undefined) return + const harness = entry.harness + const logId = logIdOf({ tenantId, logId: threadId }) + const state = await loadLogState({ + store: logStore, + logId, + rootThread: threadId, + }) + const root = sessionOf(state, threadId) + /** The sessions here stopped the operations that they know. */ + const isKnownHere = (operationId: string) => + open.some((session) => session.operation(operationId) !== undefined) + const offer = ( + inputId: string, + record: Parameters[0]['record'], + ) => logStore.offer({ logId, harness, inputId, record }) + for (const logged of root.inputs.values()) { + const isTheirs = + logged.principal?.id === userId && WORK_OPS.has(logged.input.op) + const hasEnded = + logged.status === 'settled' || + logged.status === 'rejected' || + logged.status === 'joined' + if (!isTheirs || hasEnded) continue + const { operationId } = logged + if (operationId !== undefined && isKnownHere(operationId)) continue + const input: HarnessInput = + operationId === undefined + ? { op: 'cancelInput', inputId: logged.inputId } + : { op: 'cancel', operationId } + const inputId = createInputId() + await offer(inputId, offeredInputRecord({ inputId, input, principal })) + } + for (const [scheduleId, schedule] of root.thread.schedules) { + if (schedule.principal?.id !== userId || !schedule.input) continue + // The id that the schedule offers its input with when it fires. + await offer(`schedule:${scheduleId}`, { + type: 'harness.schedule.done', + scheduleId, + }) + } + } catch (error) { + throw new Error( + `The access of ${JSON.stringify(userId)} to thread ${JSON.stringify(threadId)} is gone, but their work did not stop: ${String(error)}. Call host.revoke again to stop it.`, + { cause: error }, + ) + } + } + /** The thread ids of the index entries with `threadId` as parent. */ const childrenOf = async (tenantId: string | undefined, threadId: string) => { if (!sessionStore) return [] @@ -2063,10 +2342,63 @@ export function createHarnessHost

( return children } + /** + * Tell the plugins of a thread that it ends for good (`ThreadEnds`), so + * they free what lives as long as the thread: an owned sandbox, a shadow + * repository. The sessions of the thread are closed already, so no turn + * runs. A short session with only the log and the session plugins runs + * the handlers, and closes again. A child thread lives in the log of its + * root, so the session reads that log. It needs the harness definition + * that the index entry names, from the `harnesses` option or an open on + * this host. Without it, it warns and runs nothing. With no entry, there + * is nothing to run. A failed handler only warns: the purge of the data + * goes on. + */ + const endThread = async (tenantId: string | undefined, threadId: string) => { + const entry = await index.get({ tenantId, threadId }) + const name = entry?.harness + if (!entry || name === undefined) return + const harness = harnessesByName.get(name) + if (!harness) { + console.warn( + `[@tanstack/ai-harness] Thread ${JSON.stringify(threadId)} is deleted, but its plugins did not end it: this host does not know the harness ${JSON.stringify(name)}. What lives as long as the thread (an owned sandbox, a shadow repository) is not freed. Pass the harness in the \`harnesses\` option of createHarnessHost, or open it on this host first.`, + ) + return + } + const root = await rootOf(entry) + // The harness of the root folds the log, as in `openChildIn`. + const rootHarness = harnessesByName.get(root.harness ?? name) ?? harness + const session = new SessionImpl({ + ...sessionDeps( + { + logId: logIdOf({ tenantId, logId: root.threadId }), + rootThread: root.threadId, + }, + rootHarness, + tenantId, + ), + harness, + threadId, + onEvent: () => {}, + onClose: () => {}, + }) + try { + await session.openToEnd() + await session.endThread() + } catch (error) { + console.warn( + `[@tanstack/ai-harness] A plugin did not end thread ${JSON.stringify(threadId)}: ${String(error)}`, + ) + } finally { + // As on a crash: the log gets no write of this short session. + await session.close({ recoverable: true }).catch(() => {}) + } + } + /** * `host.sessions.delete` after its check: close the open sessions of the - * thread, purge its children depth first, then delete its data and its - * entry. `seen` stops a loop of parent ids. + * thread, end it in its plugins, purge its children depth first, then + * delete its data and its entry. `seen` stops a loop of parent ids. */ const purge = async ( tenantId: string | undefined, @@ -2082,6 +2414,7 @@ export function createHarnessHost

( const session = await sessions.get(key)?.catch(() => undefined) if (session) await closeShared(key, session) } + await endThread(tenantId, threadId) const children = await childrenOf(tenantId, threadId) for (const child of children) await purge(tenantId, child, seen) const storeId = logIdOf({ tenantId, logId: threadId }) @@ -2188,6 +2521,163 @@ export function createHarnessHost

( } } + /** `host.transfer` without its span. */ + const transferThread = async ({ + from, + to, + envelope, + principal, + transferId: id, + }: TransferOptions

): Promise => { + if (isClosed) throw hostClosed() + const transferId = id ?? crypto.randomUUID() + // The tenant of the principal is the tenant of both threads. + const tenantId = principal?.tenantId + const threadId = to.threadId ?? `thread-${transferId}` + if (threadId === from.threadId) { + throw new Error('A transfer needs another thread: to is from.') + } + const user = access ? need(principal, 'host.transfer') : undefined + const sourceEntry = await index.get({ tenantId, threadId: from.threadId }) + const targetEntry = await index.get({ tenantId, threadId }) + if (user) { + // No entry is not found: the same answer as a refusal. + if (!sourceEntry) { + throw new AccessDeniedError({ + action: 'transfer', + threadId: from.threadId, + }) + } + await allow({ + principal: user, + action: 'transfer', + threadId: from.threadId, + entry: sourceEntry, + }) + // A target that exists has members already: the principal must + // be one who may send to it. A retry finds its own new target. + if (targetEntry) { + await allow({ + principal: user, + action: 'send', + threadId, + entry: targetEntry, + }) + } + } + checkHarness(sourceEntry, from.harness) + checkHarness(targetEntry, to.harness) + const source = { harness: from.harness.name, threadId: from.threadId } + const target = { harness: to.harness.name, threadId } + const sourceKey = sessionKey(from.harness.name, tenantId, from.threadId) + hold(sourceKey) + try { + // With no principal, as `host.fork` opens its source. The open does + // not take the log from a host that owns it. + const sourceSession = await openSession( + from.harness, + { threadId: from.threadId }, + tenantId, + ) + const targetKey = sessionKey(to.harness.name, tenantId, threadId) + hold(targetKey) + try { + // A new target gets the principal as its owner and only member. + const targetSession = await openSession( + to.harness, + { threadId, ...(principal ? { principal } : {}) }, + tenantId, + ) + // First the target, then the source: a retry after a crash in + // between finds the target end, and writes only the source end. + let end = targetSession.transferOf(transferId) + if (!end) { + const isUsed = + to.threadId === undefined && + (await targetSession.transcript()).length > 0 + if (isUsed) { + throw new Error( + `Thread ${JSON.stringify(threadId)} is not empty, and has no transfer ${JSON.stringify(transferId)}.`, + ) + } + const settled = settledMessages(await sourceSession.transcript()) + const mode = envelope.messages ?? 'none' + const moved: Array = + mode === 'all' + ? // New ids: the target may have these messages from an + // earlier transfer. + settled.map((message) => ({ + ...message, + id: createMessageId(), + })) + : mode === 'summary' + ? [ + { + id: createMessageId(), + role: 'assistant', + content: `Summary of the conversation so far:\n${await sourceSession.transferSummary(principal, settled)}`, + }, + ] + : [] + const note: ModelMessage = { + id: createMessageId(), + role: 'user', + content: envelope.note, + } + end = { + from: source, + ...storedPrincipal(principal), + at: Date.now(), + } + // One append: the messages and the record, or neither. + await targetSession.adoptMessages( + sourceSession, + [note, ...moved], + [{ type: 'harness.transferred', transferId, ...end }], + ) + } + // The same record on each retry, so the input lane takes it as a + // duplicate. + await sourceSession.recordTransfer({ + type: 'harness.transferred', + transferId, + to: target, + ...(end.principal ? { principal: end.principal } : {}), + at: end.at, + }) + await index.update({ tenantId, threadId }, (entry) => + withTransferLink(entry, 'harness:transferredFrom', { + transferId, + ...source, + }), + ) + await index.update({ tenantId, threadId: from.threadId }, (entry) => + withTransferLink(entry, 'harness:transferredTo', { + transferId, + ...target, + }), + ) + emit({ + type: 'transferred', + transferId, + ...(tenantId !== undefined ? { tenantId } : {}), + from: source, + to: target, + ...(end.principal ? { principalId: end.principal.id } : {}), + at: end.at, + }) + return { transferId, threadId } + } finally { + // A target that only this transfer opened closes on idle. + await release(targetKey) + } + } finally { + await release(sourceKey) + } + } + + const tracer = options.observability?.tracer + const host: HarnessHost

= { async open(harness, openOptions) { if (isClosed) throw hostClosed() @@ -2197,13 +2687,46 @@ export function createHarnessHost

( `${JSON.stringify(openOptions.threadId)} is the thread of a harnessAgent child. Open its root thread and read the child through it.`, ) } + // See `throughOuterThread`. An inner thread opens only through the + // outer thread that its id names. + const { threadId } = openOptions + const outerThread = accessThreads.get(openOptions) + const isOwnInner = + outerThread !== undefined && + threadId.startsWith(`${outerThread}/inner/`) + if (isInnerThread(threadId) && !isOwnInner) { + throw new Error( + `${JSON.stringify(threadId)} is the thread of a harnessText inner session. Open its outer thread, and use the harness there.`, + ) + } const tenantId = openOptions.principal?.tenantId - const key = sessionKey(harness.name, tenantId, openOptions.threadId) + const key = sessionKey(harness.name, tenantId, threadId) // Before the open, so the session never closes on idle under it. hold(key) try { - const session = await openChecked(harness, openOptions) - return handleOf({ harness, session, openOptions, key }) + const [outer, inner] = + access && isOwnInner + ? await Promise.all([ + index.get({ tenantId, threadId: outerThread }), + index.get({ tenantId, threadId }), + ]) + : [] + // A new inner thread, or one linked to this outer thread. + const isLinked = + inner === undefined || inner.metadata?.[OUTER_THREAD] === outerThread + const accessThread = outer && isLinked ? outerThread : undefined + const session = await openChecked(harness, openOptions, accessThread) + if (accessThread !== undefined && inner === undefined) { + await index.update({ tenantId, threadId }, (entry) => + entry && entry.metadata?.[OUTER_THREAD] === undefined + ? { + ...entry, + metadata: { ...entry.metadata, [OUTER_THREAD]: accessThread }, + } + : undefined, + ) + } + return handleOf({ harness, session, openOptions, key, accessThread }) } catch (error) { await release(key) throw error @@ -2213,8 +2736,9 @@ export function createHarnessHost

( if (isClosed) throw hostClosed() const tenantId = principal?.tenantId const entry = await index.get({ tenantId, threadId }) + let user: P | undefined if (access) { - const user = need(principal, 'host.observe') + user = need(principal, 'host.observe') // No entry is not found: the same answer as a refusal. if (!entry) throw new AccessDeniedError({ action: 'read', threadId }) // `allow` checks a child thread through its root thread. @@ -2224,6 +2748,9 @@ export function createHarnessHost

( // A child thread has its records in the log of its root, with its // thread on each. A thread with no entry is a root. const root = entry ? (await rootOf(entry)).threadId : threadId + // A revoke of the user on the thread (its root) ends the view: each + // later call throws, and a read that runs now stops. + const revoked = user && guardOf(user, root) const ids = { logId: logIdOf({ tenantId, logId: root }), rootThread: root, @@ -2242,7 +2769,11 @@ export function createHarnessHost

( }, check: () => { if (isClosed) throw hostClosed() + if (revoked?.isRevoked()) { + throw new AccessDeniedError({ action: 'read', threadId }) + } }, + ...(revoked ? { onRevoke: revoked.onRevoke } : {}), index, media, readLog: () => readLog(logStore, ids, harness, threadId), @@ -2264,159 +2795,52 @@ export function createHarnessHost

( } return handleOf({ harness, session: fork, openOptions, key }) }, - async transfer({ from, to, envelope, principal, transferId: id }) { - if (isClosed) throw hostClosed() - const transferId = id ?? crypto.randomUUID() - // The tenant of the principal is the tenant of both threads. - const tenantId = principal?.tenantId - const threadId = to.threadId ?? `thread-${transferId}` - if (threadId === from.threadId) { - throw new Error('A transfer needs another thread: to is from.') - } - const user = access ? need(principal, 'host.transfer') : undefined - const sourceEntry = await index.get({ tenantId, threadId: from.threadId }) - const targetEntry = await index.get({ tenantId, threadId }) - if (user) { - // No entry is not found: the same answer as a refusal. - if (!sourceEntry) { - throw new AccessDeniedError({ - action: 'transfer', - threadId: from.threadId, - }) - } - await allow({ - principal: user, - action: 'transfer', - threadId: from.threadId, - entry: sourceEntry, - }) - // A target that exists has members already: the principal must - // be one who may send to it. A retry finds its own new target. - if (targetEntry) { - await allow({ - principal: user, - action: 'send', - threadId, - entry: targetEntry, - }) - } - } - checkHarness(sourceEntry, from.harness) - checkHarness(targetEntry, to.harness) - const source = { harness: from.harness.name, threadId: from.threadId } - const target = { harness: to.harness.name, threadId } - const sourceKey = sessionKey(from.harness.name, tenantId, from.threadId) - hold(sourceKey) - try { - // With no principal, as `host.fork` opens its source. The open does - // not take the log from a host that owns it. - const sourceSession = await openSession( - from.harness, - { threadId: from.threadId }, - tenantId, - ) - const targetKey = sessionKey(to.harness.name, tenantId, threadId) - hold(targetKey) - try { - // A new target gets the principal as its owner and only member. - const targetSession = await openSession( - to.harness, - { threadId, ...(principal ? { principal } : {}) }, - tenantId, - ) - // First the target, then the source: a retry after a crash in - // between finds the target end, and writes only the source end. - let end = targetSession.transferOf(transferId) - if (!end) { - const isUsed = - to.threadId === undefined && - (await targetSession.transcript()).length > 0 - if (isUsed) { - throw new Error( - `Thread ${JSON.stringify(threadId)} is not empty, and has no transfer ${JSON.stringify(transferId)}.`, + // Each transfer runs in a `harness.transfer` span. + transfer: (transferOptions) => + tracer + ? withSpan( + { + tracer, + name: 'harness.transfer', + attributes: { + 'harness.thread_id': transferOptions.from.threadId, + ...(transferOptions.to.threadId !== undefined + ? { + 'harness.transfer.to_thread_id': + transferOptions.to.threadId, + } + : {}), + ...(transferOptions.principal?.tenantId !== undefined + ? { 'harness.tenant_id': transferOptions.principal.tenantId } + : {}), + }, + }, + async (span) => { + const result = await transferThread(transferOptions) + span.setAttribute('harness.transfer.id', result.transferId) + span.setAttribute( + 'harness.transfer.to_thread_id', + result.threadId, ) - } - const settled = settledMessages(await sourceSession.transcript()) - const mode = envelope.messages ?? 'none' - const moved: Array = - mode === 'all' - ? // New ids: the target may have these messages from an - // earlier transfer. - settled.map((message) => ({ - ...message, - id: createMessageId(), - })) - : mode === 'summary' - ? [ - { - id: createMessageId(), - role: 'assistant', - content: `Summary of the conversation so far:\n${await sourceSession.transferSummary(principal, settled)}`, - }, - ] - : [] - const note: ModelMessage = { - id: createMessageId(), - role: 'user', - content: envelope.note, - } - end = { - from: source, - ...storedPrincipal(principal), - at: Date.now(), - } - // One append: the messages and the record, or neither. - await targetSession.adoptMessages( - sourceSession, - [note, ...moved], - [{ type: 'harness.transferred', transferId, ...end }], - ) - } - // The same record on each retry, so the input lane takes it as a - // duplicate. - await sourceSession.recordTransfer({ - type: 'harness.transferred', - transferId, - to: target, - ...(end.principal ? { principal: end.principal } : {}), - at: end.at, - }) - await index.update({ tenantId, threadId }, (entry) => - withTransferLink(entry, 'harness:transferredFrom', { - transferId, - ...source, - }), - ) - await index.update({ tenantId, threadId: from.threadId }, (entry) => - withTransferLink(entry, 'harness:transferredTo', { - transferId, - ...target, - }), + return result + }, ) - emit({ - type: 'transferred', - transferId, - ...(tenantId !== undefined ? { tenantId } : {}), - from: source, - to: target, - ...(end.principal ? { principalId: end.principal.id } : {}), - at: end.at, - }) - return { transferId, threadId } - } finally { - // A target that only this transfer opened closes on idle. - await release(targetKey) - } - } finally { - await release(sourceKey) - } - }, - revoke: ({ threadId, principal }) => + : transferThread(transferOptions), + revoke: async ({ threadId, principal }) => { + // At once, so the handles of the user refuse each call from now on. revokeUser({ tenantId: principal.tenantId, threadId, principalId: principal.id, - }), + }) + // No principal stops it: the app changed its own access tables. + await stopWorkOf({ + tenantId: principal.tenantId, + threadId, + userId: principal.id, + principal: undefined, + }) + }, hasAccess: access !== undefined, async identify({ request }) { if (!access) { @@ -2759,47 +3183,34 @@ export function createHarnessHost

( updatedAt: Date.now(), } }) - // After the write, so no new work of the user gets through. The log + // After the write, so no new work of the user gets through. At + // once, so the reads of the user end before the work stops. The log // keeps who cancelled it. - const open = await openSessionsOf(tenantId, threadId) - for (const session of open) { - await session.cancelWorkOf({ userId, principal: user }) - } revokeUser({ tenantId, threadId, principalId: userId }) + await stopWorkOf({ tenantId, threadId, userId, principal: user }) return changed }, }, - async *events({ signal } = {}) { + events({ signal } = {}) { // The current statuses first. A change while the reader handles an // event waits in the queue. - const queue: Array = [...statuses.values()] - let wake = () => {} - const listener = (event: HostEvent) => { - queue.push(event) - wake() - } - const onAbort = () => wake() - listeners.add(listener) - signal?.addEventListener('abort', onAbort) - try { - while (!signal?.aborted) { - const next = queue.shift() - if (next) { - yield next - continue - } - await new Promise((resolve) => { - wake = resolve - }) - } - } finally { - // A loop that ends (break, return, throw, abort) stops listening. - listeners.delete(listener) - signal?.removeEventListener('abort', onAbort) - } + return bus.read({ + initial: [...statuses.values()], + ...(signal ? { signal } : {}), + }) }, } if (access) accesses.set(host, access) + tenantEventReaders.set(host, ({ tenantId, filter, signal }) => + bus.read({ + tenant: { tenantId }, + filter, + initial: [...statuses.values()].filter( + (event) => tenantOfHostEvent(event) === tenantId && filter(event), + ), + ...(signal ? { signal } : {}), + }), + ) sweeps.set(host, { listOrphaned: (query) => logStore.listOrphaned({ ...query, ...pinOf(ownership) }), diff --git a/packages/ai-harness/src/http.ts b/packages/ai-harness/src/http.ts index 589e3488d4..cd0edc321f 100644 --- a/packages/ai-harness/src/http.ts +++ b/packages/ai-harness/src/http.ts @@ -10,27 +10,43 @@ import { } from '@tanstack/ai' import { parseRangeHeader, resolveBlobRange } from '@tanstack/ai-persistence' import { AccessDeniedError } from './access' -import { accessOf, sessionOfTurn } from './host' +import { accessOf, sessionOfTurn, tenantEventsOf } from './host' +import { tenantOfHostEvent } from './host-event-bus' import { activeRunIn } from './log-stores' import { MediaError } from './media' +import { maxBodyBytesOf, readBodyText, storeInlineMedia } from './inline-media' +import { createMediaSigner } from './media-signing' import { - HARNESS_PROTOCOL_VERSION, applyInput, capabilitiesOf, describeForClient, + HARNESS_PROTOCOL_VERSION, + versionError, parseControlFrame, parseHarnessInput, + requestIdOf, + retriedMessage, turnReceipt, } from './protocol' +import { socketSender, sseStream, streamValues, withHeartbeat } from './sse' import { base64url, isRecord } from './utils' import { SessionMetadataError } from './session-entry' import { SessionImpl, scrubbedBy } from './session/harness-session' -import { liveSessionOf, storedRunsOf } from './session/observed-session' +import { + liveSessionOf, + mediaStoreOf, + storedRunOf, + storedRunsOf, +} from './session/observed-session' +import { wireFrame } from './wire' import type { ModelMessage, UIMessage, WebSocketLike } from '@tanstack/ai' import type { SessionIndexEntry, SessionMember } from '@tanstack/ai-persistence' import type { AnyHarness } from './define' import type { HarnessHost, HostEvent } from './host' -import type { HostFrame } from './protocol' +import type { InlineMediaRefusal } from './inline-media' +import type { MediaSignerOptions } from './media-signing' +import type { ControlFrame, HarnessErrorCode, HostFrame } from './protocol' +import type { HarnessStreamOptions, SseFrame, StreamValues } from './sse' import type { HarnessSession } from './session/index' import type { ObservedSession } from './session/observed-session' import type { Cursor, Principal, SessionEvent, UserInput } from './types' @@ -59,24 +75,37 @@ export interface HarnessHandlerOptions

{ */ openAccess?: OpenAccess

/** - * How the handler signs media URLs. Without it, the handler signs with a - * random key of its own, so a signed URL works only on this server, and - * only until it restarts. Set it when more than one server answers the - * same URL, so each server verifies the URLs of the others. + * How the handler signs media URLs (HMAC SHA-256). Without it, the handler + * signs with a random key of its own, so a signed URL works only on this + * server, and only until it restarts. Set it when more than one server + * answers the same URL, so each server verifies the URLs of the others. + * + * - `{ secret }`: one key. Use the same value on every server. + * - `{ keys }`: to rotate keys. The first key signs, and every key + * verifies, so a URL that an old key signed works until it expires. + * + * Each key must be at least 32 bytes (UTF-8). A shorter key throws at the + * start. */ - media?: { - /** - * The secret that signs media URLs (HMAC SHA-256). Use the same value on - * every server behind one URL. - */ - secret: string - } + media?: MediaSignerOptions + /** + * Set `true` when the app also serves `handleHarnessSocket`, so the + * capabilities say that clients can use a WebSocket. Default `false`. + */ + socket?: boolean + /** + * The values of the SSE streams: the heartbeat, the slow-client close, and + * the queue size. A value that you do not set keeps its default. See + * `HarnessStreamOptions`. + */ + stream?: HarnessStreamOptions } /** How long a signed media URL works. */ const SIGNED_URL_TTL_MS = 60 * 60 * 1000 -const HMAC = { name: 'HMAC', hash: 'SHA-256' } +/** The most bytes one answer to an open-ended `Range` (`bytes=N-`) sends. */ +const OPEN_RANGE_MAX_BYTES = 8 * 1024 * 1024 // Every media answer: a browser must not guess another type, and a `text/*` // document must not run as a page of this origin. @@ -103,43 +132,67 @@ const json = (body: unknown, status = 200) => const forbidden = () => json({ error: 'forbidden' }, 403) /** - * An SSE response: one event per frame, with `data` as JSON and `id` when - * the frame has one. `watch` gets the reader: aborting it ends the stream. - * When the client leaves, the reader aborts. + * The id of a request, for its error frames and its `X-Request-Id` answer + * header. The rule: the `X-Request-Id` header of the request when it is a + * short safe id (1 to 128 letters, digits, `-`, `_`, `.`, or `:`), else a + * new UUID. So a proxy can trace a request, and a client cannot put other + * text into the logs. + */ +function requestIdFor(request: Request) { + const sent = request.headers.get('X-Request-Id') + return sent !== null && /^[\w.:-]{1,128}$/.test(sent) + ? sent + : crypto.randomUUID() +} + +/** + * An SSE response from `sseStream`: pull-based, with a heartbeat and a close + * for a slow client (see `sse.ts`). `watch` gets the reader: aborting it + * ends the stream. When the stream ends or the client leaves, the reader + * aborts. */ function sseResponse( - frames: ( - signal: AbortSignal, - ) => AsyncIterable<{ id?: string; data: unknown }>, + frames: (signal: AbortSignal) => AsyncIterable, + requestId: string, + stream: StreamValues, watch?: (reader: AbortController) => void, ) { // Not `request.signal`: some servers abort it once the body is read. const reader = new AbortController() - let isCancelled = false - const body = new ReadableStream({ - cancel: () => { - isCancelled = true - reader.abort() - }, - async start(controller) { - watch?.(reader) - try { - for await (const { id, data } of frames(reader.signal)) { - const idLine = id === undefined ? '' : `id: ${id}\n` - controller.enqueue( - encoder.encode(`${idLine}data: ${JSON.stringify(data)}\n\n`), - ) - } - } finally { - if (!isCancelled) controller.close() - // The watch ends with the stream. - reader.abort() - } + watch?.(reader) + const body = sseStream({ + frames: (signal) => { + // The watch ends with the stream. + signal.addEventListener('abort', () => reader.abort(), { once: true }) + return frames(signal) }, + requestId, + signal: reader.signal, + stream, }) return new Response(body, { headers: SSE_HEADERS }) } +/** The error of a `?v=` that asks for another version. See `versionError`. */ +function versionOf(url: URL) { + const v = url.searchParams.get('v') + return versionError(v === null ? undefined : Number(v)) +} + +/** + * The answer to a stream with a `?v=` that the host does not speak: an SSE + * stream with one `unsupported_version` error frame. Then it closes. + */ +function versionRefused( + error: NonNullable>, + requestId: string, +) { + const frame: HostFrame = { type: 'harness.error', requestId, ...error } + return new Response(encoder.encode(`data: ${JSON.stringify(frame)}\n\n`), { + headers: SSE_HEADERS, + }) +} + /** * Session events as an AG-UI SSE response. Each event has its cursor as * `id`, so a client that loses the connection resumes with `Last-Event-ID`. @@ -148,16 +201,19 @@ function sseResponse( */ function cursorResponse( read: (signal: AbortSignal) => AsyncIterable, + stream: StreamValues, watch?: (reader: AbortController) => void, ) { // Not `request.signal`: some servers abort it once the body is read. const reader = new AbortController() watch?.(reader) - const cursors: Array = [] + // The stream encodes each event right after it pulls it, before the next + // pull, so the cursor of the last event is the only one it needs. + let cursor: string | undefined async function* chunks() { try { for await (const entry of read(reader.signal)) { - cursors.push(entry.cursor) + cursor = entry.cursor yield entry.event } } finally { @@ -166,16 +222,36 @@ function cursorResponse( } } return new Response( - toServerSentEventsStream(chunks(), reader, (_chunk, index) => - cursors.at(index), + withHeartbeat( + toServerSentEventsStream(chunks(), reader, () => cursor), + { heartbeatMs: stream.heartbeatMs }, ), { headers: SSE_HEADERS }, ) } +/** + * The status of a content part that `storeInlineMedia` refused: 413 for too + * many bytes, 415 for a kind the harness does not take, else 400. + */ +const REFUSAL_STATUS: Record = { + media_too_large: 413, + media_not_allowed: 415, + media_invalid: 400, + media_url_not_allowed: 400, +} + +/** True when `text` is more than `limit` bytes as UTF-8. */ +function isOverLimit(text: string, limit: number) { + // A UTF-16 code unit is 1 to 3 bytes, so only some texts need encoding. + if (text.length > limit) return true + if (text.length * 3 <= limit) return false + return encoder.encode(text).byteLength > limit +} + /** * A handler error as JSON: a refused action is 403, a `MediaError` keeps its - * status, the rest is 400. + * status (413 for a body over `media.maxBodyBytes`), the rest is 400. */ const failed = (error: unknown) => { if (error instanceof AccessDeniedError) return forbidden() @@ -258,16 +334,6 @@ function parseSessionOp(body: unknown) { } } -/** Base64url to bytes, or `undefined` for a value that is not base64url. */ -function fromBase64url(value: string) { - try { - const binary = atob(value.replace(/-/g, '+').replace(/_/g, '/')) - return Uint8Array.from(binary, (char) => char.charCodeAt(0)) - } catch { - return undefined - } -} - /** * The bytes of a media file as a response, with `Range` support: 206 for a * slice, 416 for a range outside the file, and 404 for an id this thread @@ -278,24 +344,39 @@ async function mediaResponse( view: ObservedSession, id: string, ) { - const record = await view.getMedia(id) + const record = await view.getMedia({ id }) if (!record) return json({ error: 'not found' }, 404) - const range = parseRangeHeader(request.headers.get('Range'), record.size) - if (range === 'unsatisfiable') { + const asked = parseRangeHeader(request.headers.get('Range'), record.size) + if (asked === 'unsatisfiable') { return new Response(null, { status: 416, headers: { ...MEDIA_HEADERS, 'Content-Range': `bytes */${record.size}` }, }) } - const bytes = await view.loadMedia({ id, range }) + // A range with no end (`bytes=N-`) gets at most OPEN_RANGE_MAX_BYTES. A + // 206 may send less than was asked: the client asks again for the rest. + const range = + asked && asked.length === undefined + ? { + offset: asked.offset, + length: Math.min(OPEN_RANGE_MAX_BYTES, record.size - asked.offset), + } + : asked + const length = range + ? resolveBlobRange(record.size, range).length + : record.size + // A stream, so a big file is not held in memory. + const body = + (await mediaStoreOf(view)?.open(id, range)) ?? + (await view.loadMedia({ id, range })) const headers = { ...MEDIA_HEADERS, 'Content-Type': record.mimeType, - 'Content-Length': String(bytes.byteLength), + 'Content-Length': String(length), } - if (!range) return new Response(bytes, { headers }) - const { offset, length } = resolveBlobRange(record.size, range) - return new Response(bytes, { + if (!range) return new Response(body, { headers }) + const { offset } = range + return new Response(body, { status: 206, headers: { ...headers, @@ -384,9 +465,18 @@ function watchRevoke

( onRevoke: () => void, ) { if (principal === undefined) return + const isRevoke = (event: HostEvent) => isRevokeOf(event, principal, threadId) + // Only the revokes of this principal, from the bus of its tenant: the + // filter runs before the queue, so other events cost this watch nothing, + // and the first revoke ends the watch, so its queue does not fill. + // A host that `createHarnessHost` did not make has `host.events()`. + const readTenant = tenantEventsOf(host) + const events = readTenant + ? readTenant({ tenantId: principal.tenantId, filter: isRevoke, signal }) + : host.events({ signal }) void (async () => { - for await (const event of host.events({ signal })) { - if (!isRevokeOf(event, principal, threadId)) continue + for await (const event of events) { + if (!isRevoke(event)) continue onRevoke() return } @@ -464,11 +554,21 @@ function watchRevoke

( * The streams of a user close when `host.revoke` or `unshare` takes their * access. Each chat input runs as the principal of its request. * - * Signed media URLs use `media.secret`. Without it, they work only on this - * server (see `HarnessHandlerOptions.media`). + * Signed media URLs use the keys of `media`. Without it, they work only on + * this server (see `HarnessHandlerOptions.media`). A URL of a deleted file + * is 404. + * + * A request body over `media.maxBodyBytes` (default 32 MiB) is 413. Inline + * media of `POST run` goes to the media store, as in `POST control`. A + * refused part is 413 (`media_too_large`), 415 (`media_not_allowed`), or + * 400 (`media_invalid`, `media_url_not_allowed`), with the reason as + * `error`. + * + * `stream` sets the values of the SSE streams (see `HarnessStreamOptions`). * * Throws at the start for a host without `access` and without - * `openAccess`, and for a host with both. + * `openAccess`, for a host with both, and for a `stream` value that is not + * a finite number above 0. * * @example * ```ts @@ -481,11 +581,7 @@ export function createHarnessHandler

( const { host, harness, media } = options const usersOf = requestUsers('createHarnessHandler', host, options.openAccess) const hasAccess = accessOf(host) !== undefined - if (media !== undefined && media.secret.length === 0) { - throw new Error( - 'createHarnessHandler needs a media.secret that is not empty.', - ) - } + const stream = streamValues('createHarnessHandler', options.stream) const openFor = (principal: P | undefined, threadId: string) => host.open(harness, { threadId, ...(principal ? { principal } : {}) }) @@ -550,18 +646,17 @@ export function createHarnessHandler

( } } - // ponytail: without `media.secret`, a random key per handler, so a signed - // URL works only on this server, until it restarts. - const secret = - media === undefined - ? crypto.getRandomValues(new Uint8Array(32)) - : encoder.encode(media.secret) - let key: Promise | undefined - const mediaKey = () => - (key ??= crypto.subtle.importKey('raw', secret, HMAC, false, [ - 'sign', - 'verify', - ])) + // ponytail: without `media`, a random key per handler, so a signed URL + // works only on this server, until it restarts. + const signer = createMediaSigner( + media ?? { secret: base64url(crypto.getRandomValues(new Uint8Array(32))) }, + ) + /** + * A request body as JSON, read up to `media.maxBodyBytes`. A bigger body is + * a `MediaError` with 413. + */ + const readJson = async (request: Request): Promise => + JSON.parse(await readBodyText(request, harness.media)) /** What a signed media URL holds: who may read which file, until when. */ interface SignedMedia { tenantId: string @@ -570,25 +665,18 @@ export function createHarnessHandler

( userId: string exp: string } + // The harness name and the user are in the signed text, so a URL works + // only for this harness and for the user that it was signed for. const signed = (media: SignedMedia) => - encoder.encode( - [ - harness.name, - media.tenantId, - media.threadId, - media.id, - media.userId, - media.exp, - ].join('\n'), - ) - const sign = async (media: SignedMedia) => { - const signature = await crypto.subtle.sign( - 'HMAC', - await mediaKey(), - signed(media), - ) - return base64url(new Uint8Array(signature)) - } + [ + harness.name, + media.tenantId, + media.threadId, + media.id, + media.userId, + media.exp, + ].join('\n') + const sign = (media: SignedMedia) => signer.sign(signed(media)) /** * Serve a signed media URL. The URL names its user: the read is checked as @@ -603,17 +691,10 @@ export function createHarnessHandler

( userId: param('userId'), exp: param('exp'), } - const signature = fromBase64url(param('sig')) const isValid = - signature !== undefined && Number(payload.exp) > Date.now() && // `verify` compares in constant time. - (await crypto.subtle.verify( - 'HMAC', - await mediaKey(), - signature, - signed(payload), - )) + (await signer.verify(signed(payload), param('sig'))) if (!isValid) return forbidden() const { tenantId, userId } = payload // A signed URL keeps the id and the tenant of its user, as the log does, @@ -663,6 +744,7 @@ export function createHarnessHandler

( const from = resumeFrom ?? startedCursor ?? '0' return cursorResponse( (signal) => turn.events({ from, signal }), + stream, closeOnRevoke(principal, threadId), ) } @@ -678,7 +760,7 @@ export function createHarnessHandler

( const threadId = named const view = await viewFor(principal, threadId) if (!view) return json({ error: 'unknown run' }, 404) - const record = (await storedRunsOf(view)).get(runId) + const { record, start } = await storedRunOf(view, runId) const isChatTurn = record !== undefined && record.threadId === threadId && @@ -689,14 +771,15 @@ export function createHarnessHandler

( if (entry?.harness !== undefined && entry.harness !== harness.name) { return json({ error: 'unknown run' }, 404) } - // ponytail: from the start of the log without a cursor. Store the - // started position on the run record if long threads make this slow. + // Without a cursor, from the start of the run: its first input record. + // A fold snapshot from before input positions does not know it: from 0. return cursorResponse( (signal) => followOperation( - view.events({ from: resumeFrom ?? '0', signal }), + view.events({ from: resumeFrom ?? start ?? '0', signal }), runId, ), + stream, closeOnRevoke(principal, threadId), ) } @@ -704,33 +787,54 @@ export function createHarnessHandler

( /** * The `host.events()` that `principal` may read. With access, a thread * needs the `read` action. The decision of each thread is kept for the - * stream, until a revoke or an index change of that thread. + * stream, until a revoke or a delete of that thread. An index change of + * the thread drops only a decision that is not `true` (for example after + * a share), so the index changes of each turn read no store. `accepts` + * reads no store, so the host event bus runs it before the queue. + * `mayWatch` runs after the queue. */ const watchedEvents = (principal: P | undefined) => { const user = principal ? { principal } : {} const decisions = new Map>() + /** The threads whose kept decision is `true`. */ + const readable = new Set() + const forget = (threadId: string) => { + decisions.delete(threadId) + readable.delete(threadId) + } const mayRead = (threadId: string) => { const known = decisions.get(threadId) if (known) return known - const decision = host.sessions.get({ threadId, ...user }).then( - (entry) => entry !== undefined, - () => false, - ) + const decision: Promise = host.sessions + .get({ threadId, ...user }) + .then( + (entry) => { + // A revoke while the read ran dropped this decision. + const isKept = decisions.get(threadId) === decision + if (entry !== undefined && isKept) readable.add(threadId) + return entry !== undefined + }, + () => false, + ) decisions.set(threadId, decision) return decision } - return async (event: HostEvent) => { - const tenantId = - event.type === 'session' || event.type === 'session-deleted' - ? event.entry.tenantId - : event.tenantId - if (tenantId !== principal?.tenantId) return false + const accepts = (event: HostEvent) => { + if (tenantOfHostEvent(event) !== principal?.tenantId) return false if (event.type === 'revoked') { - decisions.delete(event.threadId) + forget(event.threadId) // A user hears of their own revokes only. return event.principalId === principal?.id } - if (!hasAccess) return true + if (event.type === 'session-deleted') forget(event.threadId) + // A user who may read the thread keeps that until a revoke. + if (event.type === 'session' && !readable.has(event.threadId)) { + forget(event.threadId) + } + return true + } + const mayWatch = async (event: HostEvent) => { + if (!hasAccess || event.type === 'revoked') return true // A transfer names two threads: the user must read both. if (event.type === 'transferred') { const reads = await Promise.all([ @@ -739,9 +843,7 @@ export function createHarnessHandler

( ]) return reads.every(Boolean) } - if (event.type === 'session') decisions.delete(event.threadId) if (event.type === 'session-deleted') { - decisions.delete(event.threadId) // The entry is gone: decide on the removed one. const access = accessOf(host) return principal !== undefined && access !== undefined @@ -754,19 +856,24 @@ export function createHarnessHandler

( } return mayRead(event.threadId) } + return { accepts, mayWatch } } - const handle = async (request: Request, principal: P | undefined) => { + const handle = async ( + request: Request, + principal: P | undefined, + requestId: string, + ) => { const url = new URL(request.url) const route = url.pathname.split('/').at(-1) const user = principal ? { principal } : {} if (request.method === 'GET' && route === 'capabilities') { - return json(capabilitiesOf(harness)) + return json(capabilitiesOf(harness, { socket: options.socket === true })) } if (request.method === 'POST' && route === 'run') { - const params = await chatParamsFromRequestBody(await request.json()) + const params = await chatParamsFromRequestBody(await readJson(request)) const message = lastUserInput(params.messages) const session = await openFor(principal, params.threadId) // The turn runs as the request's run id, so an AG-UI client matches @@ -787,29 +894,62 @@ export function createHarnessHandler

( } if (!receipt.operationId) return json({ error: 'rejected' }, 409) const { operationId } = receipt + // Without a cursor, from the start of the resumed run, not from + // the start of the log. + const from = + lastEventId ?? + session + .snapshot() + .activeOperations.find((item) => item.id === operationId) + ?.startedCursor return cursorResponse( (signal) => followOperation( - session.events({ - ...(lastEventId ? { from: lastEventId } : {}), - signal, - }), + session.events({ ...(from ? { from } : {}), signal }), operationId, ), + stream, keep(watch), ) } if (message === undefined) { return json({ error: 'no user message' }, 400) } + // Inline bytes go to the media store, and the message keeps a + // `harness-media:` reference, as `POST control` does. A retried run + // stores no media again. + const retry = retriedMessage(session, params.runId, message) + const checked = + retry !== undefined + ? { parts: retry } + : Array.isArray(message) + ? await storeInlineMedia(message, { + options: harness.media, + putMedia: (upload) => session.putMedia(upload), + }) + : { parts: message } + if ('refused' in checked) { + return json( + { error: checked.refused }, + REFUSAL_STATUS[checked.refused], + ) + } // AG-UI `forwardedProps` are the input's context. const { forwardedProps } = params - const operation = session.send(message, { + // The span of the turn links the `traceparent` of the caller. + const traceparent = request.headers.get('traceparent') + const sent = { ...ids, ...(Object.keys(forwardedProps).length > 0 ? { context: forwardedProps } : {}), - }) + ...(traceparent ? { traceparent } : {}), + } + // Only the session class takes the run id. `host.open` gives one. + const operation = + session instanceof SessionImpl + ? session.sendForRun(checked.parts, sent) + : session.send(checked.parts, sent) // A turn that the host refused throws, and the answer is 403. const receipt = await turnReceipt(operation) if (receipt.status === 'rejected') { @@ -819,9 +959,14 @@ export function createHarnessHandler

( // sends the same request with `Last-Event-ID`, maybe to another // instance: the run id finds the same turn, and it streams from // there. - const from = lastEventId || '0' + // Without it, the operation reads from its own start. return cursorResponse( - (signal) => operation.events({ from, signal }), + (signal) => + operation.events({ + ...(lastEventId ? { from: lastEventId } : {}), + signal, + }), + stream, keep(watch), ) }) @@ -851,29 +996,62 @@ export function createHarnessHandler

( request.headers.get('Last-Event-ID') ?? url.searchParams.get('from') ?? undefined + const refused = versionOf(url) + if (refused) return versionRefused(refused, requestId) + const hello: HostFrame = { + type: 'harness.hello', + v: HARNESS_PROTOCOL_VERSION, + threadId, + } return sseResponse( async function* (signal) { + yield { data: JSON.stringify(hello) } const events = view.events({ ...(from ? { from } : {}), signal }) + // Each event goes out as an AG-UI wire chunk. for await (const entry of events) { - const frame: HostFrame = { type: 'harness.event', ...entry } - yield { id: entry.cursor, data: frame } + yield { id: entry.cursor, data: wireFrame(entry) } } }, + requestId, + stream, closeOnRevoke(principal, threadId), ) } if (request.method === 'GET' && route === 'host-events') { - const mayWatch = watchedEvents(principal) - return sseResponse(async function* (signal) { - for await (const event of host.events({ signal })) { - if (await mayWatch(event)) yield { data: event } - } + const refused = versionOf(url) + if (refused) return versionRefused(refused, requestId) + const { accepts, mayWatch } = watchedEvents(principal) + // The events of the principal's tenant only, from the host event bus. + // A host that `createHarnessHost` did not make has `host.events()`. + const readTenant = tenantEventsOf(host) + // The hello of `host-events` names no thread. + const hello = JSON.stringify({ + type: 'harness.hello', + v: HARNESS_PROTOCOL_VERSION, }) + return sseResponse( + async function* (signal) { + yield { data: hello } + const events = readTenant + ? readTenant({ + tenantId: principal?.tenantId, + filter: accepts, + signal, + }) + : host.events({ signal }) + for await (const event of events) { + if (!readTenant && !accepts(event)) continue + if (await mayWatch(event)) yield { data: JSON.stringify(event) } + } + }, + requestId, + stream, + ) } if (request.method === 'POST' && route === 'control') { - const body: unknown = await request.json() + const body = await readJson(request) if (typeof body !== 'object' || body === null || !('threadId' in body)) { return json({ error: 'threadId is required' }, 400) } @@ -882,8 +1060,10 @@ export function createHarnessHandler

( return json({ error: 'threadId is required' }, 400) const input = parseHarnessInput('input' in body ? body.input : undefined) const session = await openFor(principal, threadId) + const traceparent = request.headers.get('traceparent') + const traced = traceparent ? { traceparent } : {} return withSession(session, async () => - json(await applyInput({ harness, session, input, ...user })), + json(await applyInput({ harness, session, input, ...user, ...traced })), ) } @@ -914,7 +1094,7 @@ export function createHarnessHandler

( return json({ error: 'subagentRunId is required' }, 400) const view = await viewFor(principal, threadId) if (!view) return json({ error: 'not found' }, 404) - return json(await view.agentTranscript(subagentRunId)) + return json(await view.agentTranscript({ subagentRunId })) } if (request.method === 'GET' && route === 'snapshot') { @@ -957,7 +1137,7 @@ export function createHarnessHandler

( } if (request.method === 'POST' && route === 'sessions') { - const op = parseSessionOp(await request.json()) + const op = parseSessionOp(await readJson(request)) if (!op) return json({ error: 'not a valid sessions op' }, 400) const target = { threadId: op.threadId, ...user } const entry = await host.sessions.get(target) @@ -1047,7 +1227,7 @@ export function createHarnessHandler

( const view = await viewFor(principal, threadId) if (!view) return json({ error: 'not found' }, 404) if (route === 'media') return mediaResponse(request, view, id) - if (!(await view.getMedia(id))) { + if (!(await view.getMedia({ id }))) { return json({ error: 'not found' }, 404) } const expiresAt = Date.now() + SIGNED_URL_TTL_MS @@ -1069,7 +1249,7 @@ export function createHarnessHandler

( return json({ error: 'not found' }, 404) } - return async (request: Request) => { + const answer = async (request: Request, requestId: string) => { const url = new URL(request.url) const isSignedMedia = request.method === 'GET' && @@ -1080,11 +1260,19 @@ export function createHarnessHandler

( if (isSignedMedia) return await serveSigned(request, url) const user = await usersOf(request) if (!user) return json({ error: 'unauthorized' }, 401) - return await handle(request, user.principal) + return await handle(request, user.principal, requestId) } catch (error) { return failed(error) } } + + return async (request: Request) => { + const requestId = requestIdFor(request) + const response = await answer(request, requestId) + // Every answer names its request, so a client can find it in the logs. + response.headers.set('X-Request-Id', requestId) + return response + } } /** @@ -1146,8 +1334,17 @@ export interface HarnessSocketOptions

{ request: Request /** For a host without `access` only. See `HarnessHandlerOptions`. */ openAccess?: OpenAccess

+ /** + * The values of the socket: the heartbeat, the slow-client close, and the + * buffer size. A value that you do not set keeps its default. See + * `HarnessStreamOptions`. + */ + stream?: HarnessStreamOptions } +/** The close code of a socket after an `unsupported_version` error frame. */ +const UNSUPPORTED_VERSION_CLOSE_CODE = 4400 + /** The text of an error frame: a refused action is `forbidden`. */ const errorMessage = (error: unknown) => error instanceof AccessDeniedError @@ -1168,8 +1365,17 @@ const errorMessage = (error: unknown) => * session. The first input opens the session, and the socket holds it until * it closes. * + * `stream` sets the values of the socket (see `HarnessStreamOptions`). + * + * A frame over `media.maxBodyBytes` gets an error frame, and the socket + * closes with 1009. Your WebSocket server has read the whole frame by then, + * so this check saves only the parse. The memory cap is the frame limit of + * your server, for example `maxPayload` of `ws`. Set it to + * `media.maxBodyBytes` or lower. + * * Throws at the start for a host without `access` and without `openAccess`, - * and for a host with both. + * for a host with both, and for a `stream` value that is not a finite number + * above 0. * * @example * ```ts @@ -1181,7 +1387,17 @@ export function handleHarnessSocket

( ) { const { host, harness, socket, request } = options const usersOf = requestUsers('handleHarnessSocket', host, options.openAccess) + const values = streamValues('handleHarnessSocket', options.stream) const reader = new AbortController() + // The id of the upgrade request, with the rule of `requestIdFor`. An error + // frame of the socket itself, or of a frame with no `requestId`, names it. + const socketId = requestIdFor(request) + const sender = socketSender({ + socket, + requestId: socketId, + signal: reader.signal, + stream: values, + }) const who = usersOf(request).catch(() => null) // Set by the first subscribe, before its awaits, so a second subscribe is // refused and an input sent right after the subscribe waits for it. It @@ -1235,10 +1451,12 @@ export function handleHarnessSocket

( from: Cursor | undefined, ) => { let cursor = from + // Each event goes out as an AG-UI wire chunk. The send waits while the + // socket buffer is full, so a slow client stops the read of the log. const forward = async (events: AsyncIterable) => { for await (const entry of events) { cursor = entry.cursor - send({ type: 'harness.event', ...entry }) + if (!(await send(wireFrame(entry)))) return } } const readOptions = () => ({ @@ -1254,30 +1472,84 @@ export function handleHarnessSocket

( if (!session) return await forward(session.events(readOptions())) } - const send = (frame: HostFrame) => { + /** Send a frame in order. False when the socket is gone. */ + const send = async (frame: HostFrame | string) => { + if (await sender(frame)) return true + reader.abort() + return false + } + const fail = (code: HarnessErrorCode, message: string, requestId: string) => + void send({ type: 'harness.error', message, code, requestId }) + /** + * Send a last error frame and close. It goes out at once, before the + * frames that wait in the sender: the socket closes next. + */ + const closeWith = ( + frame: Extract, + closeCode: number, + ) => { try { socket.send(JSON.stringify(frame)) } catch { - reader.abort() + // The socket is gone already. } - } - const refuse = (message: 'unauthorized' | 'forbidden') => { - send({ type: 'harness.error', message }) - socket.close(message === 'unauthorized' ? 4401 : 4403, message) + socket.close(closeCode, frame.code) reader.abort() } + const refuse = (message: 'unauthorized' | 'forbidden', requestId: string) => + closeWith( + { type: 'harness.error', message, code: message, requestId }, + message === 'unauthorized' ? 4401 : 4403, + ) socket.addEventListener('close', () => reader.abort()) socket.addEventListener('error', () => reader.abort()) + const frameLimit = maxBodyBytesOf(harness.media) + // The frames of the socket run one after another, as the stdio worker + // runs them: an input with media to store does not let a later input + // reach the session first. + let frames = Promise.resolve() socket.addEventListener('message', (message) => { - void (async () => { + const data = String(message.data) + // A frame over `media.maxBodyBytes` is not parsed, so its error frame + // names the socket. The socket closes with 1009 (message too big). The + // server read the frame already: its own frame limit caps the memory. + if (isOverLimit(data, frameLimit)) { + closeWith( + { + type: 'harness.error', + message: `The frame is bigger than the limit of ${frameLimit} bytes.`, + code: 'invalid_frame', + requestId: socketId, + }, + 1009, + ) + return + } + // Read before the full parse, so a frame with a bad `input` still names + // its request. + const requestId = requestIdOf(data) ?? socketId + frames = frames.then(async () => { try { const user = await who - if (!user) return refuse('unauthorized') + if (!user) return refuse('unauthorized', requestId) const { principal } = user - const frame = parseControlFrame(String(message.data)) + let frame: ControlFrame + try { + frame = parseControlFrame(data) + } catch (error) { + return fail('invalid_frame', errorMessage(error), requestId) + } if (frame.type === 'harness.subscribe') { - if (subscribed) throw new Error('Already subscribed.') + if (subscribed) + return fail('invalid_frame', 'Already subscribed.', requestId) + const refused = versionError(frame.v) + if (refused) { + return closeWith( + { type: 'harness.error', requestId, ...refused }, + UNSUPPORTED_VERSION_CLOSE_CODE, + ) + } const { threadId } = frame subscribed = observeToRead(host, harness, { threadId, @@ -1290,22 +1562,25 @@ export function handleHarnessSocket

( }, ) const subscription = await subscribed - if (!subscription) return refuse('forbidden') + if (!subscription) return refuse('forbidden', requestId) watchRevoke( host, { principal, threadId, signal: reader.signal }, - () => refuse('forbidden'), + () => refuse('forbidden', requestId), ) - send({ + void send({ type: 'harness.hello', v: HARNESS_PROTOCOL_VERSION, - threadId: frame.threadId, + threadId, }) void stream(subscription.view, frame.from) + // The close waits for this frame: it goes out first. .catch((error: unknown) => send({ type: 'harness.error', message: errorMessage(scrubbedBy(holding, error)), + code: 'internal', + requestId, }), ) // The feed ended while the socket is open: close it, so the @@ -1316,13 +1591,18 @@ export function handleHarnessSocket

( return } const subscription = await subscribed - if (!subscription) throw new Error('Send harness.subscribe first.') + if (!subscription) + return fail( + 'invalid_frame', + 'Send harness.subscribe first.', + requestId, + ) if (frame.type === 'harness.snapshot') { // The held session, else the view: a snapshot opens nothing. const live = holding?.snapshot() const snapshot = live ?? (await subscription.view?.snapshot()) if (!snapshot) throw new Error('not found') - send({ type: 'harness.snapshot', snapshot }) + void send({ type: 'harness.snapshot', snapshot }) return } const session = await hold(subscription.threadId, principal) @@ -1331,20 +1611,20 @@ export function handleHarnessSocket

( session, input: frame.input, ...(principal ? { principal } : {}), + ...(frame.traceparent ? { traceparent: frame.traceparent } : {}), }) - send({ + void send({ type: 'harness.receipt', requestId: frame.requestId, ...receipt, }) } catch (error) { + const code = + error instanceof AccessDeniedError ? 'forbidden' : 'internal' // The held session scrubs its secrets. Before the first input, // no session is open, so there is none to scrub. - send({ - type: 'harness.error', - message: errorMessage(scrubbedBy(holding, error)), - }) + fail(code, errorMessage(scrubbedBy(holding, error)), requestId) } - })() + }) }) } diff --git a/packages/ai-harness/src/index.ts b/packages/ai-harness/src/index.ts index 7a162eaee6..93fbf8d415 100644 --- a/packages/ai-harness/src/index.ts +++ b/packages/ai-harness/src/index.ts @@ -166,6 +166,7 @@ export type { export type { ObservedSession } from './session/observed-session' export type { SessionUsage, UsageCounts } from './usage' export type { OnHarnessEvent } from './feed' +export type { ObservabilityOptions, OperationReport } from './observability' export type { FinishContext, @@ -191,6 +192,7 @@ export type { EventGap, ForkPoint, HandoffEnvelope, + HarnessEventValues, HarnessInput, HarnessRunContext, InputSettlement, diff --git a/packages/ai-harness/src/inline-media.ts b/packages/ai-harness/src/inline-media.ts new file mode 100644 index 0000000000..5c6f368eae --- /dev/null +++ b/packages/ai-harness/src/inline-media.ts @@ -0,0 +1,239 @@ +// The checks on media that a client sends inside an input. Inline bytes go +// to the media store, so the transcript, the log, and every viewer get a +// small `harness-media:` reference. External URLs are refused unless the +// harness allows them. +import { MEDIA_URL_PREFIX, kindOf, mediaPart } from './media-ref' +import { MediaError } from './media' +import type { ContentPart } from '@tanstack/ai' +import type { MediaOptions } from './media' +import type { MediaRecord, MediaUpload } from './types' + +/** The default of `media.maxBodyBytes`: 32 MiB. */ +export const DEFAULT_MAX_BODY_BYTES = 32 * 1024 * 1024 + +/** + * Why a client content part was refused. It is the `reason` of the receipt. + * + * - `media_too_large`: the bytes are over `media.maxBytes`. + * - `media_not_allowed`: a MIME type with no media kind, or a kind not in + * `media.kinds`. + * - `media_invalid`: the data is not base64. + * - `media_url_not_allowed`: a URL that `media.externalUrls` refuses, a + * URL scheme other than `http:`, `https:`, `data:`, and `harness-media:`, + * or a provider file handle without `media.allowProviderFiles`. + */ +export type InlineMediaRefusal = + | 'media_too_large' + | 'media_not_allowed' + | 'media_invalid' + | 'media_url_not_allowed' + +/** What `storeInlineMedia` needs: the media options and a store. */ +export interface InlineMediaContext { + options: MediaOptions | undefined + /** Store the bytes, as `session.putMedia` does. */ + putMedia: (upload: MediaUpload) => Promise +} + +type MediaContentPart = Exclude + +/** A checked part, with the bytes that it must store, if any. */ +type CheckedPart = + | { part: ContentPart; upload?: undefined } + | { part: MediaContentPart; upload: MediaUpload } + +const BASE64 = /^[A-Za-z0-9+/]*={0,2}$/ +// `data:[;params];base64,`. Only base64 data URLs carry bytes. +const DATA_URL = /^data:([^,]*?);base64,(.*)$/s + +/** A refused part. `reason` goes to the receipt. */ +function refuse(reason: InlineMediaRefusal) { + return { refused: reason } +} + +/** The bytes of base64 `value`, or `undefined` when it is not base64. */ +function fromBase64(value: string) { + const clean = value.replace(/\s/g, '') + if (clean.length % 4 === 1 || !BASE64.test(clean)) return undefined + const NodeBuffer = globalThis.Buffer + if (typeof NodeBuffer === 'function') { + return new Uint8Array(NodeBuffer.from(clean, 'base64')) + } + return Uint8Array.from(atob(clean), (char) => char.charCodeAt(0)) +} + +/** The number of bytes that base64 `value` holds, without decoding it. */ +function base64Size(value: string) { + const clean = value.replace(/\s/g, '') + const padding = clean.endsWith('==') ? 2 : clean.endsWith('=') ? 1 : 0 + return Math.floor((clean.length * 3) / 4) - padding +} + +/** True when `media.externalUrls` lets a part point to `url`. */ +function isUrlAllowed(url: URL, options: MediaOptions | undefined) { + const isWeb = url.protocol === 'http:' || url.protocol === 'https:' + if (!isWeb) return false + const allow = options?.externalUrls ?? false + return typeof allow === 'function' ? allow(url) : allow +} + +/** `part`, checked, with the bytes it must store, or why it was refused. */ +function checkBytes( + part: MediaContentPart, + data: { value: string; mimeType: string }, + options: MediaOptions | undefined, +) { + const kind = kindOf(data.mimeType) + const isKindAllowed = + kind !== undefined && (options?.kinds ?? [kind]).includes(kind) + if (!isKindAllowed) return refuse('media_not_allowed') + const isTooLarge = + options?.maxBytes !== undefined && base64Size(data.value) > options.maxBytes + if (isTooLarge) return refuse('media_too_large') + const bytes = fromBase64(data.value) + if (bytes === undefined) return refuse('media_invalid') + const checked: CheckedPart = { + part, + upload: { body: bytes, mimeType: data.mimeType, name: kind }, + } + return checked +} + +/** One part, checked, or the reason it was refused. Stores nothing. */ +function checkPart( + part: ContentPart, + options: MediaOptions | undefined, +): CheckedPart | { refused: InlineMediaRefusal } { + if (part.type === 'text') return { part } + const { source } = part + if (source.type === 'data') return checkBytes(part, source, options) + // A provider file handle is read with the server's provider key. + if (source.type === 'file') { + return options?.allowProviderFiles === true + ? { part } + : refuse('media_url_not_allowed') + } + if (source.value.startsWith(MEDIA_URL_PREFIX)) return { part } + const dataUrl = DATA_URL.exec(source.value) + if (dataUrl) { + const [, mimeType = '', value = ''] = dataUrl + return checkBytes(part, { value, mimeType }, options) + } + const url = URL.canParse(source.value) ? new URL(source.value) : undefined + if (url === undefined || !isUrlAllowed(url, options)) { + return refuse('media_url_not_allowed') + } + return { part } +} + +/** + * Check the content parts that a client sent, at the trust boundary. Inline + * bytes (a `data` source, or a base64 `data:` URL) go to the media store + * with `putMedia`, and the part becomes a `harness-media:` reference. So the + * bytes are checked against `media.kinds` and `media.maxBytes`, and the model + * check of `harness-media:` parts runs on them. An `http:` or `https:` URL + * stays only when `media.externalUrls` allows it. A provider file handle + * stays only when `media.allowProviderFiles` is `true`. Other URL schemes are + * refused. Text and `harness-media:` parts stay as they are. + * + * Resolves to `{ parts }`, or to `{ refused }` with the reason of the first + * part that was refused. Every part is checked before any part is stored, so + * a refused input stores nothing. + * + * @example + * ```ts + * const checked = await storeInlineMedia(parts, { + * options: harness.media, + * putMedia: (upload) => session.putMedia(upload), + * }) + * if ('refused' in checked) return reject(checked.refused) + * session.send(checked.parts) + * ``` + */ +export async function storeInlineMedia( + parts: ReadonlyArray, + context: InlineMediaContext, +) { + const pending: Array = [] + for (const part of parts) { + const result = checkPart(part, context.options) + if ('refused' in result) return result + pending.push(result) + } + const checked: Array = [] + for (const entry of pending) { + if (entry.upload === undefined) { + checked.push(entry.part) + continue + } + try { + const record = await context.putMedia(entry.upload) + // Other fields, like provider metadata, stay on the part. + const stored: ContentPart = { ...entry.part, ...mediaPart(record) } + checked.push(stored) + } catch (error) { + // ponytail: the store refuses only what the checks above let through + // when its own limits differ (its 100 MB default `maxBytes`). Parts + // stored before such a refusal stay in the store. + if (!(error instanceof MediaError)) throw error + return refuse( + error.status === 413 ? 'media_too_large' : 'media_not_allowed', + ) + } + } + return { parts: checked } +} + +/** + * The largest request body or socket frame, in bytes, that a transport + * reads: `media.maxBodyBytes`, or 32 MiB. A socket checks a frame after the + * WebSocket server read it, so this saves only the parse. + * + * @example + * ```ts + * if (frame.length > maxBodyBytesOf(harness.media)) socket.close(1009) + * ``` + */ +export function maxBodyBytesOf(options: MediaOptions | undefined) { + return options?.maxBodyBytes ?? DEFAULT_MAX_BODY_BYTES +} + +/** + * The body of `request` as text, read up to `media.maxBodyBytes`. Throws a + * `MediaError` with 413 when the `Content-Length` or the bytes read are over + * the limit. It stops reading at the limit, so a large body is not buffered. + * + * @example + * ```ts + * const body: unknown = JSON.parse(await readBodyText(request, harness.media)) + * ``` + */ +export async function readBodyText( + request: Request, + options: MediaOptions | undefined, +) { + const limit = maxBodyBytesOf(options) + const tooLarge = () => + new MediaError(413, `The body is bigger than the limit of ${limit} bytes.`) + const declared = Number(request.headers.get('content-length') ?? 0) + if (declared > limit) throw tooLarge() + if (request.body === null) return '' + const reader = request.body.getReader() + const chunks: Array = [] + let size = 0 + for (;;) { + const { done, value } = await reader.read() + if (done) break + size += value.byteLength + if (size > limit) { + await reader.cancel() + throw tooLarge() + } + chunks.push(value) + } + const decoder = new TextDecoder() + return ( + chunks.map((chunk) => decoder.decode(chunk, { stream: true })).join('') + + decoder.decode() + ) +} diff --git a/packages/ai-harness/src/log.ts b/packages/ai-harness/src/log.ts index e184500036..4c1097db6a 100644 --- a/packages/ai-harness/src/log.ts +++ b/packages/ai-harness/src/log.ts @@ -5,8 +5,9 @@ import { OwnershipLostError, UNOWNED_EPOCH, } from '@tanstack/ai-persistence' +import { gapEvent } from './feed' import { LogOwnedElsewhereError } from './ownership' -import { HARNESS_EVENTS } from './types' +import { EventRing } from './ring' import { THREAD_RECORD_TYPES, applyRunNode, @@ -54,7 +55,7 @@ import type { } from './log-records' import type { Cursor, - EventGap, + HarnessEventValues, HarnessInput, InputSettlement, SettlementOutcome, @@ -66,7 +67,11 @@ import type { SessionUsage, UsageCall, UsageCounts } from './usage' /** The records the harness writes to a session log. */ export type HarnessRecord = - | { type: 'harness.event'; operationId: string; event: StreamChunk } + | { + type: 'harness.event' + operationId: string + event: StreamChunk + } | { type: 'harness.transcript'; keep: number; add: Array } | { type: 'harness.input' @@ -74,6 +79,8 @@ export type HarnessRecord = input: HarnessInput principal?: { id: string; tenantId?: string } at: number + /** The W3C `traceparent` of the caller. */ + traceparent?: string } | { type: 'harness.input.applied' @@ -84,6 +91,12 @@ export type HarnessRecord = } | { type: 'harness.input.joined'; inputId: string; into: string } | { type: 'harness.input.rejected'; inputId: string; reason: string } + /** + * The owner refused an offer from the input lane, and the log has another + * input with its id. It does not change that input. Only the host that + * offered it ends its wait with the refusal. + */ + | { type: 'harness.offer.rejected'; inputId: string; reason: string } | { type: 'harness.input.abort'; inputId: string } | { type: 'harness.turn.retry' @@ -115,13 +128,28 @@ export type HarnessRecord = value: unknown } | ({ type: 'harness.usage' } & UsageCall) - /** A revert that stands, or `null` when it ended. */ - | { type: 'harness.revert'; revert: RevertState | null } + /** + * A revert that stands, or `null` when it ended. `principal` is who + * reverted, on the record of `session.revert`. + */ + | { + type: 'harness.revert' + revert: RevertState | null + principal?: { id: string; tenantId?: string } + } /** * A background job started, or it ended. `backendId` is the id of its * process in the workspace backend, so a later host can stop it. + * `principal` is who moved a running call to the background + * (`session.background`). */ - | { type: 'harness.job'; jobId: string; ended: boolean; backendId?: string } + | { + type: 'harness.job' + jobId: string + ended: boolean + backendId?: string + principal?: { id: string; tenantId?: string } + } /** A note for the transcript. `noteId` is the id of its message. */ | { type: 'harness.note'; noteId: string; note: string } /** Thread state: settings, config, plugin state, and the rest. */ @@ -168,6 +196,7 @@ const HARNESS_RECORD_TYPES = new Set([ 'harness.input.applied', 'harness.input.joined', 'harness.input.rejected', + 'harness.offer.rejected', 'harness.input.abort', 'harness.turn.retry', 'harness.input.delivery', @@ -244,6 +273,11 @@ export interface InputState { input: HarnessInput principal?: { id: string; tenantId?: string } at: number + /** + * The position of the input record. The events of its run come after it. + * Unknown in a fold snapshot from before this field. + */ + seq?: number status: 'pending' | 'applied' | 'joined' | 'rejected' | 'settled' operationId?: string attempt: number @@ -259,6 +293,8 @@ export interface InputState { appliedAt?: number settlement?: InputSettlement reason?: string + /** The W3C `traceparent` of the caller. The span of its turn links it. */ + traceparent?: string } /** The fold of a session log. */ @@ -295,10 +331,73 @@ export interface LogState { jobs: Map /** Notes that are not in the transcript yet, by message id. */ notes: Map + /** + * The last host records of the types that the plugins keep (their + * `keepRecords`), oldest first. The snapshot keeps them, so a truncate of + * the log does not drop what `session.revert` reads. + */ + keptRecords: Array /** The thread state that the thread-state records fold to. */ thread: ThreadStateFold } +/** + * The host records that a fold keeps, by type, with how many of each it + * keeps. The `keepRecords` of the plugins give it. + */ +export type KeptRecords = ReadonlyMap + +/** The count of each type in a `keptRecords` list, made at its first use. */ +const keptCounts = new WeakMap, Map>() + +/** Keep `entry` in `state` when `keep` keeps its type. */ +function keepHostRecord( + state: LogState, + entry: LogEntry, + keep: KeptRecords | undefined, +) { + const { type } = entry.record + const limit = keep?.get(type) + if (limit === undefined) return + const kept = state.keptRecords + let counts = keptCounts.get(kept) + if (!counts) { + counts = new Map() + for (const { record } of kept) { + counts.set(record.type, (counts.get(record.type) ?? 0) + 1) + } + keptCounts.set(kept, counts) + } + kept.push(entry) + let count = (counts.get(type) ?? 0) + 1 + // ponytail: the splice is O(n), only when a type is full. A list for + // each type if many types are kept. + for (; count > limit; count--) { + kept.splice( + kept.findIndex((each) => each.record.type === type), + 1, + ) + } + counts.set(type, count) +} + +/** + * The cursor before the first event of run `operationId`: before its + * first input record. A run that answers go on has more inputs, so the + * oldest one counts. `undefined` when the fold does not know it. + */ +export function runStartOf( + state: LogState, + operationId: string, +): Cursor | undefined { + let first: number | undefined + for (const input of state.inputs.values()) { + if (input.operationId !== operationId || input.seq === undefined) continue + if (first === undefined || input.seq < first) first = input.seq + } + return first === undefined ? undefined : String(first - 1) +} + export function emptyLogState() { const state: LogState = { seq: 0, @@ -312,6 +411,7 @@ export function emptyLogState() { runUsage: new Map(), jobs: new Map(), notes: new Map(), + keptRecords: [], thread: emptyThreadState(), } return state @@ -369,7 +469,9 @@ export function commonPrefix( ) { const length = Math.min(a.length, b.length) let index = 0 - while (index < length && sameJson(a[index], b[index])) index += 1 + // The same message object is the same JSON, so it skips the stringify. + const isSame = (i: number) => a[i] === b[i] || sameJson(a[i], b[i]) + while (index < length && isSame(index)) index += 1 return index } @@ -442,10 +544,12 @@ export function foldEntry( state: LogState, entry: LogEntry, project?: ProjectRecord, + keep?: KeptRecords, ) { state.seq = entry.seq const { record } = entry if (!record.type.startsWith('harness.')) { + keepHostRecord(state, entry, keep) const next = project ? hostFold( entry.seq, @@ -460,6 +564,7 @@ export function foldEntry( if (!isHarnessRecord(record)) return switch (record.type) { case 'harness.event': + case 'harness.offer.rejected': return case 'harness.transcript': state.messages = [ @@ -485,7 +590,9 @@ export function foldEntry( inputId: record.inputId, input: record.input, ...(record.principal ? { principal: record.principal } : {}), + ...(record.traceparent ? { traceparent: record.traceparent } : {}), at: record.at, + seq: entry.seq, status: 'pending', attempt: 0, abortRequested: false, @@ -713,6 +820,7 @@ interface FoldOptions { rootThread: string project?: ProjectRecord reduce?: ReduceOptions + keep?: KeptRecords } /** Fold one entry of a shared log. */ @@ -739,6 +847,7 @@ export function foldLogEntry( sessionOf(state, threadOf(record, options.rootThread)), entry, options.project, + options.keep, ) } @@ -879,6 +988,7 @@ function serialize(state: SharedLogState, versions: SnapshotVersions) { backendId === undefined ? jobId : [jobId, backendId], ), notes: [...session.notes.entries()], + keptRecords: session.keptRecords, thread: serializeThread(session.thread), }, ]), @@ -919,8 +1029,11 @@ function parseSession(value: unknown, seq: number) { : [], ), steps: new Map(steps as Array<[string, unknown]>), - // A state with no usage totals gets an empty value. - usage: isSessionUsage(value.usage) ? value.usage : emptyUsage(), + // A state with no usage totals gets an empty value, and totals from + // before `bySource` get an empty map. + usage: isSessionUsage(value.usage) + ? { bySource: {}, ...value.usage } + : emptyUsage(), // A state with no run sums gets an empty value. runUsage: new Map( Array.isArray(value.runUsage) @@ -942,6 +1055,10 @@ function parseSession(value: unknown, seq: number) { ? (value.notes as Array<[string, string]>) : [], ), + // A state with no kept records gets an empty value. + keptRecords: Array.isArray(value.keptRecords) + ? (value.keptRecords as Array) + : [], thread, } return session @@ -1014,8 +1131,10 @@ export async function loadLogState(options: { rootThread: string project?: ProjectOptions reduce?: ReduceOptions + /** The host records to keep. See {@link KeptRecords}. */ + keep?: KeptRecords }) { - const { store, logId, rootThread, project, reduce } = options + const { store, logId, rootThread, project, reduce, keep } = options // A snapshot that does not load only makes the fold longer. const snapshot = await store.getSnapshot({ logId }).catch(() => null) const state = @@ -1025,6 +1144,7 @@ export async function loadLogState(options: { rootThread, ...(project ? { project: project.record } : {}), ...(reduce ? { reduce } : {}), + ...(keep ? { keep } : {}), } await readAll(store, logId, state, (entry) => foldLogEntry(state, entry, fold), @@ -1250,6 +1370,9 @@ export interface LogWriter extends EventFeed { rest: () => Promise } +/** A new owner of a log. See `LogViewHooks.onOwned`. */ +export type OwnershipChange = HarnessEventValues['harness.ownership.changed'] + /** What a session gives its view of a shared log. */ export interface LogViewHooks { /** A write failed, so the log stopped. */ @@ -1273,6 +1396,22 @@ export interface LogViewHooks { * host offered still waits. */ onOffers: () => void + /** + * An event of the session that another host wrote, after this log folded + * it. Own events never come here: the session saw them at publish. + */ + onRemoteEvent?: (entry: SessionEvent) => void + /** + * This log folded records that another host wrote: on a notification, or + * when this host acquires the log. It comes before the `onRemoteEvent` + * calls, so a session reads the new fold first. + */ + onRemoteFold?: () => void + /** + * This host became the owner of the log, before its first append. Only + * one session of the log gets it: the root session when it is open. + */ + onOwned?: (change: OwnershipChange) => void /** * Scrub the secrets of the session from a value. Each record that the * view appends, stages, commits, or offers goes through it first. It @@ -1359,6 +1498,8 @@ export class SharedLog { * before its first append. `undefined` while it does not own the log. */ private epoch: number | undefined + /** The last epoch that this host had, also after it gave the log up. */ + private heldEpoch: number | undefined /** The open views, by thread. */ private readonly views = new Map< string, @@ -1388,19 +1529,13 @@ export class SharedLog { private chain: Promise = Promise.resolve() private failure: unknown /** - * The newest events, like SessionFeed: at most `eventBytes` of them (the - * JSON length of each). Older cursors page from the store. + * The newest events: at most `eventBytes` of them (the JSON length of + * each). Older cursors page from the store. */ - private readonly tail: Array<{ - thread: string - entry: SessionEvent - size: number - }> = [] - /** The sum of the sizes in `tail`. */ - private tailBytes = 0 + private tail: EventRing + /** While `catchUp` folds: the events that other writers appended. */ + private remote: Array<{ thread: string; entry: SessionEvent }> | undefined private readonly eventBytes: number - /** Every event after this position is in `tail`. */ - private tailFrom: number private readonly waiters = new Set<() => void>() /** The open readers. Like the views, they keep the log open. */ private readers = 0 @@ -1413,6 +1548,13 @@ export class SharedLog { */ private hasLeftWork = false private writing = false + /** + * The store notifications that came while this host appended. The store + * calls each instance once for each append, also for its own. A second + * call is another host's offer, so the write looks at the offers after + * its append. + */ + private callsInWrite = 0 private readonly unsubscribe: () => void private readonly stopLost: () => void private readonly stopWake: () => void @@ -1433,6 +1575,8 @@ export class SharedLog { eventBytes: number project?: ProjectOptions reduce?: ReduceOptions + /** The host records to keep. See {@link KeptRecords}. */ + keep?: KeptRecords /** When the owner writes a snapshot. */ snapshot?: SnapshotTiming /** Which records the log keeps after a snapshot. Default `'all'`. */ @@ -1458,21 +1602,25 @@ export class SharedLog { rootThread: options.rootThread, ...(options.project ? { project: options.project.record } : {}), ...(options.reduce ? { reduce: options.reduce } : {}), + ...(options.keep ? { keep: options.keep } : {}), } this.versions = versionsOf(options.project, options.reduce) - this.tailFrom = options.state.seq + this.tail = new EventRing({ + maxBytes: options.eventBytes, + from: options.state.seq, + }) // Records from another writer join the fold. A call during this - // writer's own append is its own records. + // writer's own append can be its own records or another host's offer: + // the write looks at the offers after its append. this.unsubscribe = options.store.subscribe({ logId: options.logId, listener: () => { - if (this.writing) return - this.catchUpLater() - // The call can be a new offer. The owner applies it. A host that - // offered applies its own offers when no live host owns the log. - if (this.epoch !== undefined || this.offeredHere.size > 0) { - this.notifyOffers() + if (this.writing) { + this.callsInWrite += 1 + return } + this.catchUpLater() + this.onOffersChanged() }, }) this.stopLost = options.ownership.onLost(options.logId, (error) => @@ -1707,6 +1855,16 @@ export class SharedLog { await this.ownership.release(this.logId, epoch) } + /** + * A store notification can be a new offer. The owner applies it. A host + * that offered applies its own offers when no live host owns the log. + */ + private onOffersChanged() { + if (this.epoch !== undefined || this.offeredHere.size > 0) { + this.notifyOffers() + } + } + /** Fold the records of another writer, after the queued writes. */ private catchUpLater() { // ponytail: a read that fails leaves the fold behind until the next @@ -1951,9 +2109,34 @@ export class SharedLog { // The new owner applies the inputs that wait for an owner. Later, so the // write that acquired the log goes first. queueMicrotask(() => this.notifyOffers()) + // The next epoch after the one that this host had: no other host owned + // the log in between, so the owner did not change. + const held = this.heldEpoch + this.heldEpoch = epoch + if (held === undefined || epoch > held + 1) { + // The store says whether the old owner stopped: its ownership expired + // with no release. + const reason = this.ownership.tookOver(this.logId) + ? 'taken_over' + : 'acquired' + queueMicrotask(() => this.notifyOwned({ epoch, reason })) + } return epoch } + /** + * Call `onOwned` of the root session of the log, or of another open + * session when the root is not open here. One call for each change. + */ + private notifyOwned(change: OwnershipChange) { + const root = this.views.get(this.rootThread) + const view = + root && !root.fenced + ? root + : [...this.views.values()].find((each) => !each.fenced) + view?.hooks.onOwned?.(change) + } + /** Pin the log to `hostId`, with the epoch of this host. */ private async pin(hostId: string | null) { if (this.failure !== undefined) throw this.failure @@ -1993,6 +2176,7 @@ export class SharedLog { const epoch = await this.own() const seq = this.state.seq + 1 this.writing = true + this.callsInWrite = 0 try { await this.appendOnce({ seq, epoch, batch, consumes }) } catch (error) { @@ -2002,6 +2186,11 @@ export class SharedLog { throw error } finally { this.writing = false + // Only the owner appends, so a call is this append or an offer. + // Records of a new owner reach this host when its next renewal or + // write fails. A store that calls after the append calls outside of + // the write, and the listener looks at the offers then. + if (this.callsInWrite > 1) this.onOffersChanged() } batch.forEach((record, index) => this.apply({ seq: seq + index, record })) this.wake() @@ -2182,6 +2371,10 @@ export class SharedLog { private async catchUp() { if (this.failure !== undefined) return const before = this.state.seq + // Only this fold reads the records of other writers. Own records fold + // in `write`, so an own event never reaches `onRemoteEvent`. + const remote: Array<{ thread: string; entry: SessionEvent }> = [] + this.remote = remote try { await readAll(this.store, this.logId, this.state, (entry) => this.apply(entry), @@ -2189,12 +2382,29 @@ export class SharedLog { } catch (error) { if (!(error instanceof LogTruncatedError)) throw error await this.foldFromSnapshot(error) + } finally { + this.remote = undefined } if (this.state.seq === before) return if (this.epoch !== undefined) { this.lose(new OwnershipLostError(this.logId, this.epoch)) } this.wake() + for (const view of [...this.views.values()]) { + try { + view.hooks.onRemoteFold?.() + } catch { + // A hook never stops the fold of the log. + } + } + // After the whole fold, so a hook reads the state after the events. + for (const { thread, entry } of remote) { + try { + this.views.get(thread)?.hooks.onRemoteEvent?.(entry) + } catch { + // A hook never stops the fold of the log. + } + } } /** @@ -2215,9 +2425,7 @@ export class SharedLog { this.state.seq = loaded.seq this.state.reduced = loaded.reduced this.snapshotSeq = loaded.seq - this.tail.length = 0 - this.tailBytes = 0 - this.tailFrom = loaded.seq + this.tail = new EventRing({ maxBytes: this.eventBytes, from: loaded.seq }) await readAll(this.store, this.logId, this.state, (entry) => this.apply(entry), ) @@ -2237,19 +2445,13 @@ export class SharedLog { this.resulted.delete(`${thread}\u0000${message.toolCallId}`) } } - const event = sessionEventOf(entry) + const event = sessionEventOf(entry, this.rootThread) if (!event) return + this.remote?.push({ thread, entry: event }) // The fold has JSON values only, so this cannot throw. const size = JSON.stringify(event.event).length + // The ring keeps the newest event, so a live reader gets it from memory. this.tail.push({ thread, entry: event, size }) - this.tailBytes += size - // The newest event stays, so a live reader gets it from memory. - while (this.tailBytes > this.eventBytes && this.tail.length > 1) { - const dropped = this.tail.shift() - if (!dropped) break - this.tailBytes -= dropped.size - this.tailFrom = Number(dropped.entry.cursor) - } } // =========================== @@ -2275,17 +2477,14 @@ export class SharedLog { const accepts = (entry: SessionEvent) => (filter ? filter(entry) : true) /** The events of this session in the tail after `after`. */ const ownAfter = (after: number) => - this.tail - .filter((item) => item.thread === thread) - .map((item) => item.entry) - .filter((entry) => Number(entry.cursor) > after && accepts(entry)) + this.tail.after(after, thread).entries.filter(accepts) let after = Number(options.from ?? '0') if (!Number.isFinite(after)) after = 0 for (;;) { if (signal?.aborted) return - if (after < this.tailFrom) { + if (after < this.tail.from) { // Older than the in-memory tail: page it from the store. - const upTo = this.tailFrom + const upTo = this.tail.from let page: Array try { page = await this.store.read({ @@ -2298,14 +2497,14 @@ export class SharedLog { // A truncate dropped the events after `after`: say so once, then // go on from the first stored record. const resumeAt = error.firstSeq - 1 - yield gapEvent(String(after), String(resumeAt)) + yield gapEvent(thread, String(after), String(resumeAt), 'truncated') after = resumeAt continue } for (const entry of page) { if (entry.seq > upTo || signal?.aborted) break after = entry.seq - const event = sessionEventOf(entry) + const event = sessionEventOf(entry, this.rootThread) const isOwn = threadOf(entry.record, this.rootThread) === thread if (event && isOwn && accepts(event)) yield event } @@ -2371,29 +2570,18 @@ function mergeDelta(first: StreamChunk, second: StreamChunk): StreamChunk { return second } -/** The `harness.gap` event for a reader at `from`. Its cursor is `resumeAt`. */ -function gapEvent(from: Cursor, resumeAt: Cursor): SessionEvent { - const value: EventGap = { from, resumeAt, reason: 'truncated' } - return { - cursor: resumeAt, - // No operation made it. Like a session input, it belongs to the session. - operationId: 'session', - event: { - type: EventType.CUSTOM, - name: HARNESS_EVENTS.gap, - value, - timestamp: Date.now(), - }, - } -} - -function sessionEventOf(entry: LogEntry) { +/** + * The session event of `entry`, or `undefined` when it is not an event. A + * record with no `thread` is an event of `rootThread`. + */ +function sessionEventOf(entry: LogEntry, rootThread: string) { const { record } = entry if (!isHarnessRecord(record) || record.type !== 'harness.event') { return undefined } const event: SessionEvent = { cursor: String(entry.seq), + threadId: threadOf(record, rootThread), operationId: record.operationId, event: record.event, } @@ -2440,6 +2628,18 @@ export function engineMessageStore(options: SessionLogView) { const shared = sessionMessageStore(options) const fold = () => [...writer.state.messages] let held = fold() + /** + * The list that the last `beforeModel` gave the model, and the fold after + * its commit. A second call of the same model call (`steering`, then + * `syncFold`) with the same message objects and the same fold has + * nothing to commit, so it skips the compare and the commit. + */ + let committed: + | { + list: ReadonlyArray + fold: ReadonlyArray + } + | undefined /** * What to write for an engine save of `list`, when the engine held `base`: * @@ -2466,11 +2666,13 @@ export function engineMessageStore(options: SessionLogView) { const engine = { loadThread: async (threadId: string) => { if (threadId !== writer.threadId) return shared.loadThread(threadId) + committed = undefined held = fold() return [...held] }, saveThread: (threadId: string, list: Array) => { if (threadId !== writer.threadId) return shared.saveThread(threadId, list) + committed = undefined const target = rebase(held, list) held = [...list] return writer.commit({ messages: target }) @@ -2485,13 +2687,22 @@ export function engineMessageStore(options: SessionLogView) { list: ReadonlyArray, records: ReadonlyArray = [], ) => { + const isUnchanged = + records.length === 0 && + committed !== undefined && + committed.fold === writer.state.messages && + committed.list.length === list.length && + committed.list.every((message, index) => message === list[index]) + if (isUnchanged) return undefined const target = rebase(held, [...list]) held = [...target] await writer.commit({ messages: target, records }) const isSame = target.length === list.length && commonPrefix(target, list) === list.length - return isSame ? undefined : [...target] + const result = isSame ? undefined : [...target] + committed = { list: result ?? [...list], fold: writer.state.messages } + return result }, /** * Host records from a middleware of the running turn. The messages that @@ -2507,6 +2718,7 @@ export function engineMessageStore(options: SessionLogView) { list: ReadonlyArray, records: ReadonlyArray, ) => { + committed = undefined const kept = commonPrefix(held, list) if (kept < held.length || list.length === held.length) { return writer.append(records) diff --git a/packages/ai-harness/src/media-signing.ts b/packages/ai-harness/src/media-signing.ts new file mode 100644 index 0000000000..259a59367d --- /dev/null +++ b/packages/ai-harness/src/media-signing.ts @@ -0,0 +1,114 @@ +import { base64url } from './utils' + +/** + * How a media signer gets its keys. + * + * - `secret`: one key. It signs and verifies. + * - `keys`: a key list, to rotate keys. The first key signs. Every key + * verifies, so a URL that an old key signed works until it expires. + * + * Each key must be at least {@link MIN_KEY_BYTES} bytes (UTF-8). + */ +export type MediaSignerOptions = + | { secret: string } + | { keys: ReadonlyArray } + +/** The smallest key length in bytes. HMAC SHA-256 wants 32 bytes or more. */ +export const MIN_KEY_BYTES = 32 + +const HMAC = { name: 'HMAC', hash: 'SHA-256' } + +const encoder = new TextEncoder() + +/** Base64url to bytes, or `undefined` for a value that is not base64url. */ +function fromBase64url(value: string) { + try { + const binary = atob(value.replace(/-/g, '+').replace(/_/g, '/')) + return Uint8Array.from(binary, (char) => char.charCodeAt(0)) + } catch { + return undefined + } +} + +/** The keys of `options`, first the key that signs. */ +function keysOf(options: MediaSignerOptions) { + return 'secret' in options ? [options.secret] : options.keys +} + +/** `bytes` as an HMAC SHA-256 key that signs and verifies. */ +function importHmacKey(bytes: Uint8Array) { + return crypto.subtle.importKey('raw', bytes, HMAC, false, ['sign', 'verify']) +} + +/** + * Make a signer for media URLs. It signs with HMAC SHA-256 through Web Crypto + * (`crypto.subtle`), so it runs on Node, Deno, Bun, and edge runtimes. + * + * `payload` is the text that the URL must keep unchanged, for example the + * harness name, tenant, thread, media id, user, and expiry, one per line. A + * signature is base64url without padding. + * + * Throws an `Error` when there is no key, or when a key is shorter than + * {@link MIN_KEY_BYTES} bytes (UTF-8). + * + * @example + * ```ts + * const signer = createMediaSigner({ keys: [newSecret, oldSecret] }) + * const sig = await signer.sign(payload) + * const isValid = await signer.verify(payload, sig) + * ``` + */ +export function createMediaSigner(options: MediaSignerOptions) { + const encoded = keysOf(options).map((secret) => encoder.encode(secret)) + const [signingBytes, ...otherBytes] = encoded + if (signingBytes === undefined) { + throw new Error('Media signing needs at least one key.') + } + for (const [index, bytes] of encoded.entries()) { + if (bytes.byteLength < MIN_KEY_BYTES) { + throw new Error( + `Media signing key ${index} is ${bytes.byteLength} bytes. Use a key of at least ${MIN_KEY_BYTES} bytes (UTF-8).`, + ) + } + } + // Import each key one time, on first use. + let imported: + | { signing: Promise; all: Promise> } + | undefined + const cryptoKeys = () => { + if (imported) return imported + const signing = importHmacKey(signingBytes) + imported = { + signing, + all: Promise.all([signing, ...otherBytes.map(importHmacKey)]), + } + return imported + } + + return { + /** The base64url signature of `payload`, made with the first key. */ + async sign(payload: string) { + const signature = await crypto.subtle.sign( + 'HMAC', + await cryptoKeys().signing, + encoder.encode(payload), + ) + return base64url(new Uint8Array(signature)) + }, + /** + * True when one of the keys signed `payload` as `signature`. Each key + * checks in constant time (`crypto.subtle.verify`). Every key checks, so + * the time does not tell which key matched. + */ + async verify(payload: string, signature: string) { + const bytes = fromBase64url(signature) + if (bytes === undefined) return false + const data = encoder.encode(payload) + const keys = await cryptoKeys().all + const results = await Promise.all( + keys.map((key) => crypto.subtle.verify('HMAC', key, bytes, data)), + ) + return results.includes(true) + }, + } +} diff --git a/packages/ai-harness/src/media.ts b/packages/ai-harness/src/media.ts index 620897f39e..a5b2c5edd1 100644 --- a/packages/ai-harness/src/media.ts +++ b/packages/ai-harness/src/media.ts @@ -39,6 +39,27 @@ export interface MediaOptions { accepts?: ReadonlyArray /** Turns audio into text first when the model cannot read audio. */ transcribe?: AnyTranscriptionAdapter + /** + * Whether a client input can send a content part with an `http:` or + * `https:` URL. The model provider can fetch that URL. Default `false`: the + * input is refused. A function decides for each URL. + */ + externalUrls?: boolean | ((url: URL) => boolean) + /** + * Whether a client input can send a content part with a provider file + * handle (`source.type: 'file'`). The adapter reads that file with the + * server's provider key, so a client that knows another user's file id + * could make the model read it. Default `false`: the input is refused. + * Server code can always send file handles. + */ + allowProviderFiles?: boolean + /** + * The largest request body or socket frame, in bytes. Default 32 MiB. A + * body stops being read at the limit. A socket frame is checked only after + * the WebSocket server read all of it, so set the frame limit of the server + * (for example `maxPayload` of `ws`) too. + */ + maxBodyBytes?: number } /** @@ -65,6 +86,18 @@ export type MediaStore = ReturnType const DEFAULT_MAX_BYTES = 100 * 1024 * 1024 const DATA_URL_MAX_BYTES = 1024 * 1024 +const DEFAULT_ENCODED_BYTES = 32 * 1024 * 1024 +/** The base64 LRU of `createMediaStore`, one for each blob store. */ +interface EncodedCache { + /** The most characters it keeps. */ + cap: number + /** The characters it keeps now. */ + size: number + /** By `/`, the oldest first. */ + entries: Map +} +const encodedCaches = new WeakMap() + const ALL_KINDS: ReadonlyArray = [ 'image', 'audio', @@ -176,12 +209,18 @@ export function createMediaStore({ threadId, tenantId, options = {}, + encodedBytes, }: { persistence: MediaPersistence threadId: string /** The tenant of the thread. Without it, the default tenant. */ tenantId?: string options?: MediaOptions + /** + * @internal How many characters of base64 to keep across turns, for all + * the media stores of this blob store. The largest value asked for wins. + */ + encodedBytes?: number }) { const { artifacts, blobs } = persistence.stores const storeThreadId = logIdOf({ tenantId, logId: threadId }) @@ -204,7 +243,7 @@ export function createMediaStore({ return record && kind ? toMediaRecord(record, kind, threadId) : null } - async function load(id: string, range?: BlobRange) { + async function blobOf(id: string, range?: BlobRange) { const record = await artifact(id) const blob = record ? await retrieveBlob( @@ -214,7 +253,52 @@ export function createMediaStore({ ) : null if (!blob) throw new MediaError(404, `Media ${id} was not found.`) - return new Uint8Array(await blob.arrayBuffer()) + return blob + } + + async function load(id: string, range?: BlobRange) { + return new Uint8Array(await (await blobOf(id, range)).arrayBuffer()) + } + + // ponytail: an LRU of base64 text across turns, at most `encodedBytes` + // characters (default 32 MB). The media stores of one blob store (all the + // sessions of a host) share it, so the cap is for the host, not for each + // open session. The largest `encodedBytes` asked for is the cap. The key is + // the store key of the thread and the media id: the bytes of an id never + // change. Above the cap, each turn encodes the oldest files again. Send + // adapters a URL when that gets slow. + const cap = encodedBytes ?? DEFAULT_ENCODED_BYTES + let encoded = encodedCaches.get(blobs) + if (!encoded) { + encoded = { cap, size: 0, entries: new Map() } + encodedCaches.set(blobs, encoded) + } + encoded.cap = Math.max(encoded.cap, cap) + const cache = encoded + + /** The base64 of the bytes of `id`. Call `get` first: a hit does not check the record. */ + async function base64(id: string) { + const key = `${storeThreadId}/${id}` + const hit = cache.entries.get(key) + if (hit !== undefined) { + // Delete and set again, so the Map order puts it last (newest). + cache.entries.delete(key) + cache.entries.set(key, hit) + return hit + } + const value = toBase64(await load(id)) + // Another call stored it while this one read the bytes. + if (value.length <= cache.cap && !cache.entries.has(key)) { + cache.entries.set(key, value) + cache.size += value.length + // The first keys of the Map are the oldest. + for (const [oldKey, old] of cache.entries) { + if (cache.size <= cache.cap) break + cache.entries.delete(oldKey) + cache.size -= old.length + } + } + return value } return { @@ -274,6 +358,12 @@ export function createMediaStore({ get, /** The bytes of `id`, or one `range` of them. */ load, + /** The bytes of `id`, or one `range` of them, as a stream. */ + async open(id: string, range?: BlobRange) { + const blob = await blobOf(id, range) + return blob.body ?? new Blob([await blob.arrayBuffer()]).stream() + }, + base64, /** A data URL for a file up to 1 MB. `undefined` above that. */ async dataUrl(id: string) { const record = await get(id) @@ -287,7 +377,7 @@ export function createMediaStore({ (part, index) => index === 0 || /^[\w.+-]+=[\w.+-]+$/.test(part), ) .join(';') - return `data:${mimeType};base64,${toBase64(await load(id))}` + return `data:${mimeType};base64,${await base64(id)}` }, } } @@ -436,10 +526,9 @@ export function mediaMiddleware({ transcribe?: AnyTranscriptionAdapter }) { // ponytail: one cache per run, so the loop does not load or transcribe a - // file again at every model call. A new turn does it again for history: - // each turn loads and base64-encodes every media file of the thread (up to - // `maxBytes` each). Keep the encoded parts across turns by media id (an - // LRU), or send adapters a URL, when long media threads get slow. + // file again at every model call. Across turns, the store keeps the base64 + // of each file (an LRU with a byte cap). A transcript is made again at each + // turn; keep it too when that gets slow. const runs = new WeakMap>>() async function resolve( @@ -455,24 +544,26 @@ export function mediaMiddleware({ if (!isAccepted && transcriber === undefined) { throw new Error(cannotRead(ctx.model, record.kind)) } - const bytes = await store.load(id).catch((error: unknown) => { + const notFound = (error: unknown) => { if (error instanceof MediaError && error.status === 404) return undefined throw error - }) - if (!bytes) return textPart(`[${record.kind} not found: ${record.name}]`) + } if (transcriber === undefined) { + // The store keeps the base64 across turns. + const value = await store.base64(id).catch(notFound) + if (value === undefined) { + return textPart(`[${record.kind} not found: ${record.name}]`) + } // The stored kind wins: a `video` part with an image id sends an image. const swapped: ContentPart = { ...part, type: record.kind, - source: { - type: 'data', - value: toBase64(bytes), - mimeType: record.mimeType, - }, + source: { type: 'data', value, mimeType: record.mimeType }, } return swapped } + const bytes = await store.load(id).catch(notFound) + if (!bytes) return textPart(`[${record.kind} not found: ${record.name}]`) const transcript = await generateTranscription({ adapter: transcriber, audio: new Blob([bytes], { type: record.mimeType }), diff --git a/packages/ai-harness/src/observability.ts b/packages/ai-harness/src/observability.ts new file mode 100644 index 0000000000..6a2422fb65 --- /dev/null +++ b/packages/ai-harness/src/observability.ts @@ -0,0 +1,380 @@ +// Tracing with the OpenTelemetry API. `@opentelemetry/api` is an optional +// peer: the types are type-only imports, and the runtime is loaded only when a +// host has a tracer (see `OperationTracing.ready`). +import { addCounts, emptyCounts } from './usage' +import type { + Attributes, + Context, + ContextAPI, + Span, + SpanContext, + SpanOptions, + TraceAPI, + Tracer, +} from '@opentelemetry/api' +import type { OperationKind, OperationStatus } from './types' +import type { UsageCall, UsageCounts } from './usage' + +/** What `onOperation` gets, once when an operation starts and once when it settles. */ +export interface OperationReport { + phase: 'started' | 'settled' + operation: { + id: string + kind: OperationKind + agent?: string + threadId: string + tenantId?: string + parentRunId?: string + status: OperationStatus + /** Only when `phase` is `'settled'`. */ + durationMs?: number + error?: { message: string; code?: string } + usage?: UsageCounts + } +} + +/** The `observability` option of a host. */ +export interface ObservabilityOptions { + /** + * An OpenTelemetry tracer, for example `trace.getTracer('my-app')`. + * Without a tracer, the harness writes no spans. + */ + tracer?: Tracer + /** A listener that throws is logged and ignored. */ + onOperation?: (report: OperationReport) => void | Promise +} + +/** The `SpanStatusCode` values, without a runtime import of the enum. */ +const OK = 1 +const ERROR = 2 + +const NOOP_CONTEXT: SpanContext = { + traceId: '0'.repeat(32), + spanId: '0'.repeat(16), + traceFlags: 0, +} + +// One shared span, so a host without a tracer allocates nothing per span. +const noopSpan: Span = { + setAttribute: () => noopSpan, + setAttributes: () => noopSpan, + addEvent: () => noopSpan, + addLink: () => noopSpan, + addLinks: () => noopSpan, + setStatus: () => noopSpan, + updateName: () => noopSpan, + end: () => {}, + isRecording: () => false, + recordException: () => {}, + spanContext: () => NOOP_CONTEXT, +} + +/** The tracer of a host without `observability.tracer`. It does nothing. */ +export const noopTracer: Pick = { + startSpan: () => noopSpan, +} + +let otel: { trace: TraceAPI; context: ContextAPI } | undefined + +/** + * Load `@opentelemetry/api` once, for a host with a tracer. The harness + * needs it only to make the context of a parent span. Without the package, + * spans have no parent on this host. + */ +async function loadOpenTelemetry(): Promise { + otel ??= await import('@opentelemetry/api').catch(() => undefined) +} + +/** + * Run `fn` in a span. The span ends when `fn` settles. When `fn` throws, the + * span records the error and gets the error status, and the error goes on. + * + * @example + * ```ts + * const result = await withSpan({ tracer, name: 'harness.recovery' }, () => + * recover(), + * ) + * ``` + */ +export async function withSpan( + { + tracer, + name, + ...options + }: { tracer: Pick; name: string } & SpanOptions, + fn: (span: Span) => T | PromiseLike, +): Promise { + const span = tracer.startSpan(name, options) + try { + return await fn(span) + } catch (error) { + const message = error instanceof Error ? error.message : String(error) + span.recordException(error instanceof Error ? error : message) + span.setStatus({ code: ERROR, message }) + throw error + } finally { + span.end() + } +} + +const TRACEPARENT = + /^([0-9a-f]{2})-([0-9a-f]{32})-([0-9a-f]{16})-([0-9a-f]{2})(-.*)?$/ + +/** + * Read a W3C `traceparent` header value. Returns `undefined` for a value + * that is not valid, as the W3C rules say: the caller then starts a new + * trace. + * + * @example + * ```ts + * parseTraceparent('00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01') + * // { traceId: '4bf9…', spanId: '00f0…', traceFlags: 1, isRemote: true } + * ``` + */ +export function parseTraceparent(value: unknown): SpanContext | undefined { + if (typeof value !== 'string') return undefined + const match = TRACEPARENT.exec(value) + if (!match) return undefined + const [, version, traceId = '', spanId = '', flags = '', rest] = match + // Version `ff` is not valid. Version `00` has exactly four fields. A later + // version can add fields after the four that we read. + if (version === 'ff' || (version === '00' && rest !== undefined)) + return undefined + if (/^0+$/.test(traceId) || /^0+$/.test(spanId)) return undefined + return { + traceId: traceId, + spanId: spanId, + traceFlags: parseInt(flags, 16), + isRemote: true, + } +} + +/** The W3C `traceparent` header value of a span context. */ +export function traceparentOf(context: SpanContext): string { + const flags = (context.traceFlags & 0xff).toString(16).padStart(2, '0') + return `00-${context.traceId}-${context.spanId}-${flags}` +} + +/** A context whose span is `span`: a span started in it is its child. */ +const parentContext = (span: Span | undefined): Context | undefined => + span && otel ? otel.trace.setSpan(otel.context.active(), span) : undefined + +/** What {@link OperationTracing} reads of an operation. */ +interface TracedOperation { + id: string + kind: OperationKind + agent?: string | undefined + status: () => OperationStatus + then: ( + onfulfilled: (value: unknown) => void, + onrejected: (reason: unknown) => void, + ) => unknown +} + +/** An operation of this host that did not settle yet. */ +interface OpenOperation { + span: Span + startedAt: number + usage: UsageCounts +} + +/** The message and the `code` of an error, when it has a string code. */ +function errorOf(error: unknown): { message: string; code?: string } { + const message = error instanceof Error ? error.message : String(error) + const code = + typeof error === 'object' && + error !== null && + 'code' in error && + typeof error.code === 'string' + ? error.code + : undefined + return code === undefined ? { message } : { message, code } +} + +/** + * @internal The spans and the `onOperation` reports of one session. A + * session of a host without `observability` has none, so it allocates + * nothing. + */ +export class OperationTracing { + private readonly tracer: Pick + private readonly open = new Map() + + constructor( + private readonly options: ObservabilityOptions & { + threadId: string + tenantId?: string + /** Logs a listener or a tracer that throws. */ + warn: (error: unknown) => void + }, + ) { + this.tracer = options.tracer ?? noopTracer + } + + /** Load what the tracer needs. A session awaits it before its first span. */ + async ready(): Promise { + if (this.options.tracer) await loadOpenTelemetry() + } + + /** + * Run the tracer or span calls `fn`. A tracer that throws is logged, so a + * tracing bug never fails a run or the host. + */ + safely(fn: () => void): void { + try { + fn() + } catch (error) { + this.options.warn(error) + } + } + + /** Start a span. A tracer that throws gives a span that does nothing. */ + private startSpan(name: string, options: SpanOptions, parent?: Span): Span { + let span = noopSpan + this.safely(() => { + span = this.tracer.startSpan(name, options, parentContext(parent)) + }) + return span + } + + /** The attributes of every span of this session, and `extra`. */ + private attributes(extra: Attributes) { + const attributes: Attributes = { + 'harness.thread_id': this.options.threadId, + } + if (this.options.tenantId !== undefined) { + attributes['harness.tenant_id'] = this.options.tenantId + } + for (const [key, value] of Object.entries(extra)) { + if (value !== undefined) attributes[key] = value + } + return attributes + } + + /** + * Start a span `name` of this session. `runId`: the run that it belongs + * to. When that run is open on this host, the span is its child. + */ + span( + name: string, + extra: Attributes, + options: { runId?: string | undefined; startTime?: number } = {}, + ): Span { + const parent = + options.runId === undefined ? undefined : this.open.get(options.runId) + return this.startSpan( + name, + { + attributes: this.attributes(extra), + ...(options.startTime !== undefined + ? { startTime: options.startTime } + : {}), + }, + parent?.span, + ) + } + + /** + * An operation started on this host. Its `harness.operation` span is a + * child of the span of its parent run when that run is open on this host. + * A valid `traceparent` of the caller becomes a link, never the parent: + * it can come from another machine. The span ends, and `onOperation` + * gets the report, when the operation settles. + */ + started( + operation: TracedOperation, + run: { parentRunId: string | undefined; traceparent: string | undefined }, + ): void { + if (this.open.has(operation.id)) return + const { parentRunId } = run + const parent = + parentRunId === undefined ? undefined : this.open.get(parentRunId) + const caller = parseTraceparent(run.traceparent) + const startedAt = Date.now() + const span = this.startSpan( + 'harness.operation', + { + attributes: this.attributes({ + 'harness.run_id': operation.id, + 'harness.parent_run_id': parentRunId, + 'harness.operation.kind': operation.kind, + 'harness.agent': operation.agent, + }), + ...(caller ? { links: [{ context: caller }] } : {}), + startTime: startedAt, + }, + parent?.span, + ) + const open: OpenOperation = { span, startedAt, usage: emptyCounts() } + this.open.set(operation.id, open) + const report = (): OperationReport['operation'] => ({ + id: operation.id, + kind: operation.kind, + ...(operation.agent !== undefined ? { agent: operation.agent } : {}), + threadId: this.options.threadId, + ...(this.options.tenantId !== undefined + ? { tenantId: this.options.tenantId } + : {}), + ...(parentRunId !== undefined ? { parentRunId } : {}), + status: operation.status(), + }) + this.report({ phase: 'started', operation: report() }) + const settle = (error?: unknown) => { + this.open.delete(operation.id) + const status = operation.status() + const failure = error === undefined ? undefined : errorOf(error) + const endedAt = Date.now() + this.safely(() => { + span.setAttribute('harness.operation.status', status) + if (status === 'failed' && failure) { + span.recordException(error instanceof Error ? error : failure.message) + span.setStatus({ code: ERROR, message: failure.message }) + } else if (status !== 'cancelled') { + span.setStatus({ code: OK }) + } + span.end(endedAt) + }) + this.report({ + phase: 'settled', + operation: { + ...report(), + durationMs: endedAt - startedAt, + ...(failure ? { error: failure } : {}), + usage: open.usage, + }, + }) + } + operation.then( + () => settle(), + (error: unknown) => settle(error), + ) + } + + /** + * One model call of the run `call.runId`. It counts for that run and for + * each run above it that is open on this host, as the run usage of the + * log does. `parentOf` gives the parent run of a run. + */ + usage(call: UsageCall, parentOf: (runId: string) => string | undefined) { + const seen = new Set() + let runId = call.runId + while (runId !== undefined && !seen.has(runId)) { + seen.add(runId) + const open = this.open.get(runId) + if (open) addCounts(open.usage, call.usage) + runId = parentOf(runId) + } + } + + /** Call `onOperation`. A listener that throws is logged and ignored. */ + private report(report: OperationReport) { + const listener = this.options.onOperation + if (!listener) return + try { + const done = listener(report) + if (done) done.then(undefined, this.options.warn) + } catch (error) { + this.options.warn(error) + } + } +} diff --git a/packages/ai-harness/src/operation.ts b/packages/ai-harness/src/operation.ts index 8221763f91..de273686d6 100644 --- a/packages/ai-harness/src/operation.ts +++ b/packages/ai-harness/src/operation.ts @@ -23,7 +23,10 @@ export function createOperationId(kind: OperationKind): string { /** The session-side implementation of {@link Operation}. */ export class OperationImpl implements Operation { readonly id: string - /** The feed head when the operation was made. Its events come after it. */ + /** + * The feed head when the operation was made. Its events come after it. A + * rebuilt operation gets the start of its run in the log. + */ readonly startedCursor: Cursor readonly abortController = new AbortController() readonly receipt: Promise @@ -51,9 +54,11 @@ export class OperationImpl implements Operation { readonly agent?: string, /** A stored operation id, when a restart rebuilds this operation. */ id?: string, + /** The cursor before the first event of a rebuilt operation. */ + startedCursor?: Cursor, ) { this.id = id ?? createOperationId(kind) - this.startedCursor = feed.head() + this.startedCursor = startedCursor ?? feed.head() // A settled operation sends no more events, so a listener is not added, // and its listeners go once it settled. After a tick, so the events it // sends as it settles (the finished event) still reach them. @@ -103,7 +108,9 @@ export class OperationImpl implements Operation { return this.ended } + /** Set the status. A settled operation keeps its final status. */ setStatus(status: OperationStatus): void { + if (this.isSettled()) return this.current = status } @@ -116,7 +123,12 @@ export class OperationImpl implements Operation { this.settleReceipt(receipt) } + /** + * Settle with `result`. The first settle wins: when the operation settled + * already, this call does nothing. + */ finish(status: 'completed' | 'interrupted', result: TResult): void { + if (this.isSettled()) return this.current = status this.ended = true this.resolveResult(result) @@ -124,7 +136,12 @@ export class OperationImpl implements Operation { this.feed.wake() } + /** + * Settle with `error`. The first settle wins: when the operation settled + * already, this call does nothing. + */ fail(status: 'failed' | 'cancelled', error: unknown): void { + if (this.isSettled()) return this.current = status this.ended = true this.rejectResult(error) @@ -137,6 +154,9 @@ export class OperationImpl implements Operation { }): AsyncIterable { return this.feed.read({ ...options, + // Every event of this operation comes after its start, so the read + // does not walk the older events of the feed. + from: options?.from ?? this.startedCursor, filter: (entry) => entry.operationId === this.id, until: () => this.isSettled(), }) @@ -145,7 +165,7 @@ export class OperationImpl implements Operation { async *stream(options?: { signal?: AbortSignal }): AsyncIterable { - for await (const entry of this.events({ from: '0', ...options })) { + for await (const entry of this.events(options)) { yield entry.event } } diff --git a/packages/ai-harness/src/ownership.ts b/packages/ai-harness/src/ownership.ts index 26145c865d..9228111787 100644 --- a/packages/ai-harness/src/ownership.ts +++ b/packages/ai-harness/src/ownership.ts @@ -4,6 +4,12 @@ import type { LogStore } from '@tanstack/ai-persistence' /** How long a host owns a log after each renewal, and how often it renews. */ export const OWNERSHIP = { ttlMs: 30_000, renewMs: 10_000 } +/** + * How long a sweep waits before it tries again a log whose thread did not + * open. The log goes back with this wake time. + */ +export const OPEN_RETRY_MS = 30_000 + /** * Ownership timing for a host, and its stable id. A missing timing field * uses {@link OWNERSHIP}. @@ -22,10 +28,53 @@ export interface OwnershipOptions { * store clock. Default 30 000. */ ttlMs?: number - /** How often the host renews every log it owns, in milliseconds. Default 10 000. */ + /** + * How often the host renews every log it owns, in milliseconds. It must be + * less than `ttlMs`. Default 10 000, or a third of `ttlMs` when you set + * only `ttlMs`. + */ renewMs?: number } +/** + * @internal For tests only. Timing with this key is not checked, so a test + * host can own a log for a short time and never renew it, like a process + * that stopped. + */ +export const UNCHECKED_TIMING = Symbol('harness.uncheckedTiming') + +/** + * The `ttlMs` and `renewMs` of `timing`, with the defaults. Throws a + * `RangeError` for a value that is not more than 0, and when `renewMs` is + * not less than `ttlMs`: then each ownership expires before its renewal, + * and other hosts take the logs of this host while it runs. + */ +function ownershipTiming(timing: OwnershipOptions = {}) { + const ttlMs = timing.ttlMs ?? OWNERSHIP.ttlMs + const renewMs = + timing.renewMs ?? + (timing.ttlMs === undefined + ? OWNERSHIP.renewMs + : Math.max(1, Math.floor(ttlMs / 3))) + if (UNCHECKED_TIMING in timing) return { ttlMs, renewMs } + for (const [name, value] of [ + ['ttlMs', ttlMs], + ['renewMs', renewMs], + ] as const) { + if (!Number.isFinite(value) || value <= 0) { + throw new RangeError( + `ownership.${name} must be a number more than 0, not ${String(value)}.`, + ) + } + } + if (renewMs >= ttlMs) { + throw new RangeError( + `ownership.renewMs (${renewMs}) must be less than ownership.ttlMs (${ttlMs}). Else the ownership of a log expires before its renewal, and other hosts take it.`, + ) + } + return { ttlMs, renewMs } +} + /** * A write that this host cannot do: another host owns the log, or took it * from this host. Nothing was written. This host still reads the log. @@ -61,6 +110,8 @@ export class HostOwnership { private readonly renewMs: number /** The epoch of each log that this host owns, by store log id. */ private readonly owned = new Map() + /** The owned logs that the store gave over an owner that stopped. */ + private readonly takenOver = new Set() private readonly lostListeners = new Map< string, Set<(error: OwnershipLostError) => void> @@ -90,8 +141,9 @@ export class HostOwnership { this.store = options.store this.ownerId = options.ownerId this.affinity = options.affinity - this.ttlMs = options.timing?.ttlMs ?? OWNERSHIP.ttlMs - this.renewMs = options.timing?.renewMs ?? OWNERSHIP.renewMs + const { ttlMs, renewMs } = ownershipTiming(options.timing) + this.ttlMs = ttlMs + this.renewMs = renewMs } /** The epoch of `logId` while this host owns it. */ @@ -99,6 +151,14 @@ export class HostOwnership { return this.owned.get(logId) } + /** + * True when this host owns `logId` because the ownership of another + * owner expired: that owner stopped with no release. + */ + tookOver(logId: string) { + return this.takenOver.has(logId) + } + /** * Own `logId`. Returns the epoch, or `null` when another live host owns * it. A log that this host owns already keeps its epoch. @@ -115,6 +175,7 @@ export class HostOwnership { }) if (got === null) return null this.owned.set(logId, got.epoch) + if (got.expired) this.takenOver.add(logId) this.startTimer() return got.epoch } @@ -131,19 +192,21 @@ export class HostOwnership { /** Stop owning `logId` with no store call: the store refused its epoch. */ forget(logId: string) { this.owned.delete(logId) + this.takenOver.delete(logId) if (this.owned.size === 0) this.stopTimer() } /** * Give `logId` up. With `epoch`, only that ownership: a later acquire * keeps the log. A log that this host does not own does nothing. The - * store gets the wake time of {@link setWake}, when the log has one. + * store gets `wakeIn` as the wake time, else the wake time of + * {@link setWake}, when the log has one. */ - release(logId: string, epoch = this.owned.get(logId)) { + release(logId: string, epoch = this.owned.get(logId), wakeIn?: number) { const isOwned = epoch !== undefined && this.owned.get(logId) === epoch if (!isOwned) return Promise.resolve() this.forget(logId) - const wakeInMs = this.wakes.get(logId)?.() + const wakeInMs = wakeIn ?? this.wakes.get(logId)?.() const released = this.store .release({ logId, diff --git a/packages/ai-harness/src/plugins.ts b/packages/ai-harness/src/plugins.ts index 4e3f7b5089..08fd7f3ee9 100644 --- a/packages/ai-harness/src/plugins.ts +++ b/packages/ai-harness/src/plugins.ts @@ -10,6 +10,8 @@ import type { AnyTool, Capability, CapabilityHandle, + ChatResult, + InferSchemaType, KeyedAdapter, MetadataStore, ModelMessage, @@ -22,7 +24,7 @@ import type { AgentRegistryView, AnyAgent, } from './agents' -import type { TurnInfo } from './types' +import type { RunContext, TurnInfo } from './types' import type { CredentialsAccess } from './auth' import type { AnyCommand, PluginSessionApi } from './commands' import type { ConfigOption } from './config' @@ -30,6 +32,7 @@ import type { ExtensionItem, ExtensionPoint, PluginEvent } from './extensions' import type { MetadataOf } from './session-entry' import type { SharedHostPlugin } from './host-plugins' import type { WorkspaceBinding } from './log-records' +import type { UsageCall } from './usage' /** * How long a plugin's `setup` result and resources live: @@ -50,6 +53,15 @@ export type PluginLifetime = 'session' | 'turn' | 'host' export interface PluginPrompt { id: string text: string | ((turn: TurnInfo) => string) + /** + * Where the model gets the text. `'system'` (default): in the system + * prompts, in plugin order. `'context'`: in one system prompt after all + * the others, made again at each model call. Use `'context'` for text that + * changes often, such as a todo list or a mode: a change makes the + * provider prompt cache miss only from that last prompt on. The log does + * not keep the `'context'` text. + */ + placement?: 'system' | 'context' } /** What a plugin's `setup` returns. Every field is optional. */ @@ -95,6 +107,13 @@ export interface PluginContributions { adapter?: ( turn: TurnInfo, ) => AnyTextAdapter | KeyedAdapter | undefined + /** + * Runs once at the start of each turn, before the adapter pick, the tools, + * and the prompts. `turn.run` has the turn's user and that user's + * credentials. Use it to refresh state that a prompt reads. A hook that + * throws is skipped, and clients get a `harness.plugin.warning` event. + */ + onTurnStart?: (turn: TurnInfo) => void | Promise /** * 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, @@ -109,11 +128,13 @@ export interface PluginContributions { * `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`. `model` is the model id of the - * turn's adapter, so the list can differ for each model. + * turn's adapter, so the list can differ for each model. `run` is the run + * of the turn: its id and its user. */ prepareTools?: (turn: { tools: ReadonlyArray model: string + run: RunContext }) => ReadonlyArray | Promise> } @@ -197,6 +218,72 @@ export interface PluginServices { agents: AgentActions /** Called after `ctx.commands` adds or removes a command. */ commandsChanged?: () => void + /** One model call of `ctx.model.chat`. Resolves to what chat() gives. */ + modelChat: (call: PluginModelCall) => Promise +} + +/** The chat() options of one `ctx.model.chat` call. */ +export interface PluginChatOptions { + messages: Array + systemPrompts?: Array + /** A schema for the answer. The call then resolves to the checked value. */ + outputSchema?: TSchema +} + +/** What `ctx.model.chat` takes. */ +export interface PluginModelChatOptions< + TSchema extends SchemaInput | undefined = undefined, +> { + /** + * The model. A `keyedAdapter(...)` gets the key of the user of `run`. + * Default: the model of the session, as the summary of `host.transfer` + * picks it: the thread setting `model`, else the harness `adapter`. + */ + adapter?: AnyTextAdapter | KeyedAdapter + options: PluginChatOptions + /** A name for this call in its usage record, for example `judge`. */ + label?: string + /** + * The run that the call is for: its user's keys, its signal, and its + * budget. In a tool or a middleware hook, pass `ctx.context.run`. In a + * command, pass `run` of the command context. Default: the turn of a + * plugin with `lifetime: 'turn'`. + */ + run?: RunContext + /** Stop the call after this time, in milliseconds. Default 120000 (2 minutes). */ + timeoutMs?: number +} + +/** What `ctx.model.chat` resolves to: the checked value with `outputSchema`, else the text and chunks. */ +export type PluginModelChatResult = + TSchema extends SchemaInput ? InferSchemaType : ChatResult + +/** `ctx.model` of a plugin. */ +export interface PluginModel { + /** + * Call a model, for example to judge, to name, or to summarize. The call + * stops when its run stops or after `timeoutMs`, and then rejects. Its + * usage counts in the session (`session.usage().bySource` under + * `plugin:`) and in its run, and the budget of the run applies. + */ + chat: ( + options: PluginModelChatOptions, + ) => Promise> +} + +/** @internal One model call of a plugin, as the session runs it. */ +export interface PluginModelCall extends Omit< + PluginModelChatOptions, + 'run' +> { + run: RunContext + /** `plugin:`, `router`, or `transfer`. */ + source: string + /** + * Where the usage of each model call goes. Default: the session and the + * run `run`, and the budget of the run applies. + */ + count?: (call: UsageCall) => Promise } /** @@ -300,6 +387,8 @@ export interface PluginSetupContext> { session: PluginSessionApi /** Add and remove this plugin's commands while the session runs. */ commands: PluginCommands + /** Call a model, with the usage counted under `plugin:`. */ + model: PluginModel /** * The turn this plugin is set up for: its input, context, sender, and run. * Set only for a plugin with `lifetime: 'turn'`. A session plugin has no @@ -328,6 +417,19 @@ interface PluginDefinitionBase { optionalRequires?: ReadonlyArray /** Agent outputs this plugin needs the session to have. */ needs?: { produces?: ReadonlyArray } + /** + * The host records of this plugin that the log fold keeps, by record + * type, with how many of each: the last ones. The fold snapshot keeps + * them, so they stay readable after a truncate of the log. The plugins + * of the harness that opens the log first (its root thread) decide. With + * one type in two plugins, the larger count wins. + * + * @example + * ```ts + * definePlugin({ name: 'acme/steps', keepRecords: { 'acme/steps:step': 100 } }) + * ``` + */ + keepRecords?: Readonly> } /** A plugin that is set up for each session (default) or for each turn. */ @@ -385,6 +487,7 @@ const SETUP_KEYS: ReadonlySet = new Set< 'commands', 'config', 'contribute', + 'onTurnStart', 'discoverTools', 'prepareTools', 'session', @@ -468,16 +571,50 @@ interface Owned { owner: string } +/** The contributions that `withoutOwners` filters, each with its owner. */ +interface OwnedLists { + /** The `'system'` and the `'context'` prompts. */ + prompts: Array> + middleware: Array> + agentMiddleware: Array> + generationMiddleware: Array> + adapters: Array>> + /** Root agents. `ctx.agents.set` and `delete` change them after setup. */ + agents: Array> + subagents: Array> +} + +/** The values of `list`, without the owners. */ +const valuesIn = (list: ReadonlyArray>) => + list.map((entry) => entry.value) + +/** The `'system'` prompts of `prompts`, then the `'context'` prompts. */ +function splitPrompts(prompts: ReadonlyArray>) { + return { + system: prompts.filter((entry) => entry.value.placement !== 'context'), + context: prompts.filter((entry) => entry.value.placement === 'context'), + } +} + /** Everything a set of mounted plugins contributes, plus its cleanup. */ export interface MountedPlugins { tools: Array - /** Strings, or functions to call for each turn. */ + /** The `'system'` prompts: strings, or functions to call for each turn. */ prompts: Array string)> + /** The `'context'` prompts, with their owners. They go last in the system prompts. */ + contextPrompts: Array<{ + text: string | ((turn: TurnInfo) => string) + owner: string + }> commands: Map config: Map /** Contributions to extension points, by point name. */ extensions: Map> adapters: Array> + turnStarts: Array<{ + start: NonNullable + owner: string + }> discoverers: Array<{ discover: NonNullable owner: string @@ -499,12 +636,10 @@ export interface MountedPlugins { provides: Array }> tools: Array<{ name: string; owner: string }> - /** The owner of each item of `prompts`, at the same index. */ + /** The owner of each `'system'` prompt, then of each `'context'` prompt. */ prompts: Array<{ id: string; owner: string }> - /** The owner of each item of `middleware`, at the same index. */ - middleware: Array - /** The owner of each item of `adapters`, at the same index. */ - adapters: Array + /** The plugin that provided each capability of `values`. */ + capabilities: Array<{ name: string; owner: string }> } middleware: Array /** Chat middleware for every agent run. See `PluginContributions`. */ @@ -519,6 +654,16 @@ export interface MountedPlugins { values: CapabilityValues /** Dispose every plugin scope, newest first. Safe to call twice. */ dispose: () => Promise + /** + * A copy without everything that the plugins in `names` contributed: tools, + * prompts (also the `'context'` prompts), every kind of middleware, agents, subagents, commands, config, + * adapter picks, discoverers, preparers, extension items, and capabilities. + * The copy reads the commands and agents of this set at each use, so a + * later `ctx.agents.set` or `ctx.commands.set` of a kept plugin reaches + * it. It shares `dispose` with this set. It throws when + * a kept plugin requires a capability that a removed plugin provides. + */ + withoutOwners: (names: ReadonlyArray) => MountedPlugins } export interface MountEnvironment { @@ -550,6 +695,11 @@ export interface MountEnvironment { hostPlugin?: (plugin: HostPluginDefinition) => Promise /** The harness name a host plugin gets without `hostPlugin`. */ harnessName?: string + /** + * The plugins whose extension items `collect` hides now: the plugins that + * the thread removed. Read at each `collect` read. + */ + hiddenOwners?: () => ReadonlySet } const unavailable = (what: string) => () => { @@ -587,6 +737,7 @@ const NO_SESSION: PluginSessionApi = { send: unavailable('ctx.session'), note: unavailable('ctx.session'), transcript: unavailable('ctx.session'), + keptRecords: unavailable('ctx.session'), replaceTranscript: unavailable('ctx.session'), entry: unavailable('ctx.session'), updateEntry: unavailable('ctx.session'), @@ -621,6 +772,7 @@ const NO_SERVICES: PluginServices = { start: unavailable('ctx.agents.start'), group: unavailable('ctx.agents.group'), }, + modelChat: unavailable('ctx.model'), } /** @@ -639,15 +791,18 @@ export interface RevertFilesHandler { /** * Put back the files that the tool calls of `hidden` changed. Other files * stay as they are. `records` reads the host records of one type from the - * session log. Resolves to what `unrevert` needs (JSON), or `undefined` - * when no file changed. + * session log. `run` is a run of the revert in the thread of the session, + * so a workspace that is not open on this host opens again from the log. + * Resolves to what `unrevert` needs (JSON), or `undefined` when no file + * changed. */ revert: ( hidden: ReadonlyArray, records: (type: string) => Promise>>, + run: RunContext, ) => Promise - /** Put the files back as they were before `revert`. */ - unrevert: (saved: unknown) => Promise + /** Put the files back as they were before `revert`. `run` as in `revert`. */ + unrevert: (saved: unknown, run: RunContext) => Promise } /** @internal See {@link RevertFilesHandler}. */ @@ -725,8 +880,51 @@ export interface WorkspaceLogApi { * running, with a backend id, and that do not run on this host. */ lostJobs: (threadId: string) => Array<{ jobId: string; backendId: string }> + /** + * The workspace of the run `operationId` is gone. The session sends + * `harness.stranded`. `hostId`: the host of a folder binding. + */ + stranded: (input: { operationId: string; hostId?: string }) => void } +/** + * @internal Tells first-party plugins that a thread ended for good: + * `host.sessions.delete` purges it. A session close is not the end of a + * thread (the next host reattaches it), so this is the only place to free + * what lives as long as the thread, like an owned sandbox or the shadow + * repository of the snapshots. A plugin adds its handler in its setup, also + * in the `session` part of a host plugin. The session provides it. The + * package does not export it. + */ +export const ThreadEnds = createCapability()( + 'tanstack/thread-ends', +) + +/** @internal See {@link ThreadEnds}. */ +export interface ThreadEndsApi { + /** Add `handler`. Gives a function that removes it. */ + add: (handler: ThreadEndHandler) => () => void +} + +/** + * @internal A handler of {@link ThreadEnds}. `run` is a run of the purge in + * the thread, for a `reattach` that needs one. + */ +export type ThreadEndHandler = (thread: { + tenantId: string | undefined + threadId: string + run: RunContext +}) => Promise + +/** + * @internal Takes the secrets of the session out of `text`, as the session + * does for tool results. A plugin that writes tool output to a file uses it. + * The session provides it. The package does not export it. + */ +export const SessionScrub = createCapability<(text: string) => string>()( + 'tanstack/session-scrub', +) + /** Capability values provided by plugins, keyed by handle. */ export class CapabilityValues { readonly context = { capabilities: new CapabilityRegistry() } @@ -748,6 +946,131 @@ export class CapabilityValues { all(): Array { return [...(this.parent?.all() ?? []), ...this.handles] } + + /** A copy without the `removed` handles. The parent stays the same. */ + without(removed: ReadonlySet) { + const copy = new CapabilityValues(this.parent) + for (const handle of this.handles) { + if (!removed.has(handle)) copy.provide(handle, this.get(handle)) + } + return copy + } +} + +/** + * Chat middleware that provides every capability of `values` to each chat + * run, so any chat middleware can read it with `getX(ctx)`. + */ +function capabilityBridgeOf(values: CapabilityValues) { + const handles = values.handles + if (handles.length === 0) return undefined + const bridge: AnyChatMiddleware = { + name: 'harness:plugin-capabilities', + provides: handles, + setup(ctx) { + for (const handle of handles) handle[1](ctx, values.get(handle)) + }, + } + return bridge +} + +/** What `withoutOwners` needs to know about a mount. */ +interface MountRecord { + plugins: ReadonlyArray + /** The plugin that provided each capability. */ + capabilityOwners: ReadonlyMap + owned: OwnedLists +} + +/** `mounted` without what the plugins in `names` contributed. */ +function withoutOwnersOf( + mounted: MountedPlugins, + record: MountRecord, + names: ReadonlyArray, +) { + const removed = new Set(names) + const isKept = (owner: string | undefined) => + owner === undefined || !removed.has(owner) + const isKeptEntry = (entry: { owner: string }) => isKept(entry.owner) + const plugins = record.plugins.filter((plugin) => isKept(plugin.name)) + for (const plugin of plugins) { + for (const handle of plugin.requires ?? []) { + const owner = record.capabilityOwners.get(handle) + if (isKept(owner)) continue + throw new Error( + `Cannot remove plugin ${owner}: plugin ${plugin.name} requires its capability "${handle.capabilityName}".`, + ) + } + } + const removedHandles = new Set() + for (const [handle, owner] of record.capabilityOwners) { + if (!isKept(owner)) removedHandles.add(handle) + } + const values = mounted.values.without(removedHandles) + const { owners } = mounted + const keptTools = new Set( + owners.tools.filter(isKeptEntry).map((entry) => entry.name), + ) + const extensions = new Map>() + for (const [point, items] of mounted.extensions) { + extensions.set(point, items.filter(isKeptEntry)) + } + const { owned } = record + const kept: OwnedLists = { + prompts: owned.prompts.filter(isKeptEntry), + middleware: owned.middleware.filter(isKeptEntry), + agentMiddleware: owned.agentMiddleware.filter(isKeptEntry), + generationMiddleware: owned.generationMiddleware.filter(isKeptEntry), + adapters: owned.adapters.filter(isKeptEntry), + // Agents are read at each use: a kept plugin can change them later + // (`ctx.agents.set`), and this thread must see that. + get agents() { + return owned.agents.filter(isKeptEntry) + }, + get subagents() { + return owned.subagents.filter(isKeptEntry) + }, + } + const next: MountedPlugins = { + tools: mounted.tools.filter((tool) => keptTools.has(tool.name)), + prompts: splitPrompts(kept.prompts).system.map((entry) => entry.value.text), + contextPrompts: mounted.contextPrompts.filter(isKeptEntry), + middleware: valuesIn(kept.middleware), + agentMiddleware: valuesIn(kept.agentMiddleware), + generationMiddleware: valuesIn(kept.generationMiddleware), + // Commands are read at each use, as agents are (`ctx.commands.set`). + get commands() { + return new Map( + [...mounted.commands].filter(([, entry]) => isKeptEntry(entry)), + ) + }, + config: new Map( + [...mounted.config].filter(([, entry]) => isKeptEntry(entry)), + ), + extensions, + adapters: valuesIn(kept.adapters), + turnStarts: mounted.turnStarts.filter(isKeptEntry), + discoverers: mounted.discoverers.filter(isKeptEntry), + preparers: mounted.preparers.filter(isKeptEntry), + get agents() { + return valuesIn(kept.agents) + }, + get subagents() { + return valuesIn(kept.subagents) + }, + owners: { + plugins: owners.plugins.filter((plugin) => isKept(plugin.name)), + tools: owners.tools.filter(isKeptEntry), + prompts: owners.prompts.filter(isKeptEntry), + capabilities: owners.capabilities.filter(isKeptEntry), + }, + capabilityBridge: capabilityBridgeOf(values), + values, + dispose: mounted.dispose, + withoutOwners: (more) => + withoutOwnersOf(next, { ...record, plugins, owned: kept }, more), + } + return next } function checkCapabilityOrder( @@ -843,32 +1166,65 @@ export async function mountPlugins( value: tool, owner: 'the harness', })) - const prompts: Array> = [] - const middleware: Array = [] - const middlewareOwners: Array = [] - const adapterOwners: Array = [] - const agentMiddleware: Array = [] - const generationMiddleware: Array = [] + const owned: OwnedLists = { + prompts: [], + middleware: [], + agentMiddleware: [], + generationMiddleware: [], + adapters: [], + agents: [], + subagents: [], + } + const { prompts, agents: roots, subagents: subs } = owned + const capabilityOwners = new Map() 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: MountedPlugins['adapters'] = [] + const turnStarts: MountedPlugins['turnStarts'] = [] const discoverers: MountedPlugins['discoverers'] = [] const preparers: MountedPlugins['preparers'] = [] - const agents: Array = [] - const subagents: Array = [] const services = env.services ?? NO_SERVICES - const remove = (name: string, list: Array) => { - const at = list.findIndex((item) => item.name === name) + const indexIn = (name: string, list: Array>) => + list.findIndex((entry) => entry.value.name === name) + const remove = (name: string, list: Array>) => { + const at = indexIn(name, list) if (at >= 0) list.splice(at, 1) } // `ctx.commands.ready`: resolves after the last plugin is set up. A failed // mount leaves it pending, so no plugin acts on a session that never opens. let markReady = () => {} const ready = new Promise((resolve) => (markReady = resolve)) + // The shown values of each extension point, built again only when it grows + // or when the hidden owners change. Contributions are only added, so the + // length tells when a copy is old. `hiddenOwners` gives a new set when the + // thread removes other plugins, so a different set also makes it old. + const snapshots = new Map< + string, + { + length: number + hidden: ReadonlySet | undefined + values: Array + } + >() + const valuesOf = (point: string) => { + const items = extensions.get(point) ?? [] + const hidden = env.hiddenOwners?.() + const cached = snapshots.get(point) + if (cached?.length === items.length && cached.hidden === hidden) { + return cached.values + } + const shown = hidden?.size + ? items.filter((item) => !hidden.has(item.owner)) + : items + const fresh = shown.map((item) => item.value) + snapshots.set(point, { length: items.length, hidden, values: fresh }) + return fresh + } + // True after the last setup. No plugin contributes after that. + let settled = false const ownerOf = (name: string) => commands.get(name)?.owner ?? env.takenCommands?.get(name)?.owner @@ -909,6 +1265,7 @@ export async function mountPlugins( } values.provide(handle, value) provided.add(handle) + capabilityOwners.set(handle, plugin.name) }, agents: { list: () => env.registry.list(), @@ -919,30 +1276,33 @@ export async function mountPlugins( group: services.agents.group, set: (agent, options) => { env.registry.set(agent, plugin.name) - const list = options?.subagent ? subagents : agents - remove(agent.name, options?.subagent ? agents : subagents) - const at = list.findIndex((item) => item.name === agent.name) - list.splice(at < 0 ? list.length : at, 1, agent) + const list = options?.subagent ? subs : roots + remove(agent.name, options?.subagent ? roots : subs) + const at = indexIn(agent.name, list) + list.splice(at < 0 ? list.length : at, 1, { + value: agent, + owner: plugin.name, + }) }, delete: (name) => { env.registry.delete(name, plugin.name) // Still there: another owner has it. if (env.registry.get(name)) return - remove(name, agents) - remove(name, subagents) + remove(name, roots) + remove(name, subs) }, }, collect: (point: ExtensionPoint): ReadonlyArray => { - let items = extensions.get(point.name) - if (!items) { - items = [] - extensions.set(point.name, items) - } - const source = items - const live = () => source.map((item) => item.value as T) - // A live view: items contributed after this call appear too. Array - // methods like `filter` also check that an index exists, so every - // trap reads the items, not the empty target. + const name = point.name + if (!extensions.has(name)) extensions.set(name, []) + // The values are the `T` items that plugins gave for this point, + // without the items of the plugins that the thread removed. + const live = () => valuesOf(name) as Array + // After setup, no plugin contributes: give the array itself. + if (settled) return live() + // During setup, a live view: items contributed after this call + // appear too. Array methods like `filter` also check that an index + // exists, so every trap reads the items, not the empty target. return new Proxy>([], { get: (_target, key) => Reflect.get(live(), key), has: (_target, key) => Reflect.has(live(), key), @@ -953,11 +1313,17 @@ export async function mountPlugins( }, emit: (event, value) => services.emit(plugin.name, event.name, value), on: (event, handler) => { - const off = services.on(event.name, (value) => + const stop = services.on(event.name, (value) => handler(value as never), ) // The listener goes away with the plugin, for example on a reload. - scope.signal.addEventListener('abort', off) + // `off` also removes this abort listener, so many `on` and `off` + // calls do not add up listeners on the plugin signal. + const off = () => { + scope.signal.removeEventListener('abort', off) + stop() + } + scope.signal.addEventListener('abort', off, { once: true }) return off }, state: (initial) => services.state(plugin.name, initial), @@ -987,6 +1353,21 @@ export async function mountPlugins( }, ready, }, + model: { + chat: async ({ run = env.turn?.run, ...call }) => { + if (!run) { + throw new Error( + `ctx.model.chat of ${plugin.name} needs a run. In a tool or a middleware hook, pass run: ctx.context.run. In a command, pass the run of its context.`, + ) + } + // The session gives back what chat() gave for these options. + return (await services.modelChat({ + ...call, + run, + source: `plugin:${plugin.name}`, + })) as never + }, + }, ...(env.turn ? { turn: env.turn } : {}), } let parts: Array @@ -1029,12 +1410,14 @@ export async function mountPlugins( } prompts.push({ value: section, owner: plugin.name }) }) - for (const item of contributions.middleware ?? []) { - middleware.push(item) - middlewareOwners.push(plugin.name) - } - agentMiddleware.push(...(contributions.agentMiddleware ?? [])) - generationMiddleware.push(...(contributions.generationMiddleware ?? [])) + const own = (value: T) => ({ value, owner: plugin.name }) + owned.middleware.push(...(contributions.middleware ?? []).map(own)) + owned.agentMiddleware.push( + ...(contributions.agentMiddleware ?? []).map(own), + ) + owned.generationMiddleware.push( + ...(contributions.generationMiddleware ?? []).map(own), + ) // The session reads commands and config only from its session plugins. const isSessionOnly = Object.keys(contributions.commands ?? {}).length > 0 || @@ -1070,9 +1453,13 @@ export async function mountPlugins( } items.push({ value: item.value, owner: plugin.name }) } - if (contributions.adapter) { - adapters.push(contributions.adapter) - adapterOwners.push(plugin.name) + if (contributions.adapter) + owned.adapters.push(own(contributions.adapter)) + if (contributions.onTurnStart) { + turnStarts.push({ + start: contributions.onTurnStart, + owner: plugin.name, + }) } if (contributions.discoverTools) { discoverers.push({ @@ -1088,11 +1475,11 @@ export async function mountPlugins( } for (const agent of contributions.agents ?? []) { env.registry.add(agent, plugin.name) - agents.push(agent) + roots.push(own(agent)) } for (const agent of contributions.subagents ?? []) { env.registry.add(agent, plugin.name) - subagents.push(agent) + subs.push(own(agent)) } } } @@ -1118,38 +1505,38 @@ export async function mountPlugins( 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 - + settled = true markReady() let disposed: Promise | undefined - return { + const record: MountRecord = { plugins, capabilityOwners, owned } + const { system: systemPrompts, context: contextPrompts } = + splitPrompts(prompts) + const mounted: MountedPlugins = { tools: tools .filter((entry) => entry.owner !== 'the harness') .map((entry) => entry.value), - prompts: prompts.map((entry) => entry.value.text), - middleware, - agentMiddleware, - generationMiddleware, + prompts: systemPrompts.map((entry) => entry.value.text), + contextPrompts: contextPrompts.map((entry) => ({ + text: entry.value.text, + owner: entry.owner, + })), + middleware: valuesIn(owned.middleware), + agentMiddleware: valuesIn(owned.agentMiddleware), + generationMiddleware: valuesIn(owned.generationMiddleware), commands, config, extensions, - adapters, + adapters: valuesIn(owned.adapters), + turnStarts, discoverers, preparers, - agents, - subagents, + // Read at each use: `ctx.agents.set` and `delete` change them later. + get agents() { + return valuesIn(roots) + }, + get subagents() { + return valuesIn(subs) + }, owners: { plugins: plugins.map((plugin) => ({ name: plugin.name, @@ -1165,15 +1552,19 @@ export async function mountPlugins( name: entry.value.name, owner: entry.owner, })), - prompts: prompts.map((entry) => ({ + prompts: [...systemPrompts, ...contextPrompts].map((entry) => ({ id: entry.value.id, owner: entry.owner, })), - middleware: middlewareOwners, - adapters: adapterOwners, + capabilities: [...capabilityOwners].map(([handle, owner]) => ({ + name: handle.capabilityName, + owner, + })), }, - capabilityBridge, + capabilityBridge: capabilityBridgeOf(values), values, dispose: () => (disposed ??= disposeAll(scopes)), + withoutOwners: (names) => withoutOwnersOf(mounted, record, names), } + return mounted } diff --git a/packages/ai-harness/src/protocol.ts b/packages/ai-harness/src/protocol.ts index 1610500e9b..62a01f21fe 100644 --- a/packages/ai-harness/src/protocol.ts +++ b/packages/ai-harness/src/protocol.ts @@ -1,26 +1,48 @@ import { convertSchemaToJsonSchema } from '@tanstack/ai' import { AccessDeniedError } from './access' +import { parseTraceparent } from './observability' import { isRecord } from './utils' import { createInputId, errorText } from './session/helpers' -import type { StreamChunk } from '@tanstack/ai' +import { storeInlineMedia } from './inline-media' +import { SessionImpl } from './session/harness-session' +import { mediaIdOf } from './media-ref' +import type { ModelMessage, StreamChunk } from '@tanstack/ai' import type { AnyAgent } from './agents' import type { AnyHarness } from './define' +import type { InlineMediaContext } from './inline-media' import type { HarnessSession, SessionSnapshot } from './session/index' +import type { AgentRunStartOptions } from './session/types' import type { Cursor, HarnessInput, Operation, Principal, Receipt, + UserInput, } from './types' /** The session-tier protocol version. Sent in `subscribe` and `hello`. */ -export const HARNESS_PROTOCOL_VERSION = 1 +export const HARNESS_PROTOCOL_VERSION = 2 + +/** What failed, in a `harness.error` frame. */ +export type HarnessErrorCode = + | 'unauthorized' + | 'forbidden' + | 'invalid_frame' + | 'unsupported_version' + | 'slow_consumer' + | 'internal' /** Client to host. */ export type ControlFrame = | { type: 'harness.subscribe'; threadId: string; from?: Cursor; v?: number } - | { type: 'harness.input'; requestId: string; input: HarnessInput } + | { + type: 'harness.input' + requestId: string + input: HarnessInput + /** A W3C `traceparent` of the caller. An invalid value is dropped. */ + traceparent?: string + } | { type: 'harness.snapshot' } /** Host to client. */ @@ -37,11 +59,25 @@ export type HostFrame = | { type: 'harness.event' cursor: Cursor + /** The thread of the event: the thread, or a child thread of its tree. */ + threadId: string operationId: string event: StreamChunk } | { type: 'harness.snapshot'; snapshot: SessionSnapshot } - | { type: 'harness.error'; message: string } + | { + type: 'harness.error' + /** A short reason, for people. Read `code` in code. */ + message: string + code: HarnessErrorCode + /** + * The `requestId` of the frame or the request that failed. A frame or + * a request with no id gets an id from the server. + */ + requestId: string + } + /** A heartbeat. It keeps idle connections open. Ignore it. */ + | { type: 'harness.ping' } const INPUT_OPS = new Set([ 'prompt', @@ -66,6 +102,31 @@ const INPUT_OPS = new Set([ 'agentMessage', ]) +const PART_TYPES = new Set(['text', 'image', 'audio', 'video', 'document']) + +/** + * True for a content part with the fields the harness reads: a string + * `content` for text, else a `source` with a string `type` and `value`, and + * a string `mimeType` for a `data` source. + */ +function isContentPart(value: unknown) { + if (!isRecord(value) || typeof value.type !== 'string') return false + if (!PART_TYPES.has(value.type)) return false + if (value.type === 'text') return typeof value.content === 'string' + const { source } = value + if (!isRecord(source) || typeof source.type !== 'string') return false + if (typeof source.value !== 'string') return false + return source.type !== 'data' || typeof source.mimeType === 'string' +} + +/** True for a string, or an array of content parts. */ +function isMessageContent(value: unknown) { + return ( + typeof value === 'string' || + (Array.isArray(value) && value.every(isContentPart)) + ) +} + /** Check the shape of a client input. Throws with a short reason. */ export function parseHarnessInput(value: unknown): HarnessInput { if ( @@ -80,12 +141,10 @@ export function parseHarnessInput(value: unknown): HarnessInput { value.op === 'steer' || value.op === 'followUp' || value.op === 'agentMessage' - if ( - needsMessage && - typeof value.message !== 'string' && - !Array.isArray(value.message) - ) { - throw new Error(`Invalid input: ${value.op} needs a message.`) + if (needsMessage && !isMessageContent(value.message)) { + throw new Error( + `Invalid input: ${value.op} needs a message: text, or content parts.`, + ) } if ( value.op === 'prompt' && @@ -173,7 +232,9 @@ export function parseHarnessInput(value: unknown): HarnessInput { // A client sends user messages only. An assistant or tool message from a // client could fake an answer or a tool result in the model context. const isUserMessage = (message: unknown) => - isRecord(message) && message.role === 'user' + isRecord(message) && + message.role === 'user' && + (message.content === null || isMessageContent(message.content)) if ( (value.op === 'prompt' || value.op === 'continue') && value.ephemeral !== undefined && @@ -200,10 +261,14 @@ export function parseControlFrame(data: string): ControlFrame { value.type === 'harness.subscribe' && typeof value.threadId === 'string' ) { + if (value.v !== undefined && typeof value.v !== 'number') { + throw new Error('Invalid frame: v must be a number.') + } return { type: 'harness.subscribe', threadId: value.threadId, ...(typeof value.from === 'string' ? { from: value.from } : {}), + ...(value.v === undefined ? {} : { v: value.v }), } } if (value.type === 'harness.input' && typeof value.requestId === 'string') { @@ -211,6 +276,11 @@ export function parseControlFrame(data: string): ControlFrame { type: 'harness.input', requestId: value.requestId, input: parseHarnessInput(value.input), + // As the W3C rules say, an invalid `traceparent` is dropped. + ...(typeof value.traceparent === 'string' && + parseTraceparent(value.traceparent) + ? { traceparent: value.traceparent } + : {}), } } if (value.type === 'harness.snapshot') return { type: 'harness.snapshot' } @@ -227,6 +297,14 @@ export function parseControlFrame(data: string): ControlFrame { * its credentials. Rejects with `AccessDeniedError` when the host refuses * the input. * + * Media in a message, or in an `ephemeral` message, is checked first, as + * `storeInlineMedia` does: inline bytes go to the media store and the part + * becomes a `harness-media:` reference. An external URL needs + * `media.externalUrls`. A refused part gives a rejected receipt with the + * reason, for example `media_too_large`. An input with an `inputId` that the + * session has already sends the stored message again, so a retry stores no + * media. + * * @example * ```ts * const receipt = await applyInput({ harness, session, input, principal }) @@ -237,14 +315,18 @@ export async function applyInput({ session, input, principal, + traceparent, }: { harness: AnyHarness session: HarnessSession input: HarnessInput principal?: Principal + /** A W3C `traceparent` of the caller, for a message to the session. */ + traceparent?: string }): Promise { // The log keeps who sent each input, for every input kind. const user = principal ? { principal } : {} + const traced = traceparent === undefined ? {} : { traceparent } const id = { ...(input.inputId === undefined ? {} : { inputId: input.inputId }), ...user, @@ -260,29 +342,59 @@ export async function applyInput({ status: 'rejected', reason: 'not_exposed', } + const refused = (reason: string) => { + const receipt: Receipt = { + inputId: input.inputId ?? '', + status: 'rejected', + reason, + } + return receipt + } + const media: InlineMediaContext = { + options: harness.media, + putMedia: (upload) => session.putMedia(upload), + } + /** `message` with its inline media stored. A retry stores none again. */ + const checkedMessage = async (message: UserInput) => { + const retry = retriedMessage(session, input.inputId, message) + if (retry !== undefined) return { content: retry, isRetry: true } + if (typeof message === 'string') return { content: message } + const checked = await storeInlineMedia(message, media) + return 'refused' in checked ? checked : { content: checked.parts } + } switch (input.op) { - case 'prompt': + case 'prompt': { + const message = await checkedMessage(input.message) + if ('refused' in message) return refused(message.refused) + // A retry drops its ephemeral messages, as `session.send` does. + const ephemeral = message.isRetry + ? { options: {} } + : await checkedEphemeral(input.ephemeral, media) + if ('refused' in ephemeral) return refused(ephemeral.refused) return turnReceipt( - session.send(input.message, { + session.send(message.content, { ...(input.busy ? { busy: input.busy } : {}), - ...(input.ephemeral ? { ephemeral: input.ephemeral } : {}), + ...ephemeral.options, // The host checks the `operator` action of the sender. ...(input.as ? { as: input.as } : {}), ...sent, + ...traced, }), ) - case 'continue': - return turnReceipt( - session.continue({ - ...(input.ephemeral ? { ephemeral: input.ephemeral } : {}), - ...sent, - }), - ) + } + case 'continue': { + const ephemeral = await checkedEphemeral(input.ephemeral, media) + if ('refused' in ephemeral) return refused(ephemeral.refused) + return turnReceipt(session.continue({ ...ephemeral.options, ...sent })) + } case 'steer': - case 'followUp': + case 'followUp': { + const message = await checkedMessage(input.message) + if ('refused' in message) return refused(message.refused) return turnReceipt( - session.send(input.message, { ...sent, mode: input.op }), + session.send(message.content, { ...sent, ...traced, mode: input.op }), ) + } case 'resolve': return session.resolve({ resume: input.resume, ...id }) case 'cancel': @@ -290,9 +402,11 @@ export async function applyInput({ // It runs no new code, so it needs no `expose` entry, like `cancel`. It // stores no input, so there is no sender to keep. case 'background': - return session.background(input.toolCallId) + return session.background( + input.toolCallId !== undefined ? { toolCallId: input.toolCallId } : {}, + ) case 'cancelInput': - return session.cancelInput(input.inputId, user) + return session.cancelInput({ inputId: input.inputId, ...user }) case 'setDelivery': return session.setDelivery({ inputId: input.inputId, @@ -304,9 +418,12 @@ export async function applyInput({ if (!(harness.expose?.commands ?? []).includes(input.name)) { return notExposed } + // The client's id makes a resent command a duplicate, so it runs once. + const inputId = input.inputId ?? createInputId() const operation = session.command({ name: input.name, input: input.input, + inputId, ...user, }) // ponytail: a command has no receipt, so it is accepted at once. A @@ -315,11 +432,7 @@ export async function applyInput({ () => {}, () => {}, ) - return { - inputId: operation.id, - status: 'accepted', - operationId: operation.id, - } + return { inputId, status: 'accepted', operationId: operation.id } } case 'answer': return session.answer({ @@ -352,11 +465,16 @@ export async function applyInput({ (record) => types.includes(record.type) && !('thread' in record), ) if (!isExposed) return notExposed - const inputId = input.inputId ?? createInputId() + // The client's id makes a resent append a duplicate that writes + // nothing. try { - await session.append(input.records) - return { inputId, status: 'accepted' } + return await session.append({ + records: input.records, + ...(input.inputId !== undefined ? { inputId: input.inputId } : {}), + ...user, + }) } catch (error) { + const inputId = input.inputId ?? createInputId() return { inputId, status: 'rejected', reason: errorText(error) } } } @@ -376,7 +494,7 @@ export async function applyInput({ case 'unrevert': if (!(harness.expose?.commands ?? []).includes('undo')) return notExposed return input.op === 'revert' - ? session.revert(input.messageId) + ? session.revert({ messageId: input.messageId }) : session.unrevert() case 'agent': { const exposed = (harness.expose?.agents ?? []).includes(input.agent) @@ -390,15 +508,16 @@ export async function applyInput({ // The run is the client's: its credentials and keys are theirs. // ponytail: accepted at once, as a command. A run that the host // refuses fails with a `RUN_ERROR` event. - const operation = handle.start(input.input, { + // The client's id makes a resent start a duplicate, so the agent runs + // once. + const inputId = input.inputId ?? createInputId() + const options: AgentRunStartOptions = { wake: input.detached === true, + inputId, ...user, - }) - return { - inputId: operation.id, - status: 'accepted', - operationId: operation.id, } + const operation = handle.start(input.input, options) + return { inputId, status: 'accepted', operationId: operation.id } } case 'agentMessage': { const run = session.agentRun(input.operationId) @@ -418,8 +537,10 @@ export async function applyInput({ reason: 'not_exposed', } } + const message = await checkedMessage(input.message) + if ('refused' in message) return refused(message.refused) return turnReceipt( - session.send(input.message, { + session.send(message.content, { to: input.operationId, ...id, ...(input.mode ? { mode: input.mode } : {}), @@ -429,6 +550,57 @@ export async function applyInput({ } } +/** + * `message` as a retry of the input `inputId` that `session` has in its log, + * or `undefined` for a new input. A part at the place of a stored + * `harness-media:` part, of the same type, becomes the stored part. So a + * retry stores no media again, and the duplicate check of the session + * matches. A retry that differs in another way gets `conflict`. Its parts are + * not checked or stored: the session refuses it before it runs. + */ +// ponytail: a retry with other inline bytes of the same type matches too. +// The first input stands. +export function retriedMessage( + session: HarnessSession, + inputId: string | undefined, + message: UserInput, +): UserInput | undefined { + if (inputId === undefined || !(session instanceof SessionImpl)) { + return undefined + } + const known = session.knownMessage(inputId) + if (known === undefined) return undefined + if (typeof message === 'string' || typeof known === 'string') return message + return message.map((part, index) => { + const stored = known[index] + const isStoredMedia = + stored?.type === part.type && mediaIdOf(stored) !== undefined + return isStoredMedia ? stored : part + }) +} + +/** + * The `ephemeral` option of a turn, with the inline media of each message + * stored, or why a part was refused. No messages give no option. + */ +async function checkedEphemeral( + messages: ReadonlyArray | undefined, + media: InlineMediaContext, +) { + if (!messages) return { options: {} } + const ephemeral: Array = [] + for (const message of messages) { + if (!Array.isArray(message.content)) { + ephemeral.push(message) + continue + } + const checked = await storeInlineMedia(message.content, media) + if ('refused' in checked) return checked + ephemeral.push({ ...message, content: checked.parts }) + } + return { options: { ephemeral } } +} + /** * The receipt of a turn that a client started. Nobody may await the turn, * so its failure is handled here. A turn that the host refused throws its @@ -468,8 +640,20 @@ export function describeForClient( } } -/** The AG-UI capabilities document of a harness. */ -export function capabilitiesOf(harness: AnyHarness) { +/** + * The AG-UI capabilities document of a harness. Pass `socket: true` when the + * app also serves `handleHarnessSocket`, so `transport.websocket` tells + * clients that they can use a WebSocket. Default `false`. + * + * @example + * ```ts + * const capabilities = capabilitiesOf(assistant, { socket: true }) + * ``` + */ +export function capabilitiesOf( + harness: AnyHarness, + { socket = false }: { socket?: boolean } = {}, +) { const exposed = new Set(harness.expose?.agents ?? []) const subagents: ReadonlyArray = harness.subagents?.agents ?? [] const agents: ReadonlyArray = [ @@ -478,7 +662,7 @@ export function capabilitiesOf(harness: AnyHarness) { ] return { identity: { name: harness.name, type: 'tanstack-ai-harness' }, - transport: { streaming: true, websocket: true }, + transport: { streaming: true, websocket: socket }, tools: { supported: true, items: (harness.tools ?? []).map((tool) => ({ @@ -512,3 +696,36 @@ export function capabilitiesOf(harness: AnyHarness) { }, } } + +/** + * The `unsupported_version` error of a `harness.subscribe` frame or a + * `?v=` query that asks for a version other than + * {@link HARNESS_PROTOCOL_VERSION}. No version asks for the current one. + * Send the error frame, then close the stream. + */ +export function versionError(requested: number | undefined) { + if (requested === undefined || requested === HARNESS_PROTOCOL_VERSION) { + return undefined + } + return { + message: `Protocol version ${requested} is not supported. This host speaks version ${HARNESS_PROTOCOL_VERSION}.`, + code: 'unsupported_version' as const, + } +} + +/** + * The `requestId` of the raw text of a client frame, or `undefined`. It + * reads only that field, so the error frame of a frame that fails the full + * parse (for example, with a bad `input`) still names its request. It never + * throws. + */ +export function requestIdOf(data: string): string | undefined { + try { + const value: unknown = JSON.parse(data) + return isRecord(value) && typeof value.requestId === 'string' + ? value.requestId + : undefined + } catch { + return undefined + } +} diff --git a/packages/ai-harness/src/ring.ts b/packages/ai-harness/src/ring.ts new file mode 100644 index 0000000000..a1f0e20253 --- /dev/null +++ b/packages/ai-harness/src/ring.ts @@ -0,0 +1,150 @@ +import type { SessionEvent } from './types' + +/** + * A FIFO queue with O(1) push and amortized O(1) shift. A shift moves a head + * index. The array is cut only when the dropped half is the bigger half, so + * each item is moved at most once per cut. + */ +class Queue { + private items: Array = [] + private head = 0 + + get length() { + return this.items.length - this.head + } + + /** The item at `index`, counted from the oldest. */ + at(index: number) { + return this.items[this.head + index] as T + } + + push(item: T) { + this.items.push(item) + } + + shift() { + if (this.length === 0) return undefined + const item = this.items[this.head] as T + this.items[this.head] = undefined + this.head += 1 + if (this.head * 2 >= this.items.length) { + this.items = this.items.slice(this.head) + this.head = 0 + } + return item + } +} + +/** One kept event: the shape of the in-memory tail of the shared log. */ +export interface RingEntry { + thread: string + entry: SessionEvent + /** The byte size of the event (its JSON length). */ + size: number +} + +/** The events after a cursor, and whether older events are gone. */ +export interface RingRead { + entries: Array + /** + * `true` when the cursor is older than the oldest kept event: the ring + * dropped events that the reader did not get. The caller can send a gap. + */ + gap: boolean +} + +/** + * The newest events, up to `maxBytes` of them. A push that goes over the cap + * drops the oldest events, but the newest event always stays. Cursors are + * numbers that grow with each push. They do not need to be consecutive. + * + * Push and drop are O(1) (amortized). A read finds its start with a binary + * search, and a read of one thread uses the index of that thread, so it does + * not walk the events of other threads. + */ +export class EventRing { + private readonly all = new Queue() + /** The kept events of each thread, oldest first. */ + private readonly threads = new Map< + string, + Queue + >() + private total = 0 + private readonly maxBytes: number + /** Every event after this cursor is kept. */ + private floor: number + + /** + * `from`: the cursor before the first event the ring gets. A read from an + * older cursor reports a gap. Default `0`. + */ + constructor(options: { maxBytes: number; from?: number }) { + this.maxBytes = options.maxBytes + this.floor = options.from ?? 0 + } + + /** The sum of the sizes of the kept events. */ + get bytes() { + return this.total + } + + /** The number of kept events. */ + get length() { + return this.all.length + } + + /** Every event after this cursor is kept. An older cursor has a gap. */ + get from() { + return this.floor + } + + /** Keep `item`. Its cursor must be bigger than the cursor of each kept event. */ + push(item: RingEntry): void { + const kept = { ...item, cursor: Number(item.entry.cursor) } + this.all.push(kept) + let own = this.threads.get(item.thread) + if (!own) { + own = new Queue() + this.threads.set(item.thread, own) + } + own.push(kept) + this.total += item.size + while (this.total > this.maxBytes && this.all.length > 1) { + // The oldest event of all is also the oldest event of its thread. + const dropped = this.all.shift() + if (!dropped) break + const queue = this.threads.get(dropped.thread) + queue?.shift() + if (queue?.length === 0) this.threads.delete(dropped.thread) + this.total -= dropped.size + this.floor = dropped.cursor + } + } + + /** + * The kept events after `cursor` (exclusive), oldest first. With `thread`, + * only the events of that thread. + */ + after(cursor: number, thread?: string): RingRead { + const queue = thread === undefined ? this.all : this.threads.get(thread) + const entries: Array = [] + if (queue) { + for (let i = firstAfter(queue, cursor); i < queue.length; i++) { + entries.push(queue.at(i).entry) + } + } + return { entries, gap: cursor < this.floor } + } +} + +/** The index of the first item of `queue` with a cursor after `cursor`. */ +function firstAfter(queue: Queue<{ cursor: number }>, cursor: number) { + let low = 0 + let high = queue.length + while (low < high) { + const middle = (low + high) >>> 1 + if (queue.at(middle).cursor <= cursor) low = middle + 1 + else high = middle + } + return low +} diff --git a/packages/ai-harness/src/server.ts b/packages/ai-harness/src/server.ts index 2c5f647615..24845a6fd9 100644 --- a/packages/ai-harness/src/server.ts +++ b/packages/ai-harness/src/server.ts @@ -6,6 +6,7 @@ export type { HarnessSocketOptions, OpenAccess, } from './http' +export type { HarnessStreamOptions } from './sse' export { HARNESS_PROTOCOL_VERSION, diff --git a/packages/ai-harness/src/session/agents.ts b/packages/ai-harness/src/session/agents.ts index 0274e8af0e..daca47c3a2 100644 --- a/packages/ai-harness/src/session/agents.ts +++ b/packages/ai-harness/src/session/agents.ts @@ -79,6 +79,7 @@ import type { AgentRun, AgentSendOptions, ChatTurnResult, + Cursor, InputSettlement, SettlementOutcome, MediaRecord, @@ -123,7 +124,8 @@ export abstract class SessionAgents< * command, a protocol input). `undefined` when no such agent exists. */ agent(name: string): DynamicAgentHandle | undefined { - if (!this.agentRegistry.get(name)) return undefined + // An agent of a plugin that this thread removed does not exist here. + if (!this.threadAgent(name)) return undefined return this.agentHandle(name) } @@ -140,7 +142,14 @@ export abstract class SessionAgents< const follower = new OperationImpl( 'agent', this.feed, - () => this.cancel({ operationId: target }), + () => { + // It acts as the user of the run, as a cancel of the run does. + const principal = this.principalOfRun(target) + return this.cancel({ + operationId: target, + ...(principal ? { principal } : {}), + }) + }, chain?.agent, ) void this.sendToAgent({ @@ -397,7 +406,11 @@ export abstract class SessionAgents< return undefined } const entries = logged.map(({ node, settlement }) => { - const operation = this.agentOperation(first.input.agent, node.runId) + const operation = this.agentOperation( + first.input.agent, + node.runId, + this.runStart(node.runId), + ) this.endAgentRun(operation, settlement) return { operation, @@ -446,7 +459,8 @@ export abstract class SessionAgents< ) { const options = this.harness.media const agentMiddleware = [ - ...(this.sessionPlugins?.agentMiddleware ?? []), + // A plugin that this thread removed gives its agents no middleware. + ...(this.threadPlugins()?.agentMiddleware ?? []), ...(runPlugins?.agentMiddleware ?? []), ] const agents = turn?.agents @@ -510,12 +524,13 @@ export abstract class SessionAgents< }, (inner, ctx) => this.runHarnessChild(inner, ctx, { operation, principal }), + run, ) return { harness } }) return { generationMiddleware: [ - ...(this.sessionPlugins?.generationMiddleware ?? []), + ...(this.threadPlugins()?.generationMiddleware ?? []), ...(runPlugins?.generationMiddleware ?? []), ...mediaCapture({ persistence: this.media, @@ -1088,7 +1103,14 @@ export abstract class SessionAgents< for (const item of this.writer.state.thread.tree.values()) { if (item.chainInputId === first.inputId) node = item } - return chainFields(first, node) + const fields = chainFields(first, node) + // The input keeps only a parent run of agent code. The tree keeps the + // turn that started the chain too, so a run of the chain on any host + // knows that turn. + const parent = node?.parentRunId + return parent !== undefined && fields.parentRunId === undefined + ? { ...fields, turnRunId: parent } + : fields } /** @@ -1138,14 +1160,41 @@ export abstract class SessionAgents< } } - /** A new agent run operation. `id`: a stored one, after a restart. */ - protected agentOperation(agent: string, id?: string): AgentRunImpl { + /** + * The run that started the operation `operationId`: the parent of its + * chain, the same as its node in the run tree. A run that this host does + * not run (an offered run) has its node in the log. A top-level run and an + * operation that is not an agent run have none. + */ + protected parentRunOf(operationId: string): string | undefined { + const chain = this.chainOf.get(operationId) + if (chain) return chain.parentRunId ?? chain.turnRunId + return this.logWriter?.state.thread.tree.get(operationId)?.parentRunId + } + + /** + * A new agent run operation. `id`: a stored one, after a restart. + * `startedCursor`: the start of a run that the log rebuilds. + */ + protected agentOperation( + agent: string, + id?: string, + startedCursor?: Cursor, + ): AgentRunImpl { + // A call on the run itself acts as the user of the run, not as the + // principal that opened the session. A handle passes its own principal. + const user = (running: string) => { + const principal = this.principalOfRun(running) + return principal ? { principal } : {} + } const operation = new OperationImpl( 'agent', this.feed, - (running) => this.cancel({ operationId: running.id }), + (running) => + this.cancel({ operationId: running.id, ...user(running.id) }), agent, id, + startedCursor, ) const run = Object.assign( // `kind` and `agent` are fixed by the constructor call above. @@ -1153,9 +1202,18 @@ export abstract class SessionAgents< { send: ( message: UserInput, - options?: { mode?: 'steer' | 'followUp'; inputId?: string }, + options?: { + mode?: 'steer' | 'followUp' + inputId?: string + principal?: Principal + }, ) => - this.sendToAgent({ ...options, operationId: operation.id, message }), + this.sendToAgent({ + ...user(operation.id), + ...options, + operationId: operation.id, + message, + }), }, ) this.track(run) @@ -1440,7 +1498,13 @@ export abstract class SessionAgents< try { await this.runAgentOnce(operation, chain, start) } catch (error) { - // A store write threw. The run still ends. + // A store write threw. The run still ends. A run that had no receipt + // yet gets a refusal, so a reader of it does not wait. + operation.resolveReceipt({ + inputId: start.inputId, + status: 'rejected', + reason: error instanceof Error ? error.message : String(error), + }) this.failOperation(operation, error) } finally { // After a log failure, another host owns the thread. A closing session @@ -1555,6 +1619,11 @@ export abstract class SessionAgents< // An offered run runs as this operation on the owner. { runId: operation.id }, ) + // As a turn does: a refused, offered, or known input gets its receipt + // now. A new one gets it below, when the run starts or waits. + if (admission !== 'new') { + operation.resolveReceipt(this.duplicateReceipt(inputId, admission)) + } switch (admission) { case 'new': break @@ -1584,15 +1653,23 @@ export abstract class SessionAgents< // A top-level start waits for a slot of `maxConcurrent`. const isTopLevel = chain.parentRunId === undefined && chain.turnRunId === undefined + if (isTopLevel && !this.hasSlot()) { + operation.resolveReceipt({ + inputId, + status: 'queued', + operationId: operation.id, + }) + } if (isTopLevel && !(await this.takeSlot(operation, chain, inputId))) { return } } + // By name, an agent of a plugin that this thread removed is unknown. const agent: AnyAgent | undefined = - typeof target === 'string' ? this.agentRegistry.get(target) : target + typeof target === 'string' ? this.threadAgent(target) : target if (!agent) { this.releaseStart(chain) - this.reject(inputId, 'unknown_agent') + operation.resolveReceipt(this.reject(inputId, 'unknown_agent')) operation.fail('failed', new Error(`Unknown agent: ${name}`)) return } @@ -1607,7 +1684,7 @@ export abstract class SessionAgents< .map((issue) => issue.message) .join(', ')}` this.releaseStart(chain) - this.reject(inputId, 'invalid_input') + operation.resolveReceipt(this.reject(inputId, 'invalid_input')) operation.fail('failed', new Error(reason)) return } @@ -1615,6 +1692,12 @@ export abstract class SessionAgents< } operation.setStatus('running') + // A run that waited for a slot has its receipt already. + operation.resolveReceipt({ + inputId, + status: 'accepted', + operationId: operation.id, + }) // A run that goes on after its waits applied its input already. The // node of the run goes in the same append. if (!answers) { @@ -1661,10 +1744,13 @@ export abstract class SessionAgents< ...storedPrincipal(principal), }) if (answers) await runs.update(operation.id, { status: 'running' }) + // The run that started this one: a reader rebuilds the run tree from it. + const parentRunId = chain.parentRunId ?? chain.turnRunId operation.publish({ type: EventType.RUN_STARTED, runId: operation.id, threadId: this.threadId, + ...(parentRunId !== undefined ? { parentRunId } : {}), timestamp: Date.now(), }) if (!answers) { @@ -1737,6 +1823,7 @@ export abstract class SessionAgents< }, (inner, ctx) => this.runHarnessChild(inner, ctx, { operation, principal }), + run, ) const stream = runAgentStream( agent, @@ -2075,7 +2162,11 @@ export abstract class SessionAgents< const runId = input.operationId if (runId === undefined || run.subagentRunId === undefined) return this.warnLiveOnly(runId) - const operation = this.agentOperation(first.input.agent, runId) + const operation = this.agentOperation( + first.input.agent, + runId, + this.runStart(runId), + ) const sender = input.principal ? { principal: input.principal } : {} const chain = this.createChain(this.loggedFields(first), { operation, @@ -2136,6 +2227,7 @@ export abstract class SessionAgents< ), ], ) + for (const state of open) this.publishWaitSettled(state.wait.waitId) this.armSchedules() operation.publish({ type: EventType.RUN_FINISHED, @@ -2249,7 +2341,11 @@ export abstract class SessionAgents< // Recovery must not run it again. this.admitted.set(input.inputId, inputKey(input.input, input.principal)) this.warnLiveOnly(runId) - const operation = this.agentOperation(first.input.agent, runId) + const operation = this.agentOperation( + first.input.agent, + runId, + this.runStart(runId), + ) const sender = input.principal ? { principal: input.principal } : {} const chain = this.createChain(this.loggedFields(first), { operation, @@ -2405,8 +2501,16 @@ export abstract class SessionAgents< if (queued.has(id)) continue this.pendingNotes.push({ id, role: 'assistant', content: note }) } - for (const [jobId, backendId] of [...writer.state.jobs]) { - if (this.jobs.has(jobId) || backendId !== undefined) continue + const lost = [...writer.state.jobs] + .filter( + ([jobId, backendId]) => + !this.jobs.has(jobId) && backendId === undefined, + ) + .map(([jobId]) => jobId) + // Only the owner notes them. While another live host owns the log, the + // job runs there. + if (lost.length > 0 && !(await writer.own())) return + for (const jobId of lost) { const id = createMessageId() const note = `Background job ${jobId} stopped when the host restarted.` this.pendingNotes.push({ id, role: 'assistant', content: note }) diff --git a/packages/ai-harness/src/session/base.ts b/packages/ai-harness/src/session/base.ts index bb606a3c7f..01d9f4b9dd 100644 --- a/packages/ai-harness/src/session/base.ts +++ b/packages/ai-harness/src/session/base.ts @@ -20,7 +20,9 @@ import { onQuestionOf, onUsageOf, } from '../feed' +import { runStartOf } from '../log' import { createMediaStore } from '../media' +import { OperationTracing } from '../observability' import { pluginStateDeltas } from '../log-records' import { SESSION_PLUGIN_CREDENTIALS, SESSION_PLUGIN_KEYS } from '../plugins' import { HARNESS_EVENTS } from '../types' @@ -72,6 +74,7 @@ import type { HarnessPersistence, SessionIndexWriter } from '../host' import type { HarnessPlugin, MountedPlugins, + PluginModelCall, PluginServices, PluginState, } from '../plugins' @@ -87,8 +90,9 @@ import type { RunContext, SessionEvent, ThreadSettings, - SendOptions, + SendForRunOptions, TurnInfo, + UserInput, } from '../types' import type { AgentRunOptions, @@ -113,6 +117,8 @@ import type { } from './types' import type { HarnessSession } from './harness-session' +const NO_NAMES: ReadonlySet = new Set() + /** * What a session keeps for its inputs, by input id: the duplicate check, the * receipt, the settlement, the turn operation, and the inputs that joined an @@ -263,6 +269,15 @@ export abstract class SessionBase { if (!this.logWriter) throw new Error('The session log is not open yet.') return this.logWriter } + + /** + * The cursor before the first event of the logged run `id`, for an + * operation that a restart rebuilds. `undefined` for a new run. + */ + protected runStart(id: string | undefined) { + if (id === undefined || !this.logWriter) return undefined + return runStartOf(this.logWriter.state, id) + } protected feed: ListenedFeed /** The host's view of each event: its status feed. */ protected readonly notifyHost: (event: StreamChunk) => void @@ -351,7 +366,10 @@ export abstract class SessionBase { /** Background jobs that run on this host. Recovery skips them. */ protected readonly jobs = new Set() /** Running tool calls that support `detach`: id to the move. */ - protected readonly detachable = new Map void>() + protected readonly detachable = new Map< + string, + (principal: Principal | undefined) => void + >() protected activeTurn: OperationImpl | undefined /** The open skipped recovery: another host owned the thread then. */ protected recoveryDeferred = false @@ -424,6 +442,11 @@ export abstract class SessionBase { Set<(value: unknown) => void> >() protected readonly configValues = new Map() + /** + * The config values of the log when the running turn started. The turn + * uses them until it ends. + */ + protected turnConfig: ReadonlyMap | undefined /** The stored settings of this thread: `session.configure`. */ protected threadSettings: ThreadSettings = {} protected readonly questions = new Map< @@ -504,6 +527,8 @@ export abstract class SessionBase { * answer to the child that waits for it. */ protected readonly harnessChildren = new Map() + /** Spans and `onOperation`. Without host `observability`, none. */ + protected readonly tracing: OperationTracing | undefined constructor(deps: SessionDependencies) { this.harness = deps.harness as THarness @@ -533,7 +558,10 @@ export abstract class SessionBase { this.inputRecords = new InputRecords(deps.retention.settledOperations) this.notifyHost = deps.onEvent this.feed = tapFeed( - new SessionFeed({ eventBytes: deps.retention.eventBytes }), + new SessionFeed({ + threadId: deps.threadId, + eventBytes: deps.retention.eventBytes, + }), deps.onEvent, this.eventListeners, this.redactor, @@ -545,6 +573,14 @@ export abstract class SessionBase { this.openChild = deps.openChild this.childLink = deps.child this.rootThreadId = deps.child?.rootThreadId ?? deps.threadId + this.tracing = deps.observability + ? new OperationTracing({ + ...deps.observability, + threadId: deps.threadId, + ...(deps.tenantId !== undefined ? { tenantId: deps.tenantId } : {}), + warn: (error) => this.warnSession('harness:observability', error), + }) + : undefined // The session gets its log view in `open()`. this.messages = notOpenMessages const onMissing = (error: AuthRequiredError) => @@ -753,11 +789,67 @@ export abstract class SessionBase { } } - protected configValue(key: string): unknown { - if (this.configValues.has(key)) return this.configValues.get(key) + /** A config value. Plugins read the copy of the running turn, if any. */ + protected configValue( + key: string, + values: ReadonlyMap = this.turnConfig ?? this.configValues, + ): unknown { + if (values.has(key)) return values.get(key) return this.sessionPlugins?.config.get(key)?.option.default } + /** The last `removedPlugins()` and `threadPlugins()` results. */ + private threadPluginsCache: + | { + settings: ThreadSettings + removed: ReadonlySet + mounted?: MountedPlugins + plugins?: MountedPlugins + } + | undefined + + /** The plugins that the `plugins.remove` setting of this thread names. */ + protected removedPlugins(): ReadonlySet { + const settings = this.threadSettings + const remove = settings.plugins?.remove ?? [] + if (remove.length === 0) return NO_NAMES + if (this.threadPluginsCache?.settings !== settings) { + this.threadPluginsCache = { settings, removed: new Set(remove) } + } + return this.threadPluginsCache.removed + } + + /** + * The session plugins as this thread has them: without everything that + * the removed plugins contribute (tools, prompts, every kind of + * middleware, agents, subagents, commands, config, extension items, and + * capabilities). + */ + protected threadPlugins(): MountedPlugins | undefined { + const mounted = this.sessionPlugins + const removed = this.removedPlugins() + if (!mounted || removed.size === 0) return mounted + const cache = this.threadPluginsCache + if (!cache) return mounted + if (cache.mounted !== mounted || !cache.plugins) { + cache.mounted = mounted + cache.plugins = mounted.withoutOwners([...removed]) + } + return cache.plugins + } + + /** + * The agent `name` of the session registry, as this thread has it: + * `undefined` when a plugin that this thread removed added it. + */ + protected threadAgent(name: string): AnyAgent | undefined { + const owner = this.agentRegistry.ownerOf(name) + if (owner === undefined || this.removedPlugins().has(owner)) { + return undefined + } + return this.agentRegistry.get(name) + } + /** Put each plugin's saved state in the snapshot before the first read. */ protected loadPluginState() { for (const [plugin, load] of this.stateLoaders) { @@ -844,6 +936,7 @@ export abstract class SessionBase { 'session', customEvent(HARNESS_EVENTS.commandsChanged, {}), ), + modelChat: (call) => this.modelChat(call), } } @@ -1018,8 +1111,8 @@ export abstract class SessionBase { /** * @internal No input being stored, no queued or running turn, no waiting - * steer, and no running agent. The host closes an idle session that no - * handle holds. + * steer, no running agent, and no background job of this host. The host + * closes an idle session that no handle holds. */ isIdle(): boolean { // A follow-up that waits for its run keeps the thread busy. @@ -1028,6 +1121,9 @@ export abstract class SessionBase { ) if (hasFollowUps) return false return ( + // A background job of this host runs in its memory. A close stops it, + // and its end note is lost. + this.jobs.size === 0 && // A live wait (an `ask()`) waits in this host's memory: its answer // must come here, so it keeps the log. A wait that only the log keeps // does not: the session is idle, and an answer reaches any host. @@ -1097,9 +1193,10 @@ export abstract class SessionBase { } /** - * Resolves when the thread is idle: no input is being stored, and no turn, - * steer, agent run, or command runs or waits. A thread that waits for a - * human (an approval, a sign-in) is idle. When the thread is idle now, it + * Resolves when the thread is idle: no input is being stored, no turn, + * steer, agent run, or command runs or waits, and no background job of + * this host runs. A thread that waits for a human (an approval, a sign-in) + * is idle. When the thread is idle now, it * resolves at once. When `signal` aborts first, it rejects with the reason * of `signal`. * @@ -1233,12 +1330,18 @@ export abstract class SessionBase { * Track a background job of this host. The log also has its start and * end, so recovery can note a job that a crash stopped. The record goes * to `threadId`, the thread of the workspace that runs the job. + * `principal` is who moved the job to the background. */ protected logJob( - job: { jobId: string; backendId?: string; threadId?: string }, + job: { + jobId: string + backendId?: string + threadId?: string + principal?: Principal + }, ended: boolean, ) { - const { jobId, backendId, threadId } = job + const { jobId, backendId, threadId, principal } = job if (ended) this.jobs.delete(jobId) else this.jobs.add(jobId) const isOtherThread = threadId !== undefined && threadId !== this.threadId @@ -1247,10 +1350,13 @@ export abstract class SessionBase { jobId, ended, ...(!ended && backendId !== undefined ? { backendId } : {}), + ...storedPrincipal(principal), ...(isOtherThread ? { thread: threadId } : {}), } // No one waits for a job record. A failed log write fails the session. this.writer.append([record]).catch(() => undefined) + // The last job ended: the host can close the session now. + if (ended) queueMicrotask(() => this.checkIdle()) } // Implemented by a later layer of the class. @@ -1267,16 +1373,53 @@ export abstract class SessionBase { operationId?: string, ): PluginSessionApi protected abstract resumeSignIns(): Promise + /** One model call of `ctx.model.chat`. See `SessionTurnParts.modelChat`. */ + protected abstract modelChat(call: PluginModelCall): Promise /** Resume the agent runs whose waits all ended. See `SessionAgents`. */ protected abstract resumeWaitingRuns(): Promise + /** The run that started an operation, for a child run. See `SessionAgents`. */ + protected abstract parentRunOf(operationId: string): string | undefined /** Set the timer of the next schedule. See `SessionImpl.armSchedules`. */ protected abstract armSchedules(options?: { isRetry?: boolean }): void /** @internal See `SessionRequests.resolveForRun`. */ abstract resolveForRun(options: ResolveForRunOptions): Promise + /** @internal See `SessionRequests.sendForRun`. */ + abstract sendForRun( + message: UserInput, + options?: SendForRunOptions, + ): Operation + /** See `SessionRequests.applyCancelInput`. */ + protected abstract applyCancelInput(options: { + inputId: string + principal?: Principal + controlId?: string + }): Promise + /** See `SessionRequests.applySetDelivery`. */ + protected abstract applySetDelivery(options: { + inputId: string + delivery: 'steer' | 'queue' + principal?: Principal + controlId?: string + }): Promise + /** See `SessionState.applyConfig`. */ + protected abstract applyConfig(options: { + key: string + value: unknown + principal?: Principal + controlId?: string + }): Promise + /** See `SessionState.applyAnswer`. */ + protected abstract applyAnswer(options: { + questionId: string + value: unknown + principal?: Principal + controlId?: string + ref?: string + }): Promise protected abstract recoverFromLog(writer: LogWriter): Promise protected abstract startTurn( input: HarnessInput & { op: 'prompt' | 'steer' | 'followUp' }, - options?: SendOptions & { wake?: boolean }, + options?: SendForRunOptions & { wake?: boolean }, ): Operation protected abstract runAgent( target: string | AnyAgent, diff --git a/packages/ai-harness/src/session/handle.ts b/packages/ai-harness/src/session/handle.ts index 17624c517a..645d957b90 100644 --- a/packages/ai-harness/src/session/handle.ts +++ b/packages/ai-harness/src/session/handle.ts @@ -25,6 +25,47 @@ export interface HandleGuard { ) => Promise /** `host.revoke` took the access of the principal of the handle. */ isRevoked: () => boolean + /** + * Call `listener` when `host.revoke` (or `unshare`) takes the access of + * the principal of the handle. Returns a function that stops the calls. + */ + onRevoke: (listener: () => void) => () => void +} + +/** + * @internal The items of `read`, until `signal` aborts or `onRevoke` calls + * its listener. A revoke ends the read with `AccessDeniedError`, so a read + * that started before the revoke stops too. The listener is added at the + * call, not at the first pull, so a revoke between them ends the read too. + * + * ponytail: an iterable that is never read keeps its listener until the + * revoke. + */ +export function untilRevoked({ + onRevoke, + threadId, + signal, + read, +}: { + onRevoke: HandleGuard['onRevoke'] + threadId: string + signal: AbortSignal | undefined + read: (signal: AbortSignal) => AsyncIterable +}): AsyncGenerator { + const revoked = new AbortController() + const stop = onRevoke(() => revoked.abort()) + return (async function* () { + try { + yield* read( + signal ? AbortSignal.any([signal, revoked.signal]) : revoked.signal, + ) + } finally { + stop() + } + if (revoked.signal.aborted) { + throw new AccessDeniedError({ action: 'read', threadId }) + } + })() } /** How a handle reaches the shared session of its thread. The host gives one. */ @@ -403,6 +444,44 @@ export function sessionHandle({ live(action) return run((target) => checkedOn(target, call)) } + /** The operations that `guarded` wrapped, so one operation gets one wrapper. */ + const wrapped = new WeakMap() + /** + * `operation`, with its `cancel` and `send` through this handle: they act + * as the principal of the handle, after `authorize` allows them. + */ + const guarded = >(operation: T): T => { + const known = wrapped.get(operation) + if (known) return known as T + const proxy = new Proxy(operation, { + get(target, key) { + if (key === 'cancel') { + return async () => cancel({ operationId: target.id }) + } + const value: unknown = Reflect.get(target, key) + if (typeof value !== 'function') return value + if (key === 'send') { + return async ( + message: UserInput, + options?: Omit, + ) => { + const sent = send(message, { ...options, to: target.id }) + const receipt = await sent.receipt + // A refused access check throws, as the calls of the handle do. + if (receipt.status === 'rejected') { + await Promise.resolve(sent).catch((error: unknown) => { + if (error instanceof AccessDeniedError) throw error + }) + } + return receipt + } + } + return value.bind(target) + }, + }) + wrapped.set(operation, proxy) + return proxy + } /** A call that stores an input and returns its operation. */ const storingOperation = >( kind: OperationKind, @@ -410,7 +489,11 @@ export function sessionHandle({ agent?: string, ) => { live('send') - return runOperation(kind, (target) => checkedOn(target, call), agent) + return runOperation( + kind, + (target) => guarded(checkedOn(target, call)), + agent, + ) } /** * A call that stores no input: it runs after `authorize` allows `action`. @@ -434,7 +517,10 @@ export function sessionHandle({ run: (input?: unknown, options?: AgentRunOptions) => storingOperation( 'agent', - (target) => target.agentHandle(name, principal).run(input, options), + (target) => + target + .agentHandle(name, actor(options?.principal)) + .run(input, options), name, ), start: (input?: unknown, options?: AgentStartOptions) => @@ -474,18 +560,48 @@ export function sessionHandle({ live('read') const target = open() const isHeld = target !== undefined && !token.isReleased - return isHeld ? target.events(options) : heldEvents(options) + const read = (next?: EventsOptions) => + isHeld ? target.events(next) : heldEvents(next) + if (!guard) return read(options) + // A revoke ends a read that runs now, not only the next call. + return untilRevoked({ + onRevoke: guard.onRevoke, + threadId, + signal: options?.signal, + read: (signal) => read({ ...options, signal }), + }) + } + const send = (message: UserInput, options?: SendOptions | AgentSendOptions) => + storingOperation(options?.to === undefined ? 'chat' : 'agent', (target) => + options?.to === undefined + ? target.send(message, as(options ?? {})) + : target.send(message, as(options)), + ) + const sendForRun = ( + message: UserInput, + options?: Parameters['sendForRun']>[1], + ) => + storingOperation('chat', (target) => + target.sendForRun(message, as(options ?? {})), + ) + const cancel = (options?: Parameters['cancel']>[0]) => + storing('cancel', (target) => target.cancel(as(options ?? {}))) + /** The operation `id` of the session, through this handle. */ + const operationOf = >( + find: (target: Session) => T | undefined, + ) => { + live('read') + const found = find(open() ?? last) + return found && guarded(found) } const methods = new Map([ ...Object.entries({ - send: (message: UserInput, options?: SendOptions | AgentSendOptions) => - storingOperation( - options?.to === undefined ? 'chat' : 'agent', - (target) => - options?.to === undefined - ? target.send(message, as(options ?? {})) - : target.send(message, as(options)), - ), + send, + sendForRun, + // A returned operation cancels, and an agent run takes messages, + // through this handle too. + operation: (id: string) => operationOf((target) => target.operation(id)), + agentRun: (id: string) => operationOf((target) => target.agentRun(id)), tool: (options: Parameters['tool']>[0]) => storingOperation('chat', (target) => target.tool(as(options))), continue: (options?: Parameters['continue']>[0]) => @@ -505,15 +621,9 @@ export function sessionHandle({ ) => storing('answer', (target) => target.resolveForRun(as(options))), answer: (options: Parameters['answer']>[0]) => storing('answer', (target) => target.answer(as(options))), - cancel: (options?: Parameters['cancel']>[0]) => - storing('cancel', (target) => target.cancel(as(options ?? {}))), - cancelInput: ( - inputId: string, - options?: Parameters['cancelInput']>[1], - ) => - storing('cancel', (target) => - target.cancelInput(inputId, as(options ?? {})), - ), + cancel, + cancelInput: (options: Parameters['cancelInput']>[0]) => + storing('cancel', (target) => target.cancelInput(as(options))), configure: (options: Parameters['configure']>[0]) => storing('configure', (target) => target.configure(as(options))), setConfig: (options: Parameters['setConfig']>[0]) => @@ -522,13 +632,13 @@ export function sessionHandle({ agent: (name: string) => (open() ?? last).agent(name) === undefined ? undefined : agentOf(name), // These store no input, so the handle asks `authorize` itself. - revert: (messageId: string) => - acting('send', (target) => target.revert(messageId)), + revert: (options: Parameters['revert']>[0]) => + acting('send', (target) => target.revert(as(options))), unrevert: () => acting('send', (target) => target.unrevert()), - background: (toolCallId?: string) => - acting('send', (target) => target.background(toolCallId)), - append: (records: Parameters['append']>[0]) => - acting('send', (target) => target.append(records)), + background: (options?: Parameters['background']>[0]) => + acting('send', (target) => target.background(as(options ?? {}))), + append: (options: Parameters['append']>[0]) => + acting('send', (target) => target.append(as(options))), putMedia: (upload: Parameters['putMedia']>[0]) => acting('send', (target) => target.putMedia(upload)), updateEntry: (update: Parameters['updateEntry']>[0]) => diff --git a/packages/ai-harness/src/session/harness-session.ts b/packages/ai-harness/src/session/harness-session.ts index 46677e1066..eacffcb591 100644 --- a/packages/ai-harness/src/session/harness-session.ts +++ b/packages/ai-harness/src/session/harness-session.ts @@ -22,6 +22,7 @@ import { CONTROL_OPS, hasFinalAnswer, loggedExtras, + notifyListeners, offeredInputRecord, SHUTDOWN_REASON, storedPrincipal, @@ -58,12 +59,16 @@ export interface HarnessSession< | 'transferSummary' | 'onIdle' | 'resolveForRun' + | 'sendForRun' | 'checked' + | 'knownMessage' | 'cancelWorkOf' | 'agentHandle' | 'isIdle' | 'isClosing' | 'openWait' + | 'openToEnd' + | 'endThread' // A handle closes itself, not the session. | 'close' > { @@ -145,6 +150,7 @@ export class SessionImpl< /** @internal Mount session plugins and recover the work in the log. */ async open(): Promise { + await this.tracing?.ready() // First: plugins publish events while they mount. await this.openLog() try { @@ -169,6 +175,41 @@ export class SessionImpl< } } + /** + * @internal For `host.sessions.delete`: open the log and the session + * plugins, and nothing else. No recovery, no index entry, no offers, and + * no schedules, so no work of the thread runs again before the purge. Call + * `endThread`, then close it. + */ + async openToEnd(): Promise { + await this.openLog() + await this.mountSessionPlugins() + this.loadSettings() + } + + /** + * @internal The thread ends for good: each `ThreadEnds` handler of the + * plugins runs, the newest first, so a plugin that needs the workspace of + * an older plugin runs before that workspace closes. Every handler runs. + * Then the first error, if any, is thrown. + */ + async endThread(): Promise { + const run = this.runContext({ + principal: undefined, + runId: `end:${this.threadId}`, + signal: this.lifetime.signal, + }) + const errors: Array = [] + for (const handler of [...this.threadEnds].reverse()) { + await handler({ + tenantId: this.tenantId, + threadId: this.threadId, + run, + }).catch((error: unknown) => errors.push(error)) + } + if (errors.length > 0) throw errors[0] + } + /** Also waits for the last append of the log, so a reopen reads it. */ override async close(options?: { recoverable?: boolean }): Promise { // The log keeps the schedules. The release of the log gives the store @@ -206,6 +247,8 @@ export class SessionImpl< ), services: this.services, inherited: this.storeCapabilities(), + // A plugin that this thread removed gives `collect` nothing. + hiddenOwners: () => this.removedPlugins(), }, ) this.loadConfig() @@ -522,6 +565,22 @@ export class SessionImpl< (!this.isIdle() || this.jobs.size > 0 || this.recovering > 0) && this.sleeping === 0, onOffers: () => void this.drainOffers(), + // An event that another host wrote: the host status and the `onEvent` + // listeners see it, as they see the events of this host at publish. + onRemoteEvent: ({ operationId, event }) => + notifyListeners(this.notifyHost, this.eventListeners, { + operationId, + event, + }), + // Records of another host: the settings, the config, and the + // interrupted turn come from the new fold. + onRemoteFold: () => this.syncFromLog(), + // This host became the owner. The event goes in the next append. + onOwned: (change) => + this.feed.publish( + 'session', + customEvent(HARNESS_EVENTS.ownershipChanged, { ...change }), + ), redact: this.redactor.redact, }) this.logWriter = writer @@ -783,10 +842,18 @@ export class SessionImpl< } // The host keeps the log until recovery queued its work. this.recovering += 1 + // A span only when inputs of a host that stopped wait for this host. + const span = hasWork + ? this.tracing?.span('harness.recovery', {}) + : undefined try { try { await this.recoverInputs(writer, isTaken) } catch (error) { + this.tracing?.safely(() => { + span?.recordException(error instanceof Error ? error : String(error)) + span?.setStatus({ code: 2 }) + }) // Not recovered, so the next `recover()` takes them again. for (const inputId of taken) { if (!this.turnOperations.has(inputId)) this.admitted.delete(inputId) @@ -803,6 +870,7 @@ export class SessionImpl< await this.resumeWaitingRuns() } finally { this.recovering -= 1 + this.tracing?.safely(() => span?.end()) // A recovery that left no work gives the log back. void writer.rest() } diff --git a/packages/ai-harness/src/session/helpers.ts b/packages/ai-harness/src/session/helpers.ts index 59c632434c..d60afd1380 100644 --- a/packages/ai-harness/src/session/helpers.ts +++ b/packages/ai-harness/src/session/helpers.ts @@ -16,6 +16,7 @@ import type { Scope, StreamChunk, SubagentsBag, + SystemPrompt, } from '@tanstack/ai' import type { LogRecord, MessageStore } from '@tanstack/ai-persistence' import type { Redactor } from '../auth' @@ -29,7 +30,7 @@ import type { } from '../log-records' import type { MountedPlugins } from '../plugins' import type { - HarnessInput, + OfferedHarnessInput, MediaKind, OperationStatus, Principal, @@ -121,7 +122,7 @@ function namedOverrides(extras: TurnExtras) { export interface OfferedInputRecord extends LogRecord { type: 'harness.input' inputId: string - input: HarnessInput + input: OfferedHarnessInput principal?: { id: string; tenantId?: string } /** * The id that the owner runs the input with: the AG-UI run id of a turn, @@ -131,15 +132,18 @@ export interface OfferedInputRecord extends LogRecord { runId?: string /** The overrides of a turn, as its `harness.input.turn` record keeps them. */ overrides?: NamedTurnOverrides + /** The W3C `traceparent` of the caller. The owner keeps it on its record. */ + traceparent?: string } /** The offer record of `input`. See {@link OfferedInputRecord}. */ export function offeredInputRecord(offer: { inputId: string - input: HarnessInput + input: OfferedHarnessInput principal: Principal | undefined runId?: string turn?: TurnExtras + traceparent?: string }) { const named = offer.turn ? namedOverrides(offer.turn) : {} const record: OfferedInputRecord = { @@ -149,6 +153,9 @@ export function offeredInputRecord(offer: { ...storedPrincipal(offer.principal), ...(offer.runId !== undefined ? { runId: offer.runId } : {}), ...(Object.keys(named).length > 0 ? { overrides: named } : {}), + ...(offer.traceparent !== undefined + ? { traceparent: offer.traceparent } + : {}), } return record } @@ -499,34 +506,22 @@ export function changeSettings( } /** - * `mounted` without what the `removed` plugins give a turn: their tools, - * prompts, middleware, adapter picks, discoverers, and preparers. + * The owner of each tool name of `tools`: the plugin from `owners`, else + * `the harness` (also for the turn's own tools). The first owner of a name + * stays. */ -export function withoutPlugins( - mounted: MountedPlugins, - removed: ReadonlySet, -): MountedPlugins { - const { owners } = mounted - const toolOwners = new Map( - owners.tools.map((entry) => [entry.name, entry.owner]), - ) - const isKept = (owner: string | undefined) => - owner === undefined || !removed.has(owner) - return { - ...mounted, - tools: mounted.tools.filter((tool) => isKept(toolOwners.get(tool.name))), - prompts: mounted.prompts.filter((_prompt, index) => - isKept(owners.prompts[index]?.owner), - ), - middleware: mounted.middleware.filter((_item, index) => - isKept(owners.middleware[index]), - ), - adapters: mounted.adapters.filter((_pick, index) => - isKept(owners.adapters[index]), - ), - discoverers: mounted.discoverers.filter((item) => isKept(item.owner)), - preparers: mounted.preparers.filter((item) => isKept(item.owner)), +export function toolOwners( + tools: ReadonlyArray, + owners: ReadonlyArray<{ name: string; owner: string }>, +) { + const byName = new Map() + for (const entry of owners) { + if (entry.owner === 'the harness' || byName.has(entry.name)) continue + byName.set(entry.name, entry.owner) } + return new Map( + tools.map((tool) => [tool.name, byName.get(tool.name) ?? 'the harness']), + ) } /** @@ -622,6 +617,39 @@ export function ephemeralCall(pending: Array): AnyChatMiddleware { } } +/** + * Gives each model call the `'context'` plugin prompts as one system prompt, + * after all other system prompts. The text has system authority, not user + * authority. Each call takes out the prompt of the call before, so a change + * replaces it and the stable system prompts before it stay the same. + * + * ponytail: a `ModelMessage` has no `system` role, so the prompt cannot go + * after the history. A change of it changes the cached prefix from the + * system prompts on. A system message in the history needs a core change. + * + * `texts` must be scrubbed of the session secrets already: a redacted copy + * is a new object, and the next call could not find it to take it out. + */ +export function contextPromptCall( + texts: () => ReadonlyArray, +): AnyChatMiddleware { + let last: SystemPrompt | undefined + return { + name: 'harness:context-prompts', + onConfig: (ctx, config) => { + if (ctx.phase !== 'beforeModel') return undefined + const text = texts() + .filter((item) => item !== '') + .join('\n\n') + const kept = config.systemPrompts.filter((prompt) => prompt !== last) + if (text === '' && kept.length === config.systemPrompts.length) + return undefined + last = text === '' ? undefined : { content: text } + return { systemPrompts: last ? [...kept, last] : kept } + }, + } +} + /** * Gives each model call of a turn the `ephemeral` messages of its input, * after `start`: the last message of the context when the turn started. The @@ -837,6 +865,34 @@ export const notOpenMessages: MessageStore = { saveThread: notOpen, } +/** + * Give one event to `notifyHost` and each of `listeners`. A listener that + * throws, or whose promise rejects, is logged, and the others still run. + */ +export function notifyListeners( + notifyHost: (event: StreamChunk) => void, + listeners: Set, + entry: Parameters[0], +) { + const report = (error: unknown) => + console.error('[@tanstack/ai-harness] An onEvent listener threw.', error) + try { + notifyHost(entry.event) + } catch (error) { + report(error) + } + for (const listener of [...listeners]) { + // A listener that an earlier listener stopped gets nothing. + if (!listeners.has(listener)) continue + try { + const result: unknown = listener(entry) + if (result instanceof Promise) result.catch(report) + } catch (error) { + report(error) + } + } +} + /** * `feed`, then `notifyHost` and each of `listeners` after each publish. This * is how the host and `onEvent` see events. The caller keeps `listeners`, so @@ -852,27 +908,12 @@ export function tapFeed( return { publish: (operationId, raw) => { const event = redactor.redact(raw) + // ponytail: core sends a MESSAGES_SNAPSHOT only when a run pauses at + // an interrupt, so the log keeps each one. A rejoin from the log (the + // run route, `Last-Event-ID`) needs it: ai-client rebuilds the paused + // run from it. If core sends it at other times, filter here. feed.publish(operationId, event) - const report = (error: unknown) => - console.error( - '[@tanstack/ai-harness] An onEvent listener threw.', - error, - ) - try { - notifyHost(event) - } catch (error) { - report(error) - } - for (const listener of [...listeners]) { - // A listener that an earlier listener stopped gets nothing. - if (!listeners.has(listener)) continue - try { - const result: unknown = listener({ operationId, event }) - if (result instanceof Promise) result.catch(report) - } catch (error) { - report(error) - } - } + notifyListeners(notifyHost, listeners, { operationId, event }) }, listen: (listener) => { // A new entry per call: a listener added twice is called twice. diff --git a/packages/ai-harness/src/session/inputs.ts b/packages/ai-harness/src/session/inputs.ts index 680ba7b561..4b91957bd0 100644 --- a/packages/ai-harness/src/session/inputs.ts +++ b/packages/ai-harness/src/session/inputs.ts @@ -28,10 +28,12 @@ import type { AnyHarness } from '../define' import type { ChatTurnResult, HarnessInput, + OfferedHarnessInput, InputSettlement, Principal, Receipt, SessionEvent, + UserInput, WaitingInput, } from '../types' import type { OfferedInputRecord, TurnExtras } from './helpers' @@ -53,11 +55,13 @@ interface AcceptOptions { runId?: string /** The overrides, preamble, and ephemeral messages of a turn. */ turn?: TurnExtras + /** A valid W3C `traceparent` of the caller. The input record keeps it. */ + traceparent?: string /** * The input that goes to the owner when another host owns the log, if it * is not the input that this host logs. It throws to refuse the input. */ - offer?: () => Promise + offer?: () => Promise } /** The offer that the next `accept` applies. See `applyOffer`. */ @@ -118,6 +122,18 @@ export abstract class SessionInputs< } } + /** + * @internal The message of the input `inputId` in the log, or `undefined`. + * A client that sends an input again sends this message, so its inline + * media is not stored again and the duplicate check matches. + */ + // ponytail: reads the log only. A retry that comes while the first copy + // is still being stored stores its media again. + knownMessage(inputId: string): UserInput | undefined { + const input = this.writer.state.inputs.get(inputId)?.input + return input && 'message' in input ? input.message : undefined + } + /** * The run that `input` answers or cancels, and who started it, when the * session knows. Other inputs target no run. @@ -238,14 +254,39 @@ export abstract class SessionInputs< * With `project` on the host, a record can change the model context. The * running turn sees the change at its next model call. * + * With `inputId`, the append is an input, as a `send` is: a second append + * with the same id is a duplicate and writes nothing, so a client can + * send it again after a lost answer. Without it, the append gets a new + * input id. `principal` is who appends. Default: the principal that + * opened the session. The receipt has the input id, as the log keeps it. + * * @example * ```ts - * await session.append([{ type: 'app.signal', text: 'The build failed.' }]) + * await session.append({ + * records: [{ type: 'app.signal', text: 'The build failed.' }], + * }) * ``` */ - async append(records: ReadonlyArray): Promise { + async append({ + records, + inputId = createInputId(), + principal, + }: { + records: ReadonlyArray + inputId?: string + principal?: Principal + }): Promise { checkHostRecords(records) - await this.writer.append(records) + const admission = await this.accept( + inputId, + { op: 'append', records: [...records] }, + principal ?? this.principal, + ) + if (admission !== 'new') return this.duplicateReceipt(inputId, admission) + // In the append of the applied record: a crash keeps both or neither. + await this.applied(inputId, 'session', records) + // Kept, so a duplicate gets the same receipt. + return this.keep({ inputId, status: 'accepted' }) } /** @@ -259,13 +300,24 @@ export abstract class SessionInputs< * wakes an idle session. Without `toolCallId`, it moves every running * call that supports it, such as `bash` and the agent tools. * With no such call, the receipt is rejected with `not_running`. + * `principal` is who moves it. Default: the principal that opened the + * session. The `harness.job` record of each moved call keeps it. A handle + * of `host.open` checks that it is the principal of the handle. The move + * stores no input. * * @example * ```ts * await session.background() + * await session.background({ toolCallId: 'call_1' }) * ``` */ - async background(toolCallId?: string): Promise { + async background({ + toolCallId, + principal, + }: { + toolCallId?: string + principal?: Principal + } = {}): Promise { const inputId = createInputId() const ids = toolCallId === undefined ? [...this.detachable.keys()] : [toolCallId] @@ -274,7 +326,7 @@ export abstract class SessionInputs< return { inputId, status: 'rejected', reason: 'not_running' } } for (const id of ids) this.detachable.delete(id) - for (const move of moves) move() + for (const move of moves) move(principal ?? this.principal) return { inputId, status: 'accepted' } } @@ -286,7 +338,9 @@ export abstract class SessionInputs< * @example * ```ts * for (const input of session.inputs()) { - * if (input.delivery === 'queue') await session.cancelInput(input.inputId) + * if (input.delivery === 'queue') { + * await session.cancelInput({ inputId: input.inputId }) + * } * } * ``` */ @@ -348,16 +402,33 @@ export abstract class SessionInputs< ) } signal?.throwIfAborted() + // A refused input ends too: the owner on another host writes only its + // rejected record, and `knownSettlement` throws for it. const settledEvent = (entry: SessionEvent) => entry.event.type === EventType.CUSTOM && - entry.event.name === HARNESS_EVENTS.inputSettled && + (entry.event.name === HARNESS_EVENTS.inputSettled || + entry.event.name === HARNESS_EVENTS.inputRejected) && isRecord(entry.event.value) && entry.event.value.inputId === inputId - for await (const _entry of this.feed.read({ + for await (const entry of this.feed.read({ from: this.feed.head(), filter: settledEvent, ...(signal ? { signal } : {}), })) { + // The owner refused the offer of this host: the log has another input + // with this id, and only this host ends with the refusal. + const value = entry.event.type === EventType.CUSTOM && entry.event.value + if ( + isRecord(value) && + value.offered === true && + this.offeredInputs.has(inputId) + ) { + throw new InputRejectedError({ + inputId, + status: 'rejected', + reason: typeof value.reason === 'string' ? value.reason : 'conflict', + }) + } const settlement = this.knownSettlement(inputId) if (settlement) return settlement } @@ -564,6 +635,10 @@ export abstract class SessionInputs< await check({ input, run: this.runOf(input), ...(wait ? { wait } : {}) }) } const sender = storedPrincipal(principal) + const traced = + options.traceparent === undefined + ? {} + : { traceparent: options.traceparent } const known = this.knownAdmission( inputId, inputKey(input, sender.principal), @@ -571,7 +646,9 @@ export abstract class SessionInputs< const consumes = offer ? [inputId] : [] if (known) { // The input ran before: its offer goes, and it does not run again. - if (offer) await this.writer.append([], { consumes }) + // An offer with another payload is refused, so its host stops waiting. + if (offer && known === 'conflict') await this.refuseOffer(inputId) + else if (offer) await this.writer.append([], { consumes }) return known } // A closing session runs no new work. Nothing is stored. An offer stays @@ -582,7 +659,14 @@ export abstract class SessionInputs< const records = options.turn ? turnRecords(inputId, options.turn) : [] await this.writer.append( [ - { type: 'harness.input', inputId, input, at: Date.now(), ...sender }, + { + type: 'harness.input', + inputId, + input, + at: Date.now(), + ...sender, + ...traced, + }, ...records, ], { consumes }, @@ -739,19 +823,30 @@ export abstract class SessionInputs< await done return } + // A prompt with a known id ends before `accept`. One with another + // payload is refused, so the host that offered it stops waiting. + const payload = inputKey(record.input, record.principal) + if (this.knownAdmission(inputId, payload) === 'conflict') { + await this.refuseOffer(inputId) + return + } } // `host.transfer` on a host that does not own the log offers records // here: the source end of the transfer, and the usage of its summary. // They go in the log as they are. A transfer that the log has already - // only drops its offer. + // only drops its offer. An unshare or a revoke offers the end of the + // schedules of the user, so they go from the list at once. const isLaneRecord = - record.type === 'harness.usage' || record.type === 'harness.transferred' + record.type === 'harness.usage' || + record.type === 'harness.transferred' || + record.type === 'harness.schedule.done' const isKnown = record.type === 'harness.transferred' && typeof record.transferId === 'string' && this.writer.state.thread.transfers.has(record.transferId) const records = isLaneRecord && !isKnown ? [record] : [] await this.writer.append(records, { consumes: [inputId] }) + for (const each of records) this.publishTransferred(each) } /** @@ -765,6 +860,10 @@ export abstract class SessionInputs< const user = principal ? { principal } : {} const run = record.runId !== undefined ? { runId: record.runId } : {} const turn = record.overrides ? loggedExtras(record.overrides) : {} + const traced = + record.traceparent !== undefined + ? { traceparent: record.traceparent } + : {} const sent = { inputId, ...user, @@ -781,12 +880,13 @@ export abstract class SessionInputs< switch (input.op) { case 'prompt': ignore( - session.send(input.message, { + this.sendForRun(input.message, { ...sent, ...turn, ...(input.busy ? { busy: input.busy } : {}), ...(input.as ? { as: input.as } : {}), ...run, + ...traced, }), ) return @@ -794,13 +894,21 @@ export abstract class SessionInputs< ignore(session.continue({ ...sent, ...turn })) return case 'steer': - ignore(session.send(input.message, { ...sent, ...turn, mode: 'steer' })) + ignore( + session.send(input.message, { + ...sent, + ...turn, + ...traced, + mode: 'steer', + }), + ) return case 'followUp': ignore( this.startTurn(input, { ...sent, ...turn, + ...traced, mode: 'followUp', ...(input.wake ? { wake: true } : {}), }), @@ -841,11 +949,17 @@ export abstract class SessionInputs< ) return case 'cancelInput': - ignore(session.cancelInput(input.inputId, { inputId, ...user })) + ignore( + this.applyCancelInput({ + inputId: input.inputId, + controlId: inputId, + ...user, + }), + ) return case 'setDelivery': ignore( - session.setDelivery({ + this.applySetDelivery({ inputId: input.inputId, delivery: input.delivery, controlId: inputId, @@ -865,21 +979,21 @@ export abstract class SessionInputs< return case 'answer': ignore( - session.answer({ + this.applyAnswer({ questionId: input.questionId, value: input.value, ...(input.ref !== undefined ? { ref: input.ref } : {}), - inputId, + controlId: inputId, ...user, }), ) return case 'config': ignore( - session.setConfig({ + this.applyConfig({ key: input.key, value: input.value, - inputId, + controlId: inputId, ...user, }), ) @@ -896,10 +1010,13 @@ export abstract class SessionInputs< }), ) return + case 'append': + // Only an append with an input id is an input that a host offers. + ignore(session.append({ records: input.records, inputId, ...user })) + return case 'revert': case 'unrevert': case 'background': - case 'append': // These run only where they are called. No host offers them. return } @@ -975,6 +1092,28 @@ export abstract class SessionInputs< return receipt } + /** + * Refuse an offer whose id the log has with another payload. The record + * consumes the offer in the same append. The input in the log does not + * change: the event has `offered`, so only the host that offered the + * input ends its wait with the refusal. + */ + private async refuseOffer(inputId: string) { + const reason = 'conflict' + await this.writer.append( + [{ type: 'harness.offer.rejected', inputId, reason }], + { consumes: [inputId] }, + ) + this.feed.publish( + 'session', + customEvent(HARNESS_EVENTS.inputRejected, { + inputId, + reason, + offered: true, + }), + ) + } + /** * Record how inputs ended, all in one append: a turn and the inputs that * joined it settle together. Clients get a `harness.input.settled` event. @@ -1057,10 +1196,54 @@ export abstract class SessionInputs< receipt: Receipt, ) { operation.resolveReceipt(this.keep(receipt)) + // Only the owner answers a turn, so the input waits in its queue. + if (receipt.status === 'queued') { + this.publishQueued(receipt.inputId, 'queue') + } + } + + /** The input `inputId` waits behind the running turn, or the turns before it. */ + protected publishQueued(inputId: string, delivery: WaitingInput['delivery']) { + this.feed.publish( + 'session', + customEvent(HARNESS_EVENTS.inputQueued, { inputId, delivery }), + ) + } + + /** + * Publish `harness.transferred` for the transfer end `record`, after this + * host appended it. Other records get no event. + */ + protected publishTransferred(record: LogRecord) { + if (record.type !== 'harness.transferred') return + const { transferId, from, to } = record + if (typeof transferId !== 'string') return + const threadOf = (link: unknown) => + isRecord(link) && typeof link.threadId === 'string' + ? { threadId: link.threadId } + : undefined + const source = threadOf(from) + const target = threadOf(to) + this.feed.publish( + 'session', + customEvent(HARNESS_EVENTS.transferred, { + transferId, + ...(source ? { from: source } : {}), + ...(target ? { to: target } : {}), + }), + ) } protected refuse(operation: OperationImpl, receipt: Receipt) { - this.answerTurn(operation, receipt) + // A conflict refuses a second payload. The id stays with the input that + // has it, so `settled()` of that input does not read this refusal. + const id = receipt.inputId + const isOtherInput = + receipt.status === 'rejected' && + receipt.reason === 'conflict' && + (this.admitted.has(id) || this.writer.state.inputs.has(id)) + if (isOtherInput) operation.resolveReceipt(receipt) + else this.answerTurn(operation, receipt) operation.fail('failed', new InputRejectedError(receipt)) } @@ -1151,6 +1334,19 @@ export abstract class SessionInputs< ...(operation.agent ? { agent: operation.agent } : {}), }), ) + if (this.tracing) { + // A turn has its input: its record keeps the `traceparent` of the + // caller, also when another host took the input. + const inputId = this.inputRecords.operationInputs.get(operation.id) + const logged = + inputId === undefined + ? undefined + : this.logWriter?.state.inputs.get(inputId) + this.tracing.started(operation, { + parentRunId: this.parentRunOf(operation.id), + traceparent: logged?.traceparent, + }) + } } protected publishFinished( diff --git a/packages/ai-harness/src/session/observed-session.ts b/packages/ai-harness/src/session/observed-session.ts index 1f7df8d989..03e7d16157 100644 --- a/packages/ai-harness/src/session/observed-session.ts +++ b/packages/ai-harness/src/session/observed-session.ts @@ -1,7 +1,8 @@ import { loadChildTranscript } from '@tanstack/ai-persistence' -import { visibleMessages } from '../log' +import { runStartOf, visibleMessages } from '../log' import { snapshotWaits } from '../waits' import { createMediaStore } from '../media' +import { untilRevoked } from './handle' import type { ModelMessage } from '@tanstack/ai' import type { BlobRange, @@ -10,7 +11,7 @@ import type { SessionIndexEntry, } from '@tanstack/ai-persistence' import type { LogReader, ReduceOptions } from '../log' -import type { MediaPersistence } from '../media' +import type { MediaPersistence, MediaStore } from '../media' import type { SessionUsage } from '../usage' import type { AnyHarness, @@ -19,6 +20,7 @@ import type { } from '../define' import type { SessionIndexWriter } from '../host' import type { Cursor, MediaRecord, SessionEvent } from '../types' +import type { HandleGuard } from './handle' import type { SessionImpl } from './harness-session' import type { SessionSnapshot } from './types' @@ -75,9 +77,11 @@ export interface ObservedSession { * The saved messages of a child agent run of the thread, as * `session.agentTranscript()` gives them. Empty for an unknown id. */ - agentTranscript: (subagentRunId: string) => Promise> + agentTranscript: (options: { + subagentRunId: string + }) => Promise> /** The record of a media file of the thread, or `null` when it is not found. */ - getMedia: (id: string) => Promise + getMedia: (options: { id: string }) => Promise /** * The bytes of the media file `id`, or of one `range` of them. Throws a * `MediaError` with 404 when it is not found. @@ -108,8 +112,13 @@ export interface ObservedSources { tenantId: string | undefined /** The open session of the thread on this host, or `undefined`. */ live: () => SessionImpl | undefined - /** Throws when the host is closed. */ + /** Throws when the host is closed, and when a revoke took the access. */ check: () => void + /** + * With access: call `listener` when `host.revoke` (or `unshare`) takes the + * access of the user of the view. A read that runs then stops. + */ + onRevoke?: HandleGuard['onRevoke'] index: SessionIndexWriter media: MediaPersistence /** @@ -137,11 +146,26 @@ export function liveSessionOf(view: ObservedSession) { return liveSessions.get(view)?.() } +/** The media store of each view, for the HTTP handler. Not public API. */ +const mediaStores = new WeakMap() + +/** @internal The media store that `view` reads, so a route can stream bytes. */ +export function mediaStoreOf(view: ObservedSession) { + return mediaStores.get(view) +} + /** How each view reads the runs in its log, for the HTTP handler. Not public API. */ const storedRuns = new WeakMap< object, () => Promise> >() +/** How each view reads one run and its start, for the HTTP handler. */ +const storedRun = new WeakMap< + object, + ( + runId: string, + ) => Promise<{ record: RunRecord | undefined; start: Cursor | undefined }> +>() /** * @internal The runs in the log of the thread of `view`, by run id: the @@ -153,6 +177,16 @@ export function storedRunsOf(view: ObservedSession) { return read() } +/** + * @internal The run `runId` in the log of the thread of `view`, and the + * cursor before its first event (`undefined` when the log does not know it). + */ +export function storedRunOf(view: ObservedSession, runId: string) { + const read = storedRun.get(view) + if (!read) throw new Error('Not a view of host.observe.') + return read(runId) +} + /** The view of `host.observe`. See {@link ObservedSession}. */ export function observedSession( sources: ObservedSources, @@ -201,9 +235,19 @@ export function observedSession( } } + const { onRevoke } = sources const view: ObservedSession = { threadId, - events: (options) => events(options ?? {}), + // A revoke ends a read that runs now, not only the next call. + events: (options) => + onRevoke + ? untilRevoked({ + onRevoke, + threadId, + signal: options?.signal, + read: (signal) => events({ ...options, signal }), + }) + : events(options ?? {}), async transcript() { const session = open() if (session) return session.transcript() @@ -250,7 +294,7 @@ export function observedSession( } return snapshot }, - async agentTranscript(subagentRunId) { + async agentTranscript({ subagentRunId }) { check() const { reader } = await readLog() try { @@ -269,7 +313,7 @@ export function observedSession( reader.close() } }, - getMedia(id) { + getMedia({ id }) { check() return mediaStore.get(id) }, @@ -302,10 +346,18 @@ export function observedSession( }, } liveSessions.set(view, live) + mediaStores.set(view, mediaStore) // The fold changes by copy, so this map stays as it is after the read. storedRuns.set(view, () => { check() return fromLog(({ state }) => state.thread.runs) }) + storedRun.set(view, (runId) => { + check() + return fromLog(({ state }) => ({ + record: state.thread.runs.get(runId), + start: runStartOf(state, runId), + })) + }) return view } diff --git a/packages/ai-harness/src/session/requests.ts b/packages/ai-harness/src/session/requests.ts index deaefbd2e8..e5cfee2b16 100644 --- a/packages/ai-harness/src/session/requests.ts +++ b/packages/ai-harness/src/session/requests.ts @@ -1,6 +1,7 @@ import { RUN_CANCEL_REASON } from '@tanstack/ai' import { changedEntry } from '../session-entry' import { OperationImpl } from '../operation' +import { parseTraceparent } from '../observability' import { HARNESS_EVENTS, InputRejectedError } from '../types' import { timeOf } from '../waits' import { @@ -20,12 +21,14 @@ import type { AgentSendOptions, CancelOptions, ChatTurnResult, + Cursor, HarnessInput, Operation, Principal, Receipt, ResetOptions, ResolveOptions, + SendForRunOptions, SendOptions, TurnOverrides, UserInput, @@ -125,6 +128,18 @@ export abstract class SessionRequests< options?: SendOptions | AgentSendOptions, ): Operation { if (options?.to !== undefined) return this.sendToRun(message, options) + return this.sendForRun(message, options) + } + + /** + * @internal `send` of a message to the session, with the AG-UI `runId` of + * a `POST run` request: the turn runs as it. The HTTP handler and an + * offered prompt call it. + */ + override sendForRun( + message: UserInput, + options?: SendForRunOptions, + ): Operation { const context = options?.context const mode = options?.mode ?? 'prompt' const input: HarnessInput = @@ -158,7 +173,7 @@ export abstract class SessionRequests< */ protected startTurn( input: HarnessInput & { op: 'prompt' | 'steer' | 'followUp' }, - options?: SendOptions & { wake?: boolean }, + options?: SendForRunOptions & { wake?: boolean }, ): Operation { const { message } = input const inputId = options?.inputId ?? createInputId() @@ -196,11 +211,16 @@ export abstract class SessionRequests< } const operation = passed ?? this.createTurnOperation(runId) this.bindTurn(inputId, operation) + const traceparent = options?.traceparent void this.admitting(() => this.accept(inputId, stored, principal, { // An offered turn runs as the operation of this host. runId: operation.id, turn: { overrides, systemPreamble, ephemeral }, + // As the W3C rules say, an invalid `traceparent` is dropped. + ...(traceparent !== undefined && parseTraceparent(traceparent) + ? { traceparent } + : {}), }).then( async (admission) => { if (admission !== 'new') { @@ -275,6 +295,7 @@ export abstract class SessionRequests< ephemeral, ...sent, }) + this.publishQueued(inputId, 'steer') return } this.enqueueTurn({ @@ -405,7 +426,9 @@ export abstract class SessionRequests< * * @example * ```ts - * await session.append([{ type: 'app.signal', text: 'The build failed.' }]) + * await session.append({ + * records: [{ type: 'app.signal', text: 'The build failed.' }], + * }) * await session.continue({ inputId: 'after-signal' }) * ``` */ @@ -773,15 +796,15 @@ export abstract class SessionRequests< /** * @internal Cancel the queued and running work that the user `userId` - * started: `host.sessions.unshare` stops it. `principal` is who cancels, - * and the log keeps it. + * started: `host.sessions.unshare` and `host.revoke` stop it. `principal` + * is who cancels, and the log keeps it. */ async cancelWorkOf({ userId, principal, }: { userId: string - principal: Principal + principal: Principal | undefined }): Promise { const theirs = [...this.operations.values()].filter( (operation) => @@ -789,7 +812,10 @@ export abstract class SessionRequests< this.startedBy(operation.id).principalId === userId, ) for (const operation of theirs) { - await this.cancel({ operationId: operation.id, principal }) + await this.cancel({ + operationId: operation.id, + ...(principal ? { principal } : {}), + }) } } @@ -803,22 +829,34 @@ export abstract class SessionRequests< * @example * ```ts * await session.send('Then write the tests.', { mode: 'followUp', inputId: 'tests' }) - * await session.cancelInput('tests') + * await session.cancelInput({ inputId: 'tests' }) * ``` */ - async cancelInput( - target: string, - options?: { - principal?: Principal - /** @internal The id of an offered input, so a replay is a duplicate. */ - inputId?: string - }, - ): Promise { - const inputId = options?.inputId ?? createInputId() + cancelInput(options: { + inputId: string + principal?: Principal + }): Promise { + return this.applyCancelInput(options) + } + + /** + * `cancelInput`, with `controlId`: the id of an offered input, so a replay + * is a duplicate. + */ + protected override async applyCancelInput({ + inputId: target, + principal, + controlId, + }: { + inputId: string + principal?: Principal + controlId?: string + }): Promise { + const inputId = controlId ?? createInputId() const admission = await this.accept( inputId, { op: 'cancelInput', inputId: target }, - options?.principal, + principal, ) if (admission !== 'new') return this.duplicateReceipt(inputId, admission) const taken = this.takeWaiting(target) ?? this.takeQueuedContinue(target) @@ -855,7 +893,19 @@ export abstract class SessionRequests< * await session.setDelivery({ inputId: 'style', delivery: 'steer' }) * ``` */ - async setDelivery({ + setDelivery(options: { + inputId: string + delivery: WaitingInput['delivery'] + principal?: Principal + }): Promise { + return this.applySetDelivery(options) + } + + /** + * `setDelivery`, with `controlId`: the id of an offered input, so a replay + * is a duplicate. + */ + protected override async applySetDelivery({ inputId: target, delivery, principal, @@ -864,7 +914,6 @@ export abstract class SessionRequests< inputId: string delivery: WaitingInput['delivery'] principal?: Principal - /** @internal The id of an offered input, so a replay is a duplicate. */ controlId?: string }): Promise { const inputId = controlId ?? createInputId() @@ -962,8 +1011,10 @@ export abstract class SessionRequests< /** * The session API of plugins. With `principal` (a command, a turn plugin), - * it acts for that user. Without it (a session plugin), its `principal` is - * the principal that opened the session. With `operationId` (a command, a + * it acts for that user. Without it (a session plugin), it acts for the + * owner of the session entry, else for the principal that opened the + * session. A `principal` option of `send` and `note` wins: a plugin passes + * the user whose work the turn continues. With `operationId` (a command, a * turn plugin), its questions belong to that operation. */ protected override pluginApi( @@ -971,13 +1022,11 @@ export abstract class SessionRequests< metadataSchema?: SchemaInput, operationId?: string, ): PluginSessionApi { - // ponytail: D9. Without `principal`, `send`, `note`, and `ask` follow the - // running turn: its sender, and its operation for a question. PR 7 - // rebuilds them on the run that calls. - const sender = () => principal ?? this.sender() // A session plugin acts for the owner of the entry, not for whoever // opened the session first on this host. The owner is known only after - // the session read its entry, so it is read at each call. + // the session read its entry, so it is read at each call. A turn of + // `send` or `note` never follows the sender of the running turn: that + // user did not start the work of the plugin. const owner = () => principal ?? this.owner ?? this.principal return { threadId: this.threadId, @@ -986,23 +1035,24 @@ export abstract class SessionRequests< }, snapshot: () => this.snapshot(), // A plugin turn always waits its turn. It cannot see a busy rejection. - send: (text) => + send: (text, options) => this.startTurn( { op: 'prompt', message: text, busy: 'queue' }, - { principal: sender() }, + { principal: options?.principal ?? owner() }, ), note: async (text, options) => { await this.addNote(text) if (options?.wake) { const receipt = await this.startTurn( { op: 'followUp', message: text }, - { principal: sender(), wake: true }, + { principal: options.principal ?? owner(), wake: true }, ).receipt if (receipt.status === 'rejected') throw new InputRejectedError(receipt) } }, transcript: () => this.transcript(), + keptRecords: async (type) => this.keptRecords(type), replaceTranscript: (messages) => this.messages.saveThread(this.threadId, messages), entry: () => this.index.get(this.indexKey), @@ -1128,13 +1178,29 @@ export abstract class SessionRequests< // Chat turns // =========================== - protected createTurnOperation(id?: string): OperationImpl { + /** + * `id`: a stored or chosen run id. `startedCursor`: the start of a run + * that a restart rebuilds from the log. + */ + protected createTurnOperation( + id?: string, + startedCursor?: Cursor, + ): OperationImpl { const operation = new OperationImpl( 'chat', this.feed, - (target) => this.cancel({ operationId: target.id }), + (target) => { + // A cancel on the turn itself acts as its sender, not as the + // principal that opened the session. A handle passes its own. + const principal = this.principalOfRun(target.id) + return this.cancel({ + operationId: target.id, + ...(principal ? { principal } : {}), + }) + }, undefined, id, + startedCursor, ) this.track(operation as OperationImpl) return operation @@ -1213,7 +1279,11 @@ export abstract class SessionRequests< } const live = this.turnOperations.get(inputId) if (live) return live - const replay = this.createTurnOperation(logged?.operationId) + // A join of the replayed turn reads its events from its run start. + const replay = this.createTurnOperation( + logged?.operationId, + this.runStart(logged?.operationId), + ) this.answerDuplicate(replay, inputId, 'duplicate') return replay } diff --git a/packages/ai-harness/src/session/state.ts b/packages/ai-harness/src/session/state.ts index 90735763b6..62b15828ab 100644 --- a/packages/ai-harness/src/session/state.ts +++ b/packages/ai-harness/src/session/state.ts @@ -1,6 +1,5 @@ import { INTERRUPT_BINDING_METADATA_KEY, - chat, convertSchemaToJsonSchema, validateWithStandardSchema, } from '@tanstack/ai' @@ -8,11 +7,10 @@ import { UNOWNED_EPOCH, loadChildTranscript } from '@tanstack/ai-persistence' import { checkConfigValue } from '../config' import { transcriptText } from '../first-party/session-tools' import { LogOwnedElsewhereError } from '../ownership' -import { callUsage } from '../usage' import { isRecord } from '../utils' import { jsonSchemaIssues } from '../json-schema' import { changedEntry } from '../session-entry' -import { stableJson, threadOf, visibleMessages } from '../log' +import { stableJson, visibleMessages } from '../log' import { isHumanTarget } from '../log-records' import { mediaIdOf, mediaOfMessage, mediaPart } from '../media-ref' import { @@ -21,7 +19,9 @@ import { RevertFiles, RevertStanding, SessionMetadata, + SessionScrub, SessionSignal, + ThreadEnds, WorkspaceLog, } from '../plugins' import { HARNESS_EVENTS } from '../types' @@ -57,6 +57,7 @@ import type { SessionIndexEntry, } from '@tanstack/ai-persistence' import type { RevertState } from '../log' +import type { ThreadEndHandler } from '../plugins' import type { HandoffRecord, RecordPrincipal, @@ -76,7 +77,7 @@ import type { SessionEntryChange, SessionEntryUpdate } from '../session-entry' import type { ChatTurnResult, ConfigureOptions, - HarnessInput, + OfferedHarnessInput, MediaRecord, MediaUpload, PendingQuestion, @@ -120,6 +121,8 @@ export abstract class SessionState< protected readonly waitEnds = new Map void>() /** The secret answers to `step.ask` waits that this host kept, by wait id. */ protected readonly keptAnswers = new Map() + /** The handlers that plugins added to `ThreadEnds`, oldest first. */ + protected readonly threadEnds = new Set() /** The stores that first-party plugins read as capabilities. */ protected storeCapabilities() { @@ -128,6 +131,13 @@ export abstract class SessionState< if (metadata) values.provide(SessionMetadata, metadata) values.provide(RevertStanding, () => this.reverted !== undefined) values.provide(SessionSignal, this.lifetime.signal) + values.provide(SessionScrub, (text) => this.scrub(text)) + values.provide(ThreadEnds, { + add: (handler) => { + this.threadEnds.add(handler) + return () => this.threadEnds.delete(handler) + }, + }) values.provide(BackgroundJobs, { started: (job) => this.logJob(job, false), ended: (job) => this.logJob(job, true), @@ -151,6 +161,14 @@ export abstract class SessionState< ? [] : [{ jobId, backendId }], ), + stranded: ({ operationId, hostId }) => + this.feed.publish( + operationId, + customEvent(HARNESS_EVENTS.stranded, { + operationId, + ...(hostId !== undefined ? { hostId } : {}), + }), + ), }) return values } @@ -293,7 +311,7 @@ export abstract class SessionState< } /** The record of a media file of this thread, or `null` when it is not found. */ - getMedia(id: string) { + getMedia({ id }: { id: string }) { return this.mediaStore.get(id) } @@ -313,7 +331,7 @@ export abstract class SessionState< * has no server, so this is a data URL for a file up to 1 MB, and `{}` for * a bigger file (use `loadMedia`). */ - async mediaUrl(id: string) { + async mediaUrl({ id }: { id: string }) { const url = await this.mediaStore.dataUrl(id) const answer: { url?: string; expiresAt?: number } = url === undefined ? {} : { url } @@ -341,7 +359,7 @@ export abstract class SessionState< * call it. */ async adoptMessages( - source: Pick, + source: SessionImpl, messages: ReadonlyArray, records: ReadonlyArray = [], ): Promise { @@ -349,12 +367,13 @@ export abstract class SessionState< const copies = new Map() const copy = async (id: string) => { if (!copies.has(id)) { - const record = await source.getMedia(id) + const record = await source.getMedia({ id }) copies.set( id, record && (await this.putMedia({ - body: await source.loadMedia({ id }), + // A stream, so the copy does not hold the whole file in memory. + body: await source.mediaStore.open(id), mimeType: record.mimeType, name: record.name, })), @@ -394,6 +413,7 @@ export abstract class SessionState< messages: [...this.writer.state.messages, ...forked], records, }) + for (const record of records) this.publishTransferred(record) } /** @@ -433,17 +453,23 @@ export abstract class SessionState< record: Extract, ): Promise { if (this.transferOf(record.transferId)) return - await this.appendOrOffer(record, `transfer:${record.transferId}`) + const isAppended = await this.appendOrOffer( + record, + `transfer:${record.transferId}`, + ) + // An offered record gets its event from the owner, when it appends it. + if (isAppended) this.publishTransferred(record) } /** * Append `record` to the log of this thread. When another host owns the * log, offer it to that host as `inputId`: the owner appends it as it is. + * Resolves `true` when this host appended it. */ private async appendOrOffer(record: LogRecord, inputId: string) { try { await this.writer.append([record]) - return + return true } catch (error) { if (!(error instanceof LogOwnedElsewhereError)) throw error } @@ -453,13 +479,14 @@ export abstract class SessionState< `The input lane has another record with id ${JSON.stringify(inputId)}.`, ) } + return false } /** * @internal A summary of `messages` for `host.transfer`, by the model of * this harness, with the keys of `principal`. The usage of its model calls - * goes to the log of this thread at once, so it counts also when the - * transfer fails later. + * goes to the log of this thread at once, under the source `transfer` and + * the label `summary`, so it counts also when the transfer fails later. */ async transferSummary( principal: Principal | undefined, @@ -474,41 +501,40 @@ export abstract class SessionState< `A transfer with messages: 'summary' needs a model: harness ${JSON.stringify(this.harness.name)} has no adapter.`, ) } - const usage: Array = [] - try { - const { text } = await chat({ - adapter: await this.keysOf(principal).adapter(model), + // The call is part of no run. The run gives it the keys of `principal` + // and the signal of the session, and `count` takes its usage. + const run = this.runContext({ + principal, + runId: 'transfer', + signal: this.lifetime.signal, + }) + const result = await this.modelChat({ + adapter: model, + options: { messages: [ { role: 'user', content: `${TRANSFER_SUMMARY_PROMPT}\n\n${transcriptText(messages)}`, }, ], - stream: false, - middleware: [ - { - name: 'harness:transfer-usage', - onUsage: (ctx, counts) => { - usage.push({ - type: 'harness.usage', - model: `${ctx.provider}/${ctx.model}`, - ...storedPrincipal(principal), - usage: callUsage(counts), - }) - }, - }, - ], - }) - if (!text.trim()) { - throw new Error('The summary model returned no text. Nothing moved.') - } - return text - } finally { + }, + run, + source: 'transfer', + label: 'summary', // A new id for each call: a retry that calls the model again pays again. - for (const record of usage) { - await this.appendOrOffer(record, createInputId()) - } + count: async (call) => { + await this.appendOrOffer( + { type: 'harness.usage', ...call }, + createInputId(), + ) + }, + }) + const text = + isRecord(result) && typeof result.text === 'string' ? result.text : '' + if (!text.trim()) { + throw new Error('The summary model returned no text. Nothing moved.') } + return text } // =========================== @@ -528,7 +554,8 @@ export abstract class SessionState< result[key] = { option: entry.option, owner: entry.owner, - value: this.configValue(key), + // The latest value, also while a turn runs with its own copy. + value: this.configValue(key, this.configValues), } } return result @@ -539,17 +566,28 @@ export abstract class SessionState< * who changes it. Default: the principal that opened the session. The log * keeps it. */ - async setConfig({ + setConfig(options: { + key: string + value: unknown + principal?: Principal + }): Promise { + return this.applyConfig(options) + } + + /** + * `setConfig`, with `controlId`: the id of an offered input, so a replay + * is a duplicate. + */ + protected override async applyConfig({ key, value, - inputId = createInputId(), + controlId: inputId = createInputId(), principal, }: { key: string value: unknown - /** @internal The id of an offered input, so a replay is a duplicate. */ - inputId?: string principal?: Principal + controlId?: string }): Promise { const admission = await this.accept( inputId, @@ -618,15 +656,19 @@ export abstract class SessionState< principal, ) if (admission !== 'new') return this.duplicateReceipt(inputId, admission) + // The store took the input, so this host owns the log, and its fold has + // the settings that other hosts wrote. The change merges into them. + const current = this.settingsFromLog() let next: ThreadSettings try { - next = changeSettings(this.threadSettings, settings, this.settingNames()) + next = changeSettings(current, settings, this.settingNames()) + this.checkRemovable(next) } catch (error) { const reason = error instanceof Error ? error.message : String(error) this.reject(inputId, reason) return this.keep({ inputId, status: 'rejected', reason }) } - const movedFrom = this.threadSettings.cwd + const movedFrom = current.cwd this.threadSettings = next // Tell the model about a new folder, so it stops using the old paths. // A new thread gets no note: some providers refuse a conversation that @@ -656,6 +698,43 @@ export abstract class SessionState< return this.keep({ inputId, status: 'accepted' }) } + /** + * Throws when the plugins that `settings` removes provide a capability + * that a kept plugin (of any lifetime) or a harness middleware requires. + */ + protected checkRemovable(settings: ThreadSettings) { + const removed = new Set(settings.plugins?.remove ?? []) + if (removed.size === 0) return + const owners = new Map( + this.plugins.flatMap((plugin) => + (plugin.provides ?? []).map((handle) => [handle, plugin.name] as const), + ), + ) + const needs = [ + ...this.plugins + .filter((plugin) => !removed.has(plugin.name)) + .map((plugin) => ({ + who: `plugin ${plugin.name}`, + requires: plugin.requires, + })), + ...(this.harness.middleware ?? []).map((middleware) => ({ + who: middleware.name + ? `harness middleware ${middleware.name}` + : 'a harness middleware', + requires: middleware.requires, + })), + ] + for (const { who, requires } of needs) { + for (const handle of requires ?? []) { + const owner = owners.get(handle) + if (owner === undefined || !removed.has(owner)) continue + throw new Error( + `Cannot remove plugin ${owner}: ${who} requires its capability "${handle.capabilityName}".`, + ) + } + } + } + /** The model and plugin names that a setting can use. */ protected settingNames(): SettingNames { return { @@ -671,7 +750,7 @@ export abstract class SessionState< owner: string input?: unknown }> { - return [...(this.sessionPlugins?.commands ?? [])].map(([name, entry]) => ({ + return [...(this.threadPlugins()?.commands ?? [])].map(([name, entry]) => ({ name, description: entry.command.description, owner: entry.owner, @@ -734,7 +813,11 @@ export abstract class SessionState< * child of this thread too. Empty for an unknown id, and for a child of * another thread. */ - async agentTranscript(subagentRunId: string): Promise> { + async agentTranscript({ + subagentRunId, + }: { + subagentRunId: string + }): Promise> { return ( (await loadChildTranscript(this.engine, subagentRunId, this.threadId)) ?? [] @@ -775,9 +858,23 @@ export abstract class SessionState< * hidden messages for good. The revert state is saved, so it survives a * restart. Refused while the session is not idle, or another host runs * the thread (`'busy'`), and for a message that is not in the transcript - * (`'unknown_message'`). + * (`'unknown_message'`). `principal` is who reverts. Default: the + * principal that opened the session. The `harness.revert` record keeps + * it. A handle of `host.open` checks that it is the principal of the + * handle. + * + * @example + * ```ts + * await session.revert({ messageId: message.id }) + * ``` */ - async revert(messageId: string): Promise { + async revert({ + messageId, + principal, + }: { + messageId: string + principal?: Principal + }): Promise { const inputId = createInputId() if (!(await this.ownsIdleThread())) { return { inputId, status: 'rejected', reason: 'busy' } @@ -790,13 +887,15 @@ export abstract class SessionState< // A revert that stands ends first, so its files come back. await this.endRevert() const handler = this.sessionPlugins?.values.get(RevertFiles) - const files = await handler?.revert(messages.slice(at + 1), (type) => - this.hostRecords(type), + const files = await handler?.revert( + messages.slice(at + 1), + async (type) => this.keptRecords(type), + this.revertRun(principal ?? this.principal), + ) + await this.saveRevert( + { messageId, ...(files !== undefined ? { files } : {}) }, + principal ?? this.principal, ) - await this.saveRevert({ - messageId, - ...(files !== undefined ? { files } : {}), - }) return { inputId, status: 'accepted' } }) } @@ -823,11 +922,24 @@ export abstract class SessionState< if (reverted.files !== undefined) { await this.sessionPlugins?.values .get(RevertFiles) - ?.unrevert(reverted.files) + ?.unrevert(reverted.files, this.revertRun(this.principal)) } await this.saveRevert(undefined) } + /** + * A run for the files of a revert, in the thread of this session. It is + * part of no turn, so a workspace that no turn opened on this host (after + * a restart) opens again from the log. + */ + protected revertRun(principal: Principal | undefined) { + return this.runContext({ + principal, + runId: `revert:${this.threadId}`, + signal: this.lifetime.signal, + }) + } + /** Run `fn` while no new turn starts. A held turn starts after it. */ protected async holding(fn: () => Promise) { this.holdTurns++ @@ -861,11 +973,18 @@ export abstract class SessionState< await this.saveRevert(undefined) } - /** Keep the revert state in the log. */ - protected async saveRevert(revert: RevertState | undefined) { + /** Keep the revert state in the log, with who reverted. */ + protected async saveRevert( + revert: RevertState | undefined, + principal?: Principal, + ) { this.reverted = revert await this.writer.append([ - { type: 'harness.revert', revert: revert ?? null }, + { + type: 'harness.revert', + revert: revert ?? null, + ...storedPrincipal(principal), + }, ]) this.feed.publish( 'session', @@ -875,25 +994,15 @@ export abstract class SessionState< ) } - /** The host records of `type` of this thread, from the log, oldest first. */ - protected async hostRecords(type: string) { - const found: Array> = [] - let after = 0 - for (;;) { - const entries = await this.log.store.read({ - logId: this.storeLogId, - after, - limit: 256, - }) - for (const { seq, record } of entries) { - after = seq - const isOwn = threadOf(record, this.threadId) === this.threadId - if (isOwn && record.type === type) found.push(record) - } - // A store can return fewer than `limit` before the end: stop only at - // an empty page. - if (entries.length === 0) return found - } + /** + * The kept host records of `type` of this thread, oldest first: the last + * ones that the `keepRecords` of the plugins name. The log snapshot keeps + * them, so they stay after a truncate of the log. + */ + protected keptRecords(type: string): Array { + return this.writer.state.keptRecords.flatMap(({ record }) => + record.type === type ? [record] : [], + ) } /** The commands, settings, and tools of this session, for a UI. */ @@ -915,19 +1024,30 @@ export abstract class SessionState< * wait, by its id. `principal` is who answers. Default: the principal that * opened the session. The log keeps it. */ - async answer({ + answer(options: { + questionId: string + value: unknown + principal?: Principal + }): Promise { + return this.applyAnswer(options) + } + + /** + * `answer`, with `controlId` (the id of an offered input, so a replay is a + * duplicate) and `ref` (a secret answer that another host saved, see + * `answerOffer`). + */ + protected override async applyAnswer({ questionId, value, - inputId = createInputId(), + controlId: inputId = createInputId(), principal, ref, }: { questionId: string value: unknown - /** @internal The id of an offered input, so a replay is a duplicate. */ - inputId?: string principal?: Principal - /** @internal A secret answer that another host saved. See `answerOffer`. */ + controlId?: string ref?: string }): Promise { // A secret (a key or a password) stays out of the log. A restart does @@ -1087,7 +1207,7 @@ export abstract class SessionState< // ponytail: the agents of turn plugins are not known outside a turn. const names = [ ...(this.harness.agents ?? []), - ...(this.sessionPlugins?.agents ?? []), + ...(this.threadPlugins()?.agents ?? []), ].map((agent) => agent.name) if (handBack !== 'main' && !names.includes(handBack)) { return { reason: 'unknown_agent' } @@ -1181,7 +1301,7 @@ export abstract class SessionState< value: unknown inputId: string sender: Principal | undefined - }): Promise { + }): Promise { // The wait of another host: read what that host wrote first. if (!this.questions.has(questionId)) await this.writer.catchUp() const logged = this.loggedWait(questionId) @@ -1269,8 +1389,47 @@ export abstract class SessionState< ] } + /** + * Publish `harness.wait.settled` after the append that ended the wait + * `waitId`. An interrupt id has no wait in the log, so it gets no event. + */ + protected publishWaitSettled(waitId: string) { + const ended = this.writer.state.thread.waits.get(waitId) + if (!ended || ended.outcome === undefined) return + this.feed.publish( + ended.wait.runId, + customEvent(HARNESS_EVENTS.waitSettled, { + waitId, + outcome: ended.outcome, + }), + ) + } + + /** + * The `harness.wait` span of the wait `waitId` that ended. A wait opens + * and ends on different hosts, so the span is written now, from the time + * in its log record. + */ + private traceWait(waitId: string) { + const ended = this.tracing && this.writer.state.thread.waits.get(waitId) + if (!ended || ended.outcome === undefined) return + const span = this.tracing?.span( + 'harness.wait', + { + 'harness.run_id': ended.wait.runId, + 'harness.wait.id': waitId, + 'harness.wait.kind': ended.wait.kind, + 'harness.wait.outcome': ended.outcome, + }, + { runId: ended.wait.runId, startTime: ended.at }, + ) + this.tracing?.safely(() => span?.end(ended.settledAt)) + } + /** After the log ended the wait `waitId`: drop its secrets, and move the timer. */ protected async afterWaitEnd(waitId: string) { + this.traceWait(waitId) + this.publishWaitSettled(waitId) // A `step.ask` of this host that waits for it goes on. this.waitEnds.get(waitId)?.() this.armSchedules() @@ -1289,6 +1448,10 @@ export abstract class SessionState< * also when no host has the thread open. */ protected async writeWait(...records: Array) { + // A wait that the log has already (a run again) gets no second event. + const fresh = records.filter( + ({ wait }) => !this.writer.state.thread.waits.has(wait.waitId), + ) // One append: a run that waits for several opens all or none. await this.writer.append( records.flatMap((record): Array => { @@ -1307,11 +1470,27 @@ export abstract class SessionState< ] }), ) + for (const { wait } of fresh) this.publishWaitOpened(wait) if (records.some((record) => record.wait.expiresAt !== undefined)) { this.armSchedules() } } + /** Publish `harness.wait.opened`, after the append of the wait. */ + private publishWaitOpened(wait: Wait) { + const { waitId, kind, runId, message, expiresAt } = wait + this.feed.publish( + runId, + customEvent(HARNESS_EVENTS.waitOpened, { + waitId, + kind, + runId, + ...(message !== undefined ? { message } : {}), + ...(expiresAt !== undefined ? { expiresAt } : {}), + }), + ) + } + /** * @internal Open a wait that no promise holds, for the run `runId`. The * session stays idle: the log keeps the wait, and `answer` from any host @@ -1368,6 +1547,7 @@ export abstract class SessionState< { type: 'harness.wait', wait, at: Date.now() }, { type: 'harness.handoff', ...kept }, ]) + this.publishWaitOpened(wait) return kept } @@ -1605,34 +1785,69 @@ export abstract class SessionState< /** Read the config values of this thread from the log. */ protected loadConfig() { - const stored = this.writer.state.thread.config - for (const [key, value] of stored) { + for (const [key, value] of this.configSnapshot()) { + this.configValues.set(key, value) + } + } + + /** + * The config values in the log now. A turn keeps them until it ends, so a + * change while it runs applies at the next turn. + */ + protected configSnapshot(): Map { + const values = new Map() + for (const [key, value] of this.writer.state.thread.config) { const entry = this.sessionPlugins?.config.get(key) if (!entry) continue try { - this.configValues.set(key, checkConfigValue(key, entry.option, value)) + values.set(key, checkConfigValue(key, entry.option, value)) } catch { // A stored value an option no longer accepts falls back to the default. } } + return values } /** Read the stored settings of this thread from the log. */ protected loadSettings() { + this.threadSettings = this.settingsFromLog(this.threadSettings) + } + + /** The stored settings in the fold of the log, over `base`. */ + protected settingsFromLog(base: ThreadSettings = {}): ThreadSettings { const stored = this.writer.state.thread.settings const names = this.settingNames() + let settings = base for (const [key, value] of Object.entries(stored)) { try { - this.threadSettings = changeSettings( - this.threadSettings, - { [key]: value }, - names, - ) + const next = changeSettings(settings, { [key]: value }, names) + this.checkRemovable(next) + settings = next } catch { // A value the harness no longer takes (a removed model or plugin) // is dropped. } } + return settings + } + + /** + * Another host wrote to the log, and this log folded its records. Read + * the settings, the config values, and the interrupted turn from the fold + * again, so a turn and a change on this host use what that host wrote. + */ + protected syncFromLog() { + // While the plugins mount (an open or a reload), the open reads them. + if (!this.sessionPlugins) return + this.threadSettings = this.settingsFromLog() + this.configValues.clear() + for (const [key, value] of this.configSnapshot()) { + this.configValues.set(key, value) + } + const logged = this.writer.state.thread.interrupted ?? undefined + if (stableJson(logged ?? null) !== stableJson(this.interrupted ?? null)) { + this.interrupted = logged + } } /** Restore the interrupted turn that the log keeps for this thread. */ @@ -1908,6 +2123,7 @@ export abstract class SessionState< // A cancel or the close ended it before the log had it. this.writer .append(this.waitEndRecords(questionId, 'cancelled')) + .then(() => this.publishWaitSettled(questionId)) .catch(() => {}) }, (error: unknown) => { diff --git a/packages/ai-harness/src/session/turn-parts.ts b/packages/ai-harness/src/session/turn-parts.ts index 4165e88561..17d0653940 100644 --- a/packages/ai-harness/src/session/turn-parts.ts +++ b/packages/ai-harness/src/session/turn-parts.ts @@ -1,5 +1,6 @@ import { LogRecordsCapability, + chat, convertSchemaToJsonSchema, isStandardSchema, provideLogRecords, @@ -43,7 +44,7 @@ import type { LogWriter } from '../log' import type { UsageCall } from '../usage' import type { AnyHarness } from '../define' import type { TurnAdditions } from '../turn' -import type { MountedPlugins } from '../plugins' +import type { MountedPlugins, PluginModelCall } from '../plugins' import type { ChatTurnResult, MediaRecord, @@ -54,6 +55,9 @@ import type { } from '../types' import type { QueuedTurn, AgentRunImpl, AgentChain } from './types' +/** How long a `ctx.model.chat` call runs before it stops, by default. */ +const MODEL_CALL_TIMEOUT_MS = 120_000 + /** * A layer of {@link HarnessSession}: the tools, middleware, and joins a turn * uses. @@ -306,14 +310,36 @@ export abstract class SessionTurnParts< } } + /** + * Run the `onTurnStart` hook of each plugin, in order, for `turn`. A hook + * that fails is skipped, and clients get a `harness.plugin.warning` event. + */ + protected async runTurnStarts( + turnStarts: MountedPlugins['turnStarts'], + operation: OperationImpl, + turn: TurnInfo, + ) { + for (const { start, owner } of turnStarts) { + try { + await start(turn) + } catch (error) { + this.warn(operation, owner, error) + } + } + } + /** * Ask each plugin for the tools it found since the last turn, for `turn`. * A plugin that fails (for example an MCP server that is down) is skipped, * and clients get a `harness.plugin.warning` event. + * + * `taken` has the owner of each tool name. A found tool with a taken name + * is skipped: the earlier tool stays. The session warns once for each + * name and pair of owners. */ protected async discoverTools( discoverers: MountedPlugins['discoverers'], - taken: Set, + taken: Map, operation: OperationImpl, turn: TurnInfo, ): Promise> { @@ -322,8 +348,12 @@ export abstract class SessionTurnParts< try { const tools = await discover(turn) for (const tool of tools) { - if (taken.has(tool.name)) continue - taken.add(tool.name) + const first = taken.get(tool.name) + if (first !== undefined) { + this.warnToolClash(operation, tool.name, first, owner) + continue + } + taken.set(tool.name, owner) found.push(tool) } } catch (error) { @@ -333,19 +363,49 @@ export abstract class SessionTurnParts< return found } + /** The tool clashes this session warned about: name and both owners. */ + protected readonly toolClashes = new Set() + + /** Warn once that `owner` gives a tool `name` that `first` gives too. */ + protected warnToolClash( + operation: OperationImpl, + name: string, + first: string, + owner: string, + ) { + // A plugin that finds the same tool again on each turn is no clash. + if (first === owner) return + const key = JSON.stringify([name, first, owner]) + if (this.toolClashes.has(key)) return + this.toolClashes.add(key) + const label = (who: string) => + who === 'the harness' ? who : `plugin ${who}` + this.warn( + operation, + owner, + `Two tools are named "${name}": the tool of ${label(first)} is used, and the tool of ${label(owner)} is not.`, + ) + } + /** * 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. */ protected async prepareTools( preparers: MountedPlugins['preparers'], - turn: { tools: Array; model: string }, + turn: { tools: Array; model: string; run: RunContext }, operation: OperationImpl, ): Promise> { let prepared = turn.tools for (const { prepare, owner } of preparers) { try { - prepared = [...(await prepare({ tools: prepared, model: turn.model }))] + prepared = [ + ...(await prepare({ + tools: prepared, + model: turn.model, + run: turn.run, + })), + ] } catch (error) { this.warn(operation, owner, error) } @@ -402,8 +462,8 @@ export abstract class SessionTurnParts< context.runId = runId context.detach = (work) => new Promise((moved) => { - this.detachable.set(id, () => { - this.logJob({ jobId: id }, false) + this.detachable.set(id, (mover) => { + this.logJob({ jobId: id, principal: mover }, false) moved( `The job moved to the background. Its id is ${id}. You get a note when it ends.`, ) @@ -559,16 +619,108 @@ export abstract class SessionTurnParts< this.warn(operation, 'harness:usage', error) return } + this.tracing?.usage(call, (runId) => this.parentRunOf(runId)) operation.publish( customEvent(HARNESS_EVENTS.usage, { model: call.model, ...(call.principal ? { sender: call.principal.id } : {}), + ...(call.source !== undefined ? { source: call.source } : {}), + ...(call.label !== undefined ? { label: call.label } : {}), usage: call.usage, total: this.usage().total, }), ) } + /** + * One model call of `ctx.model.chat`, for the run `call.run`: with the + * keys of its user, and a signal that aborts when the run stops or after + * the timeout. Each model call counts under `call.source`, in the session + * and in the run, and the budget of the run applies. A stopped call + * rejects with the reason of the stop. + */ + protected async modelChat(call: PluginModelCall): Promise { + const { run, options } = call + // The setting names a key of `harness.models`: a setting is checked + // when it is set, and a stored one that the harness lost is dropped. + const named = this.threadSettings.model + const model = + call.adapter ?? + (named === undefined ? undefined : this.harness.models?.[named]) ?? + this.harness.adapter + if (!model) { + throw new Error( + `ctx.model.chat needs an adapter: harness ${JSON.stringify(this.harness.name)} has no adapter.`, + ) + } + const signal = AbortSignal.any([ + run.signal, + AbortSignal.timeout(call.timeoutMs ?? MODEL_CALL_TIMEOUT_MS), + ]) + const abortController = new AbortController() + const stop = () => abortController.abort(signal.reason) + if (signal.aborted) stop() + else signal.addEventListener('abort', stop, { once: true }) + const counted: Array> = [] + try { + const result = await chat({ + ...options, + adapter: await run.keys.adapter(model), + stream: false, + abortController, + middleware: [ + { + name: 'harness:model-usage', + onUsage: (ctx, usage) => { + const made: UsageCall = { + model: `${ctx.provider}/${ctx.model}`, + ...storedPrincipal(run.principal), + source: call.source, + ...(call.label !== undefined ? { label: call.label } : {}), + usage: callUsage(usage), + } + counted.push( + call.count + ? call.count(made) + : this.countModelCall(run.runId, { + ...made, + runId: run.runId, + }), + ) + }, + }, + ], + }) + // chat() can end with no error when it is aborted. + abortController.signal.throwIfAborted() + return result + } finally { + signal.removeEventListener('abort', stop) + await Promise.all(counted) + } + } + + /** + * Count one model call of `ctx.model.chat` for the run `runId`, then + * check the budgets. A run that settled, or that this session no longer + * keeps, gets the record only: a run sends no event after RUN_FINISHED. + */ + private async countModelCall(runId: string, call: UsageCall) { + const operation = this.operations.get(runId) + if (operation && !operation.isSettled()) { + await this.countUsage(operation, call) + this.checkBudget(operation) + return + } + try { + await (this.childLink + ? this.childLink.countUsage(call) + : this.writer.append([{ type: 'harness.usage', ...call }])) + } catch (error) { + this.warnSession('harness:usage', error) + } + } + /** * Take the waiting steers that join now: the prefix up to the first one * that has an abort request, or that `turn.canJoin` refuses. Without @@ -825,7 +977,7 @@ export abstract class SessionTurnParts< typeof content === 'string' ? [] : content.map(mediaIdOf).filter((id) => id !== undefined) - const records = await Promise.all(ids.map((id) => this.getMedia(id))) + const records = await Promise.all(ids.map((id) => this.getMedia({ id }))) const media = records.filter((record) => record !== null) const message: ModelMessage = { id: createMessageId(), diff --git a/packages/ai-harness/src/session/turn.ts b/packages/ai-harness/src/session/turn.ts index b84e0e8e15..dcc29b0ef8 100644 --- a/packages/ai-harness/src/session/turn.ts +++ b/packages/ai-harness/src/session/turn.ts @@ -31,11 +31,12 @@ import { storedPrincipal, turnContext, withRun, - withoutPlugins, settingTools, + toolOwners, resetMarker, resetContext, resetCut, + contextPromptCall, ephemeralCall, turnEphemeral, acceptedKinds, @@ -140,7 +141,8 @@ export abstract class SessionTurn< } return } - const entry = this.sessionPlugins?.commands.get(name) + // A command of a plugin that this thread removed is unknown here. + const entry = this.threadPlugins()?.commands.get(name) if (!entry) { this.reject(inputId, 'unknown_command') operation.fail('failed', new Error(`Unknown command: ${name}`)) @@ -188,7 +190,7 @@ export abstract class SessionTurn< input: unknown principal: Principal | undefined }): boolean { - const entry = this.sessionPlugins?.commands.get(input.name) + const entry = this.threadPlugins()?.commands.get(input.name) if (!entry) return false const operation = new OperationImpl( 'command', @@ -332,10 +334,12 @@ export abstract class SessionTurn< this.activeTurn = next.operation this.activeResume = next.resume !== undefined this.turnPrincipal = next.principal + this.turnConfig = this.configSnapshot() const ended = () => { this.activeResume = false this.activeTurn = undefined this.turnPrincipal = undefined + this.turnConfig = undefined for (const end of this.turnEnds) end() this.turnEnds.clear() this.drain() @@ -581,7 +585,11 @@ export abstract class SessionTurn< await this.flushNotes() // The stored settings of the thread, as they are when the turn starts. const stored = this.threadSettings - const removed = new Set(stored.plugins?.remove ?? []) + // A removed session plugin stays mounted, but gives this turn nothing: + // no tools, prompts, middleware, agents, extension items, or + // capabilities, also for the turn plugins below. + const removed = this.removedPlugins() + const session = this.threadPlugins() // The 'turn' plugins. A removed turn plugin is not set up. const perRun = this.plugins.filter( (plugin) => plugin.lifetime === 'turn' && !removed.has(plugin.name), @@ -593,17 +601,17 @@ export abstract class SessionTurn< registry: this.agentRegistry.fork(), harnessTools: [ ...(this.harness.tools ?? []), - ...(this.sessionPlugins?.tools ?? []), + ...(session?.tools ?? []), ], harnessProvides: (this.harness.middleware ?? []).flatMap( (middleware) => middleware.provides ?? [], ), - ...(this.sessionPlugins + ...(session ? { - inherited: this.sessionPlugins.values, - takenCommands: this.sessionPlugins.commands, - takenConfig: this.sessionPlugins.config, - inheritedExtensions: this.sessionPlugins.extensions, + inherited: session.values, + takenCommands: session.commands, + takenConfig: session.config, + inheritedExtensions: session.extensions, } : {}), // Bound to this turn's run, not read when a service is called. @@ -614,11 +622,12 @@ export abstract class SessionTurn< ) this.turnPlugins.add(heldPlugins) } - // A removed session plugin stays mounted, but gives this turn nothing. - const session = - this.sessionPlugins && removed.size > 0 - ? withoutPlugins(this.sessionPlugins, removed) - : this.sessionPlugins + // Once per turn, before the adapter pick, the tools, and the prompts. + await this.runTurnStarts( + [...(session?.turnStarts ?? []), ...(runPlugins?.turnStarts ?? [])], + operation, + info, + ) const bridges = [ session?.capabilityBridge, runPlugins?.capabilityBridge, @@ -710,7 +719,10 @@ export abstract class SessionTurn< ] const discovered = await this.discoverTools( [...(session?.discoverers ?? []), ...(runPlugins?.discoverers ?? [])], - new Set(staticTools.map((tool) => tool.name)), + toolOwners(staticTools, [ + ...(session?.owners.tools ?? []), + ...(runPlugins?.owners.tools ?? []), + ]), operation, info, ) @@ -726,6 +738,7 @@ export abstract class SessionTurn< turnToolNames, ), model: adapter.model, + run, }, operation, ) @@ -801,6 +814,15 @@ export abstract class SessionTurn< ...(sender ? { principal: sender } : {}), ...(turn.context !== undefined ? { context: turn.context } : {}), adapter, + // Its model calls are for this turn, and count under `router`. + model: { + chat: async ({ run: _run, ...call }) => + (await this.modelChat({ + ...call, + run, + source: 'router', + })) as never, + }, ...(activeAgent !== null ? { activeAgent } : {}), }) // Before the agent runs, so a takeover goes on with the new agent. @@ -828,6 +850,9 @@ export abstract class SessionTurn< /** The bag of the root agents, while a part of them runs. */ let rootBag = isRootTurn ? routedBag : undefined isRootPart = rootBag !== undefined + // A `session.tool()` turn takes no steer and no late join: it has no + // model to answer them. They wait and run as their own turn. + const isToolTurn = turn.tool !== undefined /** * The agent that answers the next part, when one agent does: `null` * for the main model. The routed message gets its name. @@ -937,9 +962,11 @@ export abstract class SessionTurn< 0, ephemeralCall(ephemeral), ) - chatMiddleware.unshift( - this.steering(ephemeral, { operationId: operation.id }), - ) + if (!isToolTurn) { + chatMiddleware.unshift( + this.steering(ephemeral, { operationId: operation.id }), + ) + } } const stream = chat({ adapter, @@ -969,7 +996,9 @@ export abstract class SessionTurn< ), // Before the harness middleware, so its change of each model // call (`onConfig`) gets the messages that join at that call. - ...(rootBag ? [] : [this.steering(ephemeral)]), + // The root agents of a routed part take steers in their own + // middleware (above). A tool turn takes none. + ...(rootBag || isToolTurn ? [] : [this.steering(ephemeral)]), // Before the harness middleware, so a compaction sees only the // context after a reset. resetCut, @@ -990,6 +1019,15 @@ export abstract class SessionTurn< ? [turnEphemeral(turn.ephemeral, start)] : []), ephemeralCall(ephemeral), + // The last system prompt: text that changes at each call. + contextPromptCall(() => + [ + ...(session?.contextPrompts ?? []), + ...(runPlugins?.contextPrompts ?? []), + ].map((prompt) => + this.redactor.scrub(resolvePrompt(prompt.text)), + ), + ), // Last, because a later middleware that returns `messages` (as // steering does) resets what the model gets. mediaMiddleware({ @@ -1275,15 +1313,18 @@ export abstract class SessionTurn< // an agent that makes no model call takes no steer. The main part // of a handoff takes no steer: it waits and runs as its own turn. // The main model claims them, so a cancel after this check cannot - // leave a run with no new input. - const hasLateJoin = isRootPart - ? (await this.joinSteers({ - operationId: operation.id, - messages: [], - ephemeral, - isRouted: true, - })) !== undefined - : !rootBag && (await this.claimJoins()) > 0 + // leave a run with no new input. A tool turn takes no late join + // (see `isToolTurn`). + const hasLateJoin = isToolTurn + ? false + : isRootPart + ? (await this.joinSteers({ + operationId: operation.id, + messages: [], + ephemeral, + isRouted: true, + })) !== undefined + : !rootBag && (await this.claimJoins()) > 0 if (!hasLateJoin) { const goesOn = await this.continueBeforeFinish( operation, diff --git a/packages/ai-harness/src/session/types.ts b/packages/ai-harness/src/session/types.ts index ba153a2b9f..134265d5d4 100644 --- a/packages/ai-harness/src/session/types.ts +++ b/packages/ai-harness/src/session/types.ts @@ -9,6 +9,7 @@ import type { UIMessage, } from '@tanstack/ai' import type { OperationImpl } from '../operation' +import type { ObservabilityOptions } from '../observability' import type { HeldPlugins } from './helpers' import type { RefreshFlights } from '../auth' import type { ToolTurn } from '../tool-turn' @@ -397,6 +398,8 @@ export interface SessionDependencies { openChild?: OpenChildSession /** Set on a child session only: what it takes from its parent. */ child?: ChildSessionLink + /** The `observability` option of the host: spans and `onOperation`. */ + observability?: ObservabilityOptions } /** What a child session of `harnessAgent` takes from its parent. */ diff --git a/packages/ai-harness/src/sse.ts b/packages/ai-harness/src/sse.ts new file mode 100644 index 0000000000..76e8f78542 --- /dev/null +++ b/packages/ai-harness/src/sse.ts @@ -0,0 +1,361 @@ +// The stream writers of the session tier: SSE and the socket. Each has a +// heartbeat, backpressure, and a close for a client that does not read. +import type { WebSocketLike } from '@tanstack/ai' +import type { HostFrame } from './protocol' + +/** The default of `HarnessStreamOptions.heartbeatMs`. */ +export const HEARTBEAT_MS = 20_000 +/** The default of `HarnessStreamOptions.slowConsumerMs`. */ +export const SLOW_CONSUMER_MS = 30_000 +/** The default of `HarnessStreamOptions.sseQueueBytes`. */ +export const SSE_QUEUE_BYTES = 1024 * 1024 +/** The default of `HarnessStreamOptions.socketBufferBytes`. */ +export const SOCKET_BUFFER_BYTES = 8 * 1024 * 1024 +/** How often a full socket is checked again. */ +export const SOCKET_RECHECK_MS = 50 +/** The close code of a socket that was too slow. */ +export const SLOW_CONSUMER_CLOSE_CODE = 4408 + +/** + * The values of the session-tier streams: the SSE routes and the socket. Each + * value must be a finite number above 0. A value that you do not set keeps + * its default. + */ +export interface HarnessStreamOptions { + /** + * After this many milliseconds, a stream sends a heartbeat, so proxies + * keep it open. Default: `20000` (20 s). + */ + heartbeatMs?: number + /** + * A client that reads nothing for this many milliseconds, while events + * wait, is closed with a `slow_consumer` error frame. Default: `30000` + * (30 s). + */ + slowConsumerMs?: number + /** + * The bytes that an SSE stream keeps ready for the client. When the queue + * is full, the stream stops reading the log. Default: `1048576` (1 MiB). + */ + sseQueueBytes?: number + /** + * The socket buffer (`bufferedAmount`) in bytes above which the socket + * stops reading the log and waits. Default: `8388608` (8 MiB). + */ + socketBufferBytes?: number +} + +/** `HarnessStreamOptions` with each default filled in. */ +export type StreamValues = Required + +/** The default stream values. */ +export const STREAM_DEFAULTS: StreamValues = { + heartbeatMs: HEARTBEAT_MS, + slowConsumerMs: SLOW_CONSUMER_MS, + sseQueueBytes: SSE_QUEUE_BYTES, + socketBufferBytes: SOCKET_BUFFER_BYTES, +} + +const STREAM_KEYS = [ + 'heartbeatMs', + 'slowConsumerMs', + 'sseQueueBytes', + 'socketBufferBytes', +] as const + +/** + * `stream` with each default filled in. Throws when a value is not a finite + * number above 0. `caller` names the function in the error. + */ +export function streamValues( + caller: string, + stream: HarnessStreamOptions = {}, +): StreamValues { + const values = { ...STREAM_DEFAULTS } + for (const key of STREAM_KEYS) { + const value: unknown = stream[key] + if (value === undefined) continue + if (typeof value !== 'number' || !Number.isFinite(value) || value <= 0) { + throw new Error( + `${caller} needs a stream.${key} that is a finite number above 0.`, + ) + } + values[key] = value + } + return values +} + +const encoder = new TextEncoder() +const PING = encoder.encode(': ping\n\n') + +/** One SSE event. `data` is JSON text, for example a `wireFrame`. */ +export interface SseFrame { + id?: string + data: string +} + +const errorFrame = ( + code: 'slow_consumer' | 'internal', + message: string, + requestId: string, +): HostFrame => ({ type: 'harness.error', message, code, requestId }) + +const errorMessage = (error: unknown) => + error instanceof Error ? error.message : String(error) + +/** + * A pull-based SSE body. It reads the next frame only when the client has + * read enough: the queue holds at most `stream.sseQueueBytes`. A slow client + * reads later, from the log. The stream also sends a `: ping` comment every + * `stream.heartbeatMs`, and it closes with a `slow_consumer` error frame when + * the queue stays full for `stream.slowConsumerMs`. When `frames` throws, the + * stream sends an `internal` error frame and closes. Aborting `signal` + * closes it. + * + * @example + * ```ts + * const body = sseStream({ + * frames: (signal) => toSse(session.events({ from, signal })), + * requestId, + * }) + * ``` + */ +export function sseStream({ + frames, + requestId, + signal, + stream = STREAM_DEFAULTS, +}: { + /** Gets the signal of the reader. The stream aborts it when it ends. */ + frames: (signal: AbortSignal) => AsyncIterable + requestId: string + signal?: AbortSignal + /** The stream values. See `streamValues`. */ + stream?: StreamValues +}): ReadableStream { + const reader = new AbortController() + let iterator: AsyncIterator | undefined + let heartbeat: ReturnType | undefined + let slow: ReturnType | undefined + let isDone = false + let onAbort: (() => void) | undefined + + const stop = () => { + isDone = true + if (onAbort) signal?.removeEventListener('abort', onAbort) + clearInterval(heartbeat) + clearTimeout(slow) + reader.abort() + void iterator?.return?.()?.catch(() => {}) + } + const end = ( + controller: ReadableStreamDefaultController, + last?: HostFrame, + ) => { + if (isDone) return + stop() + if (last) + controller.enqueue( + encoder.encode(sseText({ data: JSON.stringify(last) })), + ) + controller.close() + } + // A full queue means the client does not read. Give it time, then close. + const watchQueue = ( + controller: ReadableStreamDefaultController, + ) => { + if (slow !== undefined || (controller.desiredSize ?? 1) > 0) return + slow = setTimeout( + () => + end( + controller, + errorFrame( + 'slow_consumer', + 'The client reads too slowly.', + requestId, + ), + ), + stream.slowConsumerMs, + ) + } + + return new ReadableStream( + { + start(controller) { + iterator = frames(reader.signal)[Symbol.asyncIterator]() + heartbeat = setInterval(() => { + controller.enqueue(PING) + watchQueue(controller) + }, stream.heartbeatMs) + onAbort = () => end(controller) + if (signal?.aborted) return end(controller) + signal?.addEventListener('abort', onAbort, { once: true }) + }, + async pull(controller) { + // The client read, so it is not slow. + clearTimeout(slow) + slow = undefined + if (isDone || !iterator) return + try { + const next = await iterator.next() + if (isDone) return + if (next.done) return end(controller) + controller.enqueue(encoder.encode(sseText(next.value))) + watchQueue(controller) + } catch (error) { + end( + controller, + errorFrame('internal', errorMessage(error), requestId), + ) + } + }, + cancel: stop, + }, + new ByteLengthQueuingStrategy({ highWaterMark: stream.sseQueueBytes }), + ) +} + +function sseText({ id, data }: SseFrame) { + return `${id === undefined ? '' : `id: ${id}\n`}data: ${data}\n\n` +} + +/** + * `body` with a `: ping` comment every `heartbeatMs`, for the SSE bodies of + * the core (`POST run`, `GET run`). A comment goes in only between two + * events, never inside one. SSE parsers skip comments. + */ +export function withHeartbeat( + body: ReadableStream, + { heartbeatMs = HEARTBEAT_MS }: { heartbeatMs?: number } = {}, +): ReadableStream { + const source = body.getReader() + let heartbeat: ReturnType | undefined + // True when the bytes sent so far end with a blank line: an event ended. + let isBetweenEvents = true + return new ReadableStream({ + start(controller) { + heartbeat = setInterval(() => { + if (isBetweenEvents) controller.enqueue(PING) + }, heartbeatMs) + }, + async pull(controller) { + try { + const { done, value } = await source.read() + if (done) { + clearInterval(heartbeat) + return controller.close() + } + const size = value.byteLength + if (size > 0) + isBetweenEvents = + size >= 2 && value[size - 1] === 10 && value[size - 2] === 10 + controller.enqueue(value) + } catch (error) { + clearInterval(heartbeat) + controller.error(error) + } + }, + cancel(reason) { + clearInterval(heartbeat) + return source.cancel(reason) + }, + }) +} + +/** + * A sender for a session-tier socket. Frames go out in order. Before each + * frame it reads `socket.bufferedAmount`: above `stream.socketBufferBytes` + * it waits, and checks again every `SOCKET_RECHECK_MS`. So the caller, which + * awaits each send, stops reading the log. After `stream.slowConsumerMs` it + * sends a `slow_consumer` error frame and closes the socket with code 4408. + * It also sends a `harness.ping` frame every `stream.heartbeatMs`. Its timers + * stop when the socket closes or fails, or when `signal` aborts. + * + * A socket without `bufferedAmount` gets no backpressure. + * + * @example + * ```ts + * const send = socketSender({ socket, requestId, signal }) + * for await (const entry of events) if (!(await send(wireFrame(entry)))) break + * ``` + */ +export function socketSender({ + socket, + requestId, + signal, + stream = STREAM_DEFAULTS, +}: { + socket: WebSocketLike + /** Named in the `slow_consumer` error frame. */ + requestId: string + signal?: AbortSignal + /** The stream values. See `streamValues`. */ + stream?: StreamValues +}): (frame: HostFrame | string) => Promise { + let isClosed = false + let recheck: ReturnType | undefined + let wake: (() => void) | undefined + let queue = Promise.resolve(true) + const isFull = () => (socket.bufferedAmount ?? 0) > stream.socketBufferBytes + + const stop = () => { + isClosed = true + clearInterval(ping) + clearTimeout(recheck) + wake?.() + } + const ping = setInterval(() => { + if (isFull()) return + try { + socket.send(JSON.stringify({ type: 'harness.ping' })) + } catch { + stop() + } + }, stream.heartbeatMs) + socket.addEventListener('close', stop) + socket.addEventListener('error', stop) + if (signal?.aborted) stop() + signal?.addEventListener('abort', stop, { once: true }) + + // Resolves true when the buffer has room, false when the socket closed. + const room = () => + new Promise((resolve) => { + let waited = 0 + wake = () => resolve(false) + const check = () => { + if (isClosed) return resolve(false) + if (!isFull()) return resolve(true) + if (waited >= stream.slowConsumerMs) { + const frame = errorFrame( + 'slow_consumer', + 'The client reads too slowly.', + requestId, + ) + try { + socket.send(JSON.stringify(frame)) + socket.close(SLOW_CONSUMER_CLOSE_CODE, 'slow_consumer') + } catch { + // The socket is gone already. + } + stop() + return resolve(false) + } + waited += SOCKET_RECHECK_MS + recheck = setTimeout(check, SOCKET_RECHECK_MS) + } + check() + }) + + const sendNow = async (frame: HostFrame | string) => { + if (isClosed) return false + if (isFull() && !(await room())) return false + try { + socket.send(typeof frame === 'string' ? frame : JSON.stringify(frame)) + return true + } catch { + stop() + return false + } + } + return (frame) => (queue = queue.then(() => sendNow(frame))) +} diff --git a/packages/ai-harness/src/types.ts b/packages/ai-harness/src/types.ts index c9213308ab..fb3b4492d0 100644 --- a/packages/ai-harness/src/types.ts +++ b/packages/ai-harness/src/types.ts @@ -15,6 +15,7 @@ import type { LogRecord } from '@tanstack/ai-persistence' import type { CredentialsAccess } from './auth' import type { AnyHarness } from './define' import type { HandoffOptions } from './handoff' +import type { HandoffRecord } from './log-records' import type { UsageCounts } from './usage' /** A user message: plain text, or content parts (text, images, files). */ @@ -80,8 +81,11 @@ export type SendOptions = { ephemeral?: ReadonlyArray /** System prompts ahead of the harness prompts, from server code only. */ systemPreamble?: Array - /** @internal The AG-UI `runId` of a `POST run` request. The turn runs as it. */ - runId?: string + /** + * A W3C `traceparent` header value of the caller. The span of the turn gets + * a link to it. A value that is not a valid `traceparent` is ignored. + */ + traceparent?: string } & ( | { mode?: 'prompt' @@ -99,6 +103,13 @@ export type SendOptions = { | { mode: 'steer' | 'followUp'; busy?: never; as?: never } ) +/** + * @internal `SendOptions` with the AG-UI `runId` of a `POST run` request: the + * turn runs as it. The HTTP handler and an offered prompt pass it to + * `sendForRun`. The package index does not export it. + */ +export type SendForRunOptions = SendOptions & { runId?: string } + /** Options of `session.send` for a message to an agent run. See `send`. */ export interface AgentSendOptions { /** The operation id of the agent run. Any run of its chain works. */ @@ -136,7 +147,12 @@ export type Cursor = string /** One event in a session's ordered stream. */ export interface SessionEvent { cursor: Cursor - /** The operation that produced the event. */ + /** + * The thread of the event: the thread of the session, or a child thread + * of its tree. + */ + threadId: string + /** The operation that produced the event. It is also the AG-UI `runId`. */ operationId: string event: StreamChunk } @@ -274,16 +290,7 @@ export type HarnessInput = ( * `pendingQuestions` (a `step.ask` wait, a live MCP question), or a * `human` wait. */ - | { - op: 'answer' - questionId: string - value: unknown - /** - * @internal A secret answer that another host saved in the credential - * store, for the owner of the log. The owner reads it once. - */ - ref?: string - } + | { op: 'answer'; questionId: string; value: unknown } | { op: 'config'; key: string; value: unknown } /** Run one tool in `expose.tools` with no model. See `session.tool`. */ | { op: 'tool'; name: string; args?: unknown; meta?: Record } @@ -305,6 +312,16 @@ export type HarnessInput = ( } ) & { inputId?: string } +/** + * @internal An input that a host offers to the owner of the log. An answer + * can carry `ref`: a secret answer that another host saved in the + * credential store, for the owner of the log. The owner reads it once. The + * package index does not export it. + */ +export type OfferedHarnessInput = + | Exclude + | (Extract & { ref?: string }) + /** * An input that waits to run. `session.inputs()` lists them. Change one with * the `cancelInput` and `setDelivery` inputs. @@ -654,7 +671,10 @@ export interface Operation extends PromiseLike { * A resolved receipt means that the input survives a restart. */ readonly receipt: Promise - /** This operation's events, from `from` (exclusive) onward. */ + /** + * This operation's events, from `from` (exclusive) onward. Default: from + * the start of the operation. + */ events: (options?: { from?: Cursor signal?: AbortSignal @@ -682,11 +702,17 @@ export interface AgentRun extends Operation { * runs as a follow-up. `mode: 'followUp'` runs the agent again after the * run ends, on the run's transcript with the message added, with its * first input and options. A run that ended starts again at once. The - * receipt names the run that gets the message. + * receipt names the run that gets the message. `principal` is who sends + * it. Default: the user of the run. On a handle of `host.open`, the + * principal of the handle, after its access check. */ send: ( message: UserInput, - options?: { mode?: 'steer' | 'followUp'; inputId?: string }, + options?: { + mode?: 'steer' | 'followUp' + inputId?: string + principal?: Principal + }, ) => Promise } @@ -783,6 +809,13 @@ export interface UsageUpdate { model: string /** The principal id of the sender, when known. */ sender?: string + /** + * `plugin:` for a call of `ctx.model.chat`, and `router` for a call + * of `model.chat` in `routing.router`. Turns and agent runs have none. + */ + source?: string + /** The `label` of a `ctx.model.chat` call, when it has one. */ + label?: string /** The usage of this call. */ usage: Omit /** The new totals of the thread. */ @@ -853,6 +886,26 @@ export const HARNESS_EVENTS = { * known), the `usage` of the call, and the new thread `total`. */ usage: 'harness.usage', + /** The owner stored an input that waits behind a running turn. */ + inputQueued: 'harness.input.queued', + /** A host became the owner of the log, at its first append. */ + ownershipChanged: 'harness.ownership.changed', + /** + * The workspace that the log binds the thread to is gone, so the + * workspace tools fail. The value has the `operationId` of the run that + * tried it, and the `hostId` of a folder on a host. + */ + stranded: 'harness.stranded', + /** The conversation moved to another thread, or came from one. */ + transferred: 'harness.transferred', + /** + * A wait opened in the log: a question, a `step.ask`, a handoff to a + * person, or a wait of a background agent run. The interrupts of a turn + * come in its `RUN_FINISHED` instead. + */ + waitOpened: 'harness.wait.opened', + /** A wait ended. */ + waitSettled: 'harness.wait.settled', /** * The events after the reader's cursor are gone (a truncate dropped them), * so the stream skips to the first event it has. The value is an @@ -874,3 +927,96 @@ export interface EventGap { /** `'truncated'`: a truncate of the log dropped the events. */ reason: 'truncated' | 'evicted' } + +/** + * The `value` of each `harness.*` `CUSTOM` event, by event name. A client in + * another language can use it as the event catalog. + * + * @example + * ```ts + * function onGap(gap: HarnessEventValues['harness.gap']) { + * console.log(`Events lost. The stream goes on from ${gap.resumeAt}.`) + * } + * ``` + */ +export interface HarnessEventValues { + 'harness.operation.started': { + operationId: string + kind: OperationKind + agent?: string + } + 'harness.operation.finished': { + operationId: string + status: OperationStatus + } + 'harness.operation.resumed': { operationId: string; resumedFrom?: string } + 'harness.config.changed': { key: string; value: unknown } + 'harness.settings.changed': { settings: ThreadSettings } + 'harness.commands.changed': Record + 'harness.reloaded': { error?: string } + 'harness.question': PendingQuestion + 'harness.question.answered': { questionId: string } + 'harness.plugin.event': { plugin: string; name: string; value: unknown } + 'harness.auth_required': AuthRequest + 'harness.input.accepted': { inputId: string; op: HarnessInput['op'] } + 'harness.input.applied': { inputId: string; operationId: string } + /** + * `offered` is true when the owner refused an offer from another host, and + * the log has another input with the id. That input is not refused. + */ + 'harness.input.rejected': { + inputId: string + reason: string + offered?: true + } + 'harness.input.settled': InputSettlement + 'harness.input.delivery': { + inputId: string + delivery: WaitingInput['delivery'] + } + 'harness.media': MediaRecord + 'harness.turn.retry': { + operationId: string + retries: number + error: { message: string; code?: string; retryAfterMs?: number } + continued: boolean + } + 'harness.reset': { inputId: string; note?: string } + 'harness.revert': { messageId: string | null } + 'harness.usage': UsageUpdate + 'harness.input.queued': { + inputId: string + delivery: WaitingInput['delivery'] + } + 'harness.ownership.changed': { + /** The ownership epoch of the new owner. */ + epoch: number + /** + * `taken_over`: the ownership of the old owner expired with no release, + * so it stopped, and this host took the log. + */ + reason: 'acquired' | 'taken_over' + } + 'harness.stranded': { operationId?: string; hostId?: string } + 'harness.handoff': HandoffRecord + 'harness.transferred': { + transferId: string + /** On the target thread: the thread it came from. */ + from?: { threadId: string } + /** On the source thread: the thread it went to. */ + to?: { threadId: string } + } + 'harness.wait.opened': { + waitId: string + kind: 'approval' | 'question' | 'sign-in' | 'client-tool' | 'human' + runId: string + message?: string + /** Epoch milliseconds. */ + expiresAt?: number + } + 'harness.wait.settled': { + waitId: string + outcome: 'answered' | 'cancelled' | 'expired' + } + 'harness.gap': EventGap +} diff --git a/packages/ai-harness/src/usage.ts b/packages/ai-harness/src/usage.ts index e12b23f4d8..8430801a76 100644 --- a/packages/ai-harness/src/usage.ts +++ b/packages/ai-harness/src/usage.ts @@ -19,9 +19,11 @@ export interface UsageCounts { } /** - * The usage of a thread: every model call of its turns, their subagents, and - * its agent runs. `bySender` has the calls of a known sender only, keyed by - * the principal id, so its sum can be less than `total`. + * The usage of a thread: every model call of its turns, their subagents, its + * agent runs, the model calls of plugins (`ctx.model.chat`), of the router, + * and of a transfer summary. `bySender` has the calls of a known sender + * only, keyed by the principal id, and `bySource` has only the calls that + * have a source, so their sums can be less than `total`. */ export interface SessionUsage { total: UsageCounts @@ -29,6 +31,12 @@ export interface SessionUsage { byModel: Record /** By principal id. */ bySender: Record + /** + * By the source of a call: `plugin:` (for example + * `plugin:tanstack/title`), `router`, or `transfer`. Turns and agent runs + * have no source, so they are not in it. + */ + bySource: Record } /** One model call, as the log keeps it. A type, so it is a log record field set. */ @@ -40,6 +48,14 @@ export type UsageCall = { * counts for each run above it in the run tree. */ runId?: string + /** + * Who made a call that is not a turn or an agent run: `plugin:` + * for `ctx.model.chat`, `router` for `model.chat` of `routing.router`, + * and `transfer` for the summary of `host.transfer`. + */ + source?: string + /** The `label` of the call, for example `judge`, or `summary` for a transfer. */ + label?: string usage: Omit } @@ -56,6 +72,7 @@ export const emptyUsage = (): SessionUsage => ({ total: emptyCounts(), byModel: {}, bySender: {}, + bySource: {}, }) /** The counts that one call adds, from what the provider reported. */ @@ -104,10 +121,15 @@ export function addUsage(totals: SessionUsage, call: UsageCall) { if (call.principal) { addCounts(countsOf(totals.bySender, call.principal.id), call.usage) } + if (call.source !== undefined) { + addCounts(countsOf(totals.bySource, call.source), call.usage) + } } -/** Is `value` usage totals that the harness wrote? */ -export function isSessionUsage(value: unknown): value is SessionUsage { +/** Is `value` usage totals that the harness wrote? Old totals have no `bySource`. */ +export function isSessionUsage( + value: unknown, +): value is Omit & Partial { return ( typeof value === 'object' && value !== null && diff --git a/packages/ai-harness/src/view/index.ts b/packages/ai-harness/src/view/index.ts index ef01926446..fab18c1f59 100644 --- a/packages/ai-harness/src/view/index.ts +++ b/packages/ai-harness/src/view/index.ts @@ -1,5 +1,6 @@ // Browser-safe: runtime imports only from @tanstack/ai, @tanstack/store, -// ../types, ../media-ref, and ./reduce. Everything else is a type. +// ../types, ../media-ref, and ./reduce (it also reads ../usage). Everything +// else is a type. import { EventType, INTERRUPT_PAYLOAD_METADATA_KEY, @@ -50,7 +51,6 @@ import type { } from '../types' import type { ItemFactory } from './reduce' import type { - AgentPart, Approval, ApproveOptions, ClientToolCall, @@ -149,7 +149,7 @@ export type ViewSendOptions = { to?: undefined } & ( /** The options of `view.send`. */ export interface ViewTextOptions { - /** Stored media files (from `upload`) that go with the text. */ + /** Stored media files (from `putMedia`) that go with the text. */ attachments?: ReadonlyArray /** * `'operator'`: the reply of an operator while a handoff to a person is @@ -184,14 +184,14 @@ export interface SessionViewSource { cancel: (options?: CancelOptions) => Promise // Both sources have these. Without one, its view action shows an error. continue?: () => unknown - cancelInput?: (inputId: string) => Promise + cancelInput?: (options: { inputId: string }) => Promise setDelivery?: (change: { inputId: string delivery: WaitingInput['delivery'] }) => Promise configure?: (options: ConfigureOptions) => Promise reset?: (options?: ResetOptions) => Promise - revert?: (messageId: string) => Promise + revert?: (options: { messageId: string }) => Promise unrevert?: () => Promise answer: (answer: { questionId: string; value: unknown }) => Promise command: (options: CommandOptions) => unknown @@ -205,14 +205,17 @@ export interface SessionViewSource { transcript: () => Promise> describe: () => SessionDescription | Promise /** A URL for a media file. Without it, media parts have no `url`. */ - mediaUrl?: (id: string) => Promise<{ url?: string; expiresAt?: number }> + mediaUrl?: (options: { + id: string + }) => Promise<{ url?: string; expiresAt?: number }> /** The bytes of a media file. Without it, `load()` rejects. */ loadMedia?: (options: { id: string }) => Promise - /** Store a file: `upload` on a client, `putMedia` on a session. */ - upload?: (file: MediaUpload) => Promise + /** Store a file. Without it, `view.putMedia` rejects. */ putMedia?: (file: MediaUpload) => Promise /** The saved messages of a child agent run. Without it, a reload shows no child parts. */ - agentTranscript?: (subagentRunId: string) => Promise> + agentTranscript?: (options: { + subagentRunId: string + }) => Promise> } /** The bytes of a file to store. */ @@ -253,7 +256,7 @@ export interface SessionView { ready: Promise /** * A `/command`, a prompt, or a steer while a turn runs. `attachments` are - * stored media files (from `upload`) that go with the text. Pass them as + * stored media files (from `putMedia`) that go with the text. Pass them as * a list, or as `{ attachments }`. A `/command` ignores them. Empty text * with no attachments does nothing. `{ as: 'operator' }` sends the reply * of an operator (see {@link ViewTextOptions}). @@ -279,10 +282,10 @@ export interface SessionView { * * @example * ```ts - * const record = await view.upload({ body: bytes, mimeType: 'image/png', name: 'cat.png' }) + * const record = await view.putMedia({ body: bytes, mimeType: 'image/png', name: 'cat.png' }) * ``` */ - upload: (file: MediaUpload) => Promise + putMedia: (file: MediaUpload) => Promise /** * Run the plugin command `name` with its `input`. * @@ -305,7 +308,7 @@ export interface SessionView { /** Stop the running turn, or the operation `operationId`. */ cancel: (options?: CancelOptions) => Promise /** Drop an input from `waitingInputs` before it runs. */ - cancelInput: (inputId: string) => Promise + cancelInput: (options: { inputId: string }) => Promise /** Move a waiting input: `'steer'` joins the running turn, `'queue'` waits for its own turn. */ setDelivery: (change: { inputId: string @@ -313,15 +316,59 @@ export interface SessionView { }) => Promise /** Start a fresh model context. The model sees `note` first. */ reset: (options?: ResetOptions) => Promise - /** Go back to the message `messageId`. `unrevert` undoes it. */ - revert: (messageId: string) => Promise + /** + * Go back to the message `messageId`. `unrevert` undoes it. + * + * @example + * ```ts + * await view.revert({ messageId: message.id }) + * ``` + */ + revert: (options: { messageId: string }) => Promise unrevert: () => Promise - approve: (id: string, options?: ApproveOptions) => void - reject: (id: string, payload?: unknown) => void + /** + * Say yes to the approval `id`, maybe with changed arguments or data. + * + * @example + * ```ts + * view.approve({ id: approval.id }) + * ``` + */ + approve: (options: { id: string } & ApproveOptions) => void + /** + * Say no to the approval `id`. `payload` goes to the tool's approval + * schema, when it has one. + * + * @example + * ```ts + * view.reject({ id: approval.id }) + * ``` + */ + reject: (options: { id: string; payload?: unknown }) => void approveAll: () => void rejectAll: () => void - /** Add a line of your own to the messages. */ - notice: (text: string) => void + /** + * Add a line of your own to the messages. + * + * @example + * ```ts + * view.notice({ text: 'Saved.' }) + * ``` + */ + notice: (options: { text: string }) => void + /** + * Fill the agent part `id` of a saved message with what that child wrote. + * After a reload, a saved child shows only its card (name, status, result) + * until you call this, for example when the user opens it. Its nested + * children load the same way. Without `source.agentTranscript`, it does + * nothing. + * + * @example + * ```ts + * await view.loadAgent({ id: part.id }) + * ``` + */ + loadAgent: (options: { id: string }) => Promise /** * Listen for a view event (`'approval'`, `'turnEnd'`, ...) or a plugin event. * Returns a function that stops listening. @@ -343,10 +390,19 @@ function userInput(line: string, attachments: ReadonlyArray) { return input } +/** The shown form of a message, and the media parts it shows. */ +interface ShownMessage { + message: ViewMessage + media: Array +} + const isGap = ({ event }: SessionEvent) => event.type === EventType.CUSTOM && event.name === HARNESS_EVENTS.gap -/** Events after which the snapshot can differ. */ +/** + * Events after which the snapshot can differ. A config change is not in the + * snapshot. A usage event changes only the usage, and the reducer adds it. + */ function changesSnapshot(entry: SessionEvent) { const { event } = entry if ( @@ -358,10 +414,29 @@ function changesSnapshot(entry: SessionEvent) { } return ( event.type === EventType.CUSTOM && - /^harness\.(operation|input|question|config|usage|gap)/.test(event.name) + /^harness\.(operation|input|question|gap)/.test(event.name) ) } +/** + * A usage event that the reducer could not add, and that the view does not + * count yet: a call was missed, so the view reads the usage again. + */ +function missesUsage({ event }: SessionEvent, state: SessionViewState) { + if (event.type !== EventType.CUSTOM || event.name !== HARNESS_EVENTS.usage) + return false + const { value } = event + const total: unknown = + typeof value === 'object' && value !== null && 'total' in value + ? value.total + : undefined + const calls = + typeof total === 'object' && total !== null && 'calls' in total + ? total.calls + : undefined + return typeof calls !== 'number' || calls > (state.usage?.total.calls ?? 0) +} + /** * A receipt that `applyInput` refused before the session saw the input. The * session reports its own refusals with a `harness.input.rejected` event, so @@ -440,8 +515,14 @@ export function createSessionView(source: SessionViewSource) { const agentNames = new Map() const mediaParts = new Map() const urlRefreshes = new Map>() - /** The media ids the last commit shows. */ - let shownMedia = new Set() + /** + * The shown form of each message, and the media parts it shows, by message + * object. The reducer keeps a message that did not change, so a commit + * walks only the new messages. + */ + const shownOf = new WeakMap() + /** The agent parts of saved messages, that `loadAgent` can fill. */ + let savedAgents = new Set() /** The agent parts of a reload that loaded, or load now. */ const loadedAgents = new Set() /** The ids the first transcript has. See `replayed`. */ @@ -457,9 +538,16 @@ export function createSessionView(source: SessionViewSource) { const get = () => internal.state const commit = (next: SessionViewState) => { if (disposed) return - shownMedia = new Set() - const shown = withMediaActions(next) + const entries = next.messages.map(withActions) + const messages = keep( + next.messages, + entries.map((entry) => entry.message), + ) + const shown = messages === next.messages ? next : { ...next, messages } // A media file that no message shows now (after a revert) needs no URL. + const shownMedia = new Set( + entries.flatMap((entry) => entry.media.map((part) => part.id)), + ) for (const id of mediaParts.keys()) { if (shownMedia.has(id)) continue clearTimeout(urlRefreshes.get(id)) @@ -519,7 +607,7 @@ export function createSessionView(source: SessionViewSource) { const fetchUrl = async (id: string) => { if (!source.mediaUrl) return try { - const { url, expiresAt } = await source.mediaUrl(id) + const { url, expiresAt } = await source.mediaUrl({ id }) const part = mediaParts.get(id) if (disposed || !part) return if (url !== undefined) mediaParts.set(id, { ...part, url }) @@ -539,7 +627,6 @@ export function createSessionView(source: SessionViewSource) { } /** The view's own part for a media id: it loads from the source and gets a url once. */ const mediaActions = (part: MediaPart) => { - shownMedia.add(part.id) const known = mediaParts.get(part.id) if (known) return known const { id } = part @@ -554,33 +641,58 @@ export function createSessionView(source: SessionViewSource) { void fetchUrl(id) return shown } - const partsWithActions = (parts: Array): Array => + /** The parts with the view's media parts. Each media part goes in `found`. */ + const partsWithActions = ( + parts: Array, + found: Array, + ): Array => keep( parts, parts.map((part) => { - if (part.type === 'media') return mediaActions(part) + if (part.type === 'media') { + const shown = mediaActions(part) + found.push(shown) + return shown + } if (part.type !== 'agent') return part - const inner = partsWithActions(part.parts) + const inner = partsWithActions(part.parts, found) return inner === part.parts ? part : { ...part, parts: inner } }), ) - const messageWithActions = (message: ViewMessage) => { + const messageWithActions = ( + message: ViewMessage, + found: Array, + ) => { if (message.role === 'assistant') { - const parts = partsWithActions(message.parts) + const parts = partsWithActions(message.parts, found) return parts === message.parts ? message : { ...message, parts } } if (message.role !== 'user' || !message.media) return message - const media = keep(message.media, message.media.map(mediaActions)) + const media = keep( + message.media, + message.media.map((part) => { + const shown = mediaActions(part) + found.push(shown) + return shown + }), + ) return media === message.media ? message : { ...message, media } } - // ponytail: this walks every message on each commit. Track new media ids - // instead if a long session gets slow. - const withMediaActions = (state: SessionViewState) => { - const messages = keep( - state.messages, - state.messages.map(messageWithActions), + /** + * The shown form of `message`. A known message is walked again only when + * one of its media parts changed (a new URL) or went (after a revert). + */ + const withActions = (message: ViewMessage): ShownMessage => { + const known = shownOf.get(message) + const isCurrent = known?.media.every( + (part) => mediaParts.get(part.id) === part, ) - return messages === state.messages ? state : { ...state, messages } + if (known && isCurrent) return known + const media: Array = [] + const entry = { message: messageWithActions(message, media), media } + shownOf.set(message, entry) + shownOf.set(entry.message, entry) + return entry } /** `items` without the items that `drop` matches. */ @@ -677,7 +789,8 @@ export function createSessionView(source: SessionViewSource) { args: call?.args ?? bound?.originalArgs, ...(interrupt.message ? { message: interrupt.message } : {}), approve: (options) => decide(interrupt.id, approvedWith(options)), - reject: (payload) => decide(interrupt.id, rejectedWith(payload)), + reject: ({ payload } = {}) => + decide(interrupt.id, rejectedWith(payload)), } }, interrupt: (interrupt) => { @@ -799,31 +912,32 @@ export function createSessionView(source: SessionViewSource) { } }) + /** The ids of the agent parts in `messages` that show no parts yet. */ + const agentIdsIn = (messages: ReadonlyArray) => + emptyAgentsIn(messages).map((agent) => agent.id) + /** - * Fill the agent parts `ids` that the saved messages made, from the saved - * messages of each child, then the nested children those hold. + * Fill the saved agent part `id` from the saved messages of its child. Its + * nested children become saved agent parts too, and load on their own call. */ - const loadAgents = async (agents: Array): Promise => { - if (!source.agentTranscript) return - // ponytail: every child loads at start. Load on demand if long sessions - // get slow. - await Promise.all( - agents.map(async ({ id }) => { - // A continued child has one id for each of its calls. - if (loadedAgents.has(id)) return - loadedAgents.add(id) - try { - const messages = (await source.agentTranscript?.(id)) ?? [] - const parts = agentPartsFrom(messages) - if (disposed) return - // A child that failed before its first message has only its card. - commit(withAgentParts(get(), id, parts, endOf(messages))) - await loadAgents(emptyAgentsIn([{ id, role: 'assistant', parts }])) - } catch (error) { - fail(error) - } - }), - ) + const loadAgent = async (id: string): Promise => { + // A continued child has one id for each of its calls. + const isWaiting = savedAgents.has(id) && !loadedAgents.has(id) + if (!source.agentTranscript || !isWaiting) return + loadedAgents.add(id) + try { + const messages = await source.agentTranscript({ subagentRunId: id }) + const parts = agentPartsFrom(messages) + if (disposed) return + // A child that failed before its first message has only its card. + commit(withAgentParts(get(), id, parts, endOf(messages))) + const inner = agentIdsIn([{ id, role: 'assistant', parts }]) + for (const child of inner) savedAgents.add(child) + } catch (error) { + // A later call can try again. + loadedAgents.delete(id) + fail(error) + } } // A revert changes which saved messages show. The session is idle then. @@ -833,13 +947,49 @@ export function createSessionView(source: SessionViewSource) { if (disposed) return const messages = messagesFromTranscript(transcript) commit({ ...get(), messages }) + // The children that loaded before load again, a parent before its + // nested children. + const loaded = [...loadedAgents] loadedAgents.clear() - await loadAgents(emptyAgentsIn(messages)) + savedAgents = new Set(agentIdsIn(messages)) + for (const id of loaded) await loadAgent(id) } catch (error) { fail(error) } }) + /** + * After a gap: the events of a part of the log are gone. Load the snapshot + * and the transcript again, as `ready` does, while `read` waits. The + * stream goes on after the gap's `resumeAt`. Its events up to the new + * snapshot cursor replay only for the operations that run now, and the + * parts that the transcript has are dropped by id. + */ + const reload = async () => { + try { + // The snapshot before the transcript, as in `ready`. + const snapshot = await source.snapshot() + const transcript = await source.transcript() + if (disposed) return + commit({ ...get(), messages: messagesFromTranscript(transcript) }) + takeSnapshot(snapshot) + saved = transcriptIds(transcript) + held = [] + replayEnd = Number(snapshot.cursor) + replayOps = new Set(snapshot.activeOperations.map((item) => item.id)) + // The children that loaded before load again, a parent before its + // nested children. The events do not wait for them. + const loaded = [...loadedAgents] + loadedAgents.clear() + savedAgents = new Set(agentIdsIn(get().messages)) + void (async () => { + for (const id of loaded) await loadAgent(id) + })() + } catch (error) { + fail(error) + } + } + const fire = ( entry: SessionEvent, before: SessionViewState, @@ -897,12 +1047,8 @@ export function createSessionView(source: SessionViewSource) { event.name === HARNESS_EVENTS.commandsChanged || event.name === HARNESS_EVENTS.reloaded if (isDescriptionChange) refreshDescription() - // A gap: the events of a part of the log are gone, so the view loads - // the transcript (and the snapshot, see `changesSnapshot`) again. - const isReload = - event.name === HARNESS_EVENTS.revert || - event.name === HARNESS_EVENTS.gap - if (isReload) refreshTranscript() + // A gap loads the transcript again too, in `read` (see `reload`). + if (event.name === HARNESS_EVENTS.revert) refreshTranscript() } for (const signIn of after.signIns) if (!before.signIns.includes(signIn)) emit('signIn', signIn) @@ -917,9 +1063,12 @@ export function createSessionView(source: SessionViewSource) { }) for await (const entry of entries) { if (disposed) return - // A gap is not an event of an operation, so no replay drops it. - if (isGap(entry)) apply(entry) - else for (const ready of replayed(entry)) apply(ready) + // A gap is not an event of an operation, so no replay drops it. No + // event applies until the reload ends. + if (isGap(entry)) { + apply(entry) + await reload() + } else for (const ready of replayed(entry)) apply(ready) } } catch (error) { fail(error) @@ -933,7 +1082,7 @@ export function createSessionView(source: SessionViewSource) { const after = applyEvent(before, entry) commit(after) fire(entry, before, after) - if (changesSnapshot(entry)) refreshSnapshot() + if (changesSnapshot(entry) || missesUsage(entry, after)) refreshSnapshot() } /** @@ -1052,7 +1201,7 @@ export function createSessionView(source: SessionViewSource) { // ponytail: a background child that runs now can save a model call // before its saved messages load, and that call then shows twice. Drop // live parts by id if that shows up. - void loadAgents(emptyAgentsIn(get().messages)) + savedAgents = new Set(agentIdsIn(get().messages)) void read(snapshot ? String(from) : undefined) })() @@ -1117,9 +1266,8 @@ export function createSessionView(source: SessionViewSource) { alive() settle(source.continue ? source.continue() : missing('continue')) }, - upload: (file) => { + putMedia: (file) => { alive() - if (source.upload) return source.upload(file) if (source.putMedia) return source.putMedia(file) return Promise.reject(new Error('This view source cannot store media.')) }, @@ -1133,10 +1281,10 @@ export function createSessionView(source: SessionViewSource) { source.configure ? source.configure(options) : missing('configure'), ), cancel: (options) => act(() => source.cancel(options)), - cancelInput: (inputId) => + cancelInput: ({ inputId }) => act(() => source.cancelInput - ? source.cancelInput(inputId) + ? source.cancelInput({ inputId }) : missing('cancel inputs'), ), setDelivery: (change) => @@ -1147,18 +1295,29 @@ export function createSessionView(source: SessionViewSource) { ), reset: (options) => act(() => (source.reset ? source.reset(options) : missing('reset'))), - revert: (messageId) => - act(() => (source.revert ? source.revert(messageId) : missing('revert'))), + revert: ({ messageId }) => + act(() => + source.revert ? source.revert({ messageId }) : missing('revert'), + ), unrevert: () => act(() => (source.unrevert ? source.unrevert() : missing('unrevert'))), - approve: (id, options) => decide(id, approvedWith(options)), - reject: (id, payload) => decide(id, rejectedWith(payload)), + approve: ({ id, ...options }) => + decide( + id, + approvedWith(Object.keys(options).length > 0 ? options : undefined), + ), + reject: ({ id, payload }) => decide(id, rejectedWith(payload)), approveAll: () => decideAll(true), rejectAll: () => decideAll(false), - notice: (text) => { + notice: ({ text }) => { alive() commit(withNotice(get(), 'ui', text)) }, + loadAgent: async ({ id }) => { + alive() + await ready + await loadAgent(id) + }, on, dispose: () => { if (disposed) return diff --git a/packages/ai-harness/src/view/reduce.ts b/packages/ai-harness/src/view/reduce.ts index e8948df4d1..8e5361761e 100644 --- a/packages/ai-harness/src/view/reduce.ts +++ b/packages/ai-harness/src/view/reduce.ts @@ -2,9 +2,11 @@ import { EventType } from '@tanstack/ai' import { isMediaRecord, mediaIdOf, mediaOfMessage } from '../media-ref' import { signInOf, waitKind } from '../waits' import { HARNESS_EVENTS } from '../types' +import { addUsage } from '../usage' import type { Interrupt, ModelMessage, StreamChunk } from '@tanstack/ai' import type { SessionDescription, SessionSnapshot } from '../session/index' import type { SessionEvent } from '../types' +import type { UsageCall } from '../usage' import type { Wait } from '../waits' import type { AgentPart, @@ -206,7 +208,10 @@ function appendText( if (delta === '') return parts const last = parts.at(-1) if (last?.type === type) { - return [...parts.slice(0, -1), { type, text: last.text + delta }] + // One copy of the parts: the last part changes, the others stay. + const next = [...parts] + next[next.length - 1] = { type, text: last.text + delta } + return next } return [...parts, { type, text: delta }] } @@ -530,6 +535,35 @@ function failRunning( ) } +/** + * Add the call of a `harness.usage` event to the usage. Only the next call + * adds: a replayed call is in the usage already. The view reads the usage + * again when a call was missed. + */ +function withUsage( + state: SessionViewState, + value: Record, +): SessionViewState { + const { usage } = state + const total = recordOf(value.total) + const isNext = usage !== undefined && total.calls === usage.total.calls + 1 + if (!isNext || typeof value.model !== 'string' || !isRecord(value.usage)) + return state + const next = structuredClone(usage) + addUsage(next, { + model: value.model, + ...(typeof value.sender === 'string' + ? { principal: { id: value.sender } } + : {}), + // A plugin, router, or transfer call names its source; keep it so + // `bySource` stays correct. + ...(typeof value.source === 'string' ? { source: value.source } : {}), + // The harness sends the counts of `UsageCall['usage']`. + usage: value.usage as UsageCall['usage'], + }) + return { ...state, usage: next } +} + function withSignIn(state: SessionViewState, signIn: SignIn): SessionViewState { return { ...state, @@ -615,6 +649,7 @@ export function applyEvent( return isMediaRecord(value) ? withMedia(state, operationId, value, value.subagentRunId) : state + if (event.name === HARNESS_EVENTS.usage) return withUsage(state, value) if (event.name === HARNESS_EVENTS.operationResumed) return withNotice(state, 'info', 'Resumed a turn that a crash stopped.') // A retry that does not continue the partial answer starts the model call diff --git a/packages/ai-harness/src/view/types.ts b/packages/ai-harness/src/view/types.ts index 07f5a9f08c..a451b803a0 100644 --- a/packages/ai-harness/src/view/types.ts +++ b/packages/ai-harness/src/view/types.ts @@ -102,7 +102,7 @@ export interface Approval { message?: string approve: (options?: ApproveOptions) => void /** Say no. `payload` goes to the tool's approval schema, when it has one. */ - reject: (payload?: unknown) => void + reject: (options?: { payload?: unknown }) => void } /** diff --git a/packages/ai-harness/src/wire.ts b/packages/ai-harness/src/wire.ts new file mode 100644 index 0000000000..0359970cdb --- /dev/null +++ b/packages/ai-harness/src/wire.ts @@ -0,0 +1,27 @@ +import { encodeWsFrame } from '@tanstack/ai' +import type { SessionEvent } from './types' + +// The encoded frame of each log entry. Readers of one log get the same entry +// objects, so N readers encode an event once. +const encoded = new WeakMap() + +/** + * The `harness.event` frame of `entry` as JSON text. The event is an AG-UI + * wire chunk (`toWireChunk`): keys that are not in the AG-UI spec are gone. + * The log keeps the full chunk. Only the wire drops keys. + * + * @example + * ```ts + * for await (const entry of session.events({ from })) socket.send(wireFrame(entry)) + * ``` + */ +export function wireFrame(entry: SessionEvent): string { + const cached = encoded.get(entry) + if (cached !== undefined) return cached + const { event, ...tags } = entry + // `encodeWsFrame` with no id is the JSON of `toWireChunk(event)`. The slice + // drops the closing `}` of the tags, so the event goes in after them. + const text = `${JSON.stringify({ type: 'harness.event', ...tags }).slice(0, -1)},"event":${encodeWsFrame(event, undefined)}}` + encoded.set(entry, text) + return text +} diff --git a/packages/ai-harness/src/worker.ts b/packages/ai-harness/src/worker.ts index a2c133dfe6..643a0bd133 100644 --- a/packages/ai-harness/src/worker.ts +++ b/packages/ai-harness/src/worker.ts @@ -1,14 +1,18 @@ +import { randomUUID } from 'node:crypto' import { createInterface } from 'node:readline' import { createHarnessHost } from './host' import { scrubbedBy } from './session/harness-session' import { - HARNESS_PROTOCOL_VERSION, applyInput, + HARNESS_PROTOCOL_VERSION, + versionError, parseControlFrame, + requestIdOf, } from './protocol' +import { wireFrame } from './wire' import type { AnyHarness } from './define' import type { HarnessPersistence } from './host' -import type { HostFrame } from './protocol' +import type { ControlFrame, HarnessErrorCode, HostFrame } from './protocol' export interface HarnessWorkerOptions { input?: NodeJS.ReadableStream @@ -19,7 +23,8 @@ export interface HarnessWorkerOptions { /** * 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. + * stdin closes, or after a `subscribe` with a `v` that it does not speak. + * `artifactText` starts workers like this. * * @example * ```ts @@ -31,32 +36,62 @@ export async function runHarnessWorker( options: HarnessWorkerOptions = {}, ): Promise { const output = options.output ?? process.stdout - const send = (frame: HostFrame) => output.write(`${JSON.stringify(frame)}\n`) + const send = (frame: HostFrame | string) => + output.write( + `${typeof frame === 'string' ? frame : JSON.stringify(frame)}\n`, + ) const host = createHarnessHost( options.persistence ? { persistence: options.persistence } : {}, ) const reader = new AbortController() let session: Awaited> | undefined - /** An error frame, without the secrets of the session once it is open. */ - const sendError = (error: unknown) => { + const lines = createInterface({ + input: options.input ?? process.stdin, + crlfDelay: Infinity, + }) + + // The error frame names the request that failed, so the client can match + // it. A frame with no `requestId` gets one from the worker. The message + // has no secrets of the session once it is open. + const fail = ( + error: unknown, + requestId: string, + code: HarnessErrorCode = 'internal', + ) => { const scrubbed = scrubbedBy(session, error) send({ type: 'harness.error', message: scrubbed instanceof Error ? scrubbed.message : String(scrubbed), + code, + requestId, }) } - const lines = createInterface({ - input: options.input ?? process.stdin, - crlfDelay: Infinity, - }) try { for await (const line of lines) { if (line.trim() === '') continue + // Read the id before the full parse: a frame with a bad `input` still + // names its request. + const requestId = requestIdOf(line) ?? randomUUID() + let frame: ControlFrame + try { + frame = parseControlFrame(line) + } catch (error) { + fail(error, requestId, 'invalid_frame') + continue + } try { - const frame = parseControlFrame(line) if (frame.type === 'harness.subscribe') { - if (session) throw new Error('Already subscribed.') + if (session) { + fail(new Error('Already subscribed.'), requestId, 'invalid_frame') + continue + } + const refused = versionError(frame.v) + if (refused) { + // The worker speaks no other version, so it stops. + send({ type: 'harness.error', requestId, ...refused }) + break + } session = await host.open(harness, { threadId: frame.threadId }) send({ type: 'harness.hello', @@ -68,12 +103,19 @@ export async function runHarnessWorker( signal: reader.signal, }) void (async () => { - for await (const entry of events) - send({ type: 'harness.event', ...entry }) - })().catch((error: unknown) => sendError(error)) + // Each event goes out as an AG-UI wire chunk. + for await (const entry of events) send(wireFrame(entry)) + })().catch((error: unknown) => fail(error, requestId)) + continue + } + if (!session) { + fail( + new Error('Send harness.subscribe first.'), + requestId, + 'invalid_frame', + ) continue } - if (!session) throw new Error('Send harness.subscribe first.') if (frame.type === 'harness.snapshot') { send({ type: 'harness.snapshot', snapshot: session.snapshot() }) continue @@ -82,6 +124,7 @@ export async function runHarnessWorker( harness, session, input: frame.input, + ...(frame.traceparent ? { traceparent: frame.traceparent } : {}), }) send({ type: 'harness.receipt', @@ -89,7 +132,7 @@ export async function runHarnessWorker( ...receipt, }) } catch (error) { - sendError(error) + fail(error, requestId) } } } finally { diff --git a/packages/ai-harness/tests/access-work.test.ts b/packages/ai-harness/tests/access-work.test.ts new file mode 100644 index 0000000000..3978bd29c4 --- /dev/null +++ b/packages/ai-harness/tests/access-work.test.ts @@ -0,0 +1,385 @@ +import { describe, expect, it, vi } from 'vitest' +import { defineAgent } from '@tanstack/ai' +import { + memoryLogBackend, + memoryLogStore, + memoryPersistence, +} from '@tanstack/ai-persistence' +import { + AccessDeniedError, + createHarnessHost, + defineHarness, + definePlugin, + membersOnly, +} from '../src' +import { loadLogState, logIdOf, sessionOf } from '../src/log' +import { + mockAdapter, + storedThreadState, + text, + untilAborted, + wrapLogStore, +} from './helpers' +import type { LogStore } from '@tanstack/ai-persistence' +import type { AgentActions, Operation } from '../src' + +// Coverage runs in CI are slow, so polled checks wait up to 5 s. +const waitFor = ( + check: () => T | Promise, + options?: { timeout?: number; interval?: number }, +) => vi.waitFor(check, { timeout: 5_000, ...options }) + +interface User { + id: string + tenantId: string +} + +const olga: User = { id: 'olga', tenantId: 'acme' } +const bea: User = { id: 'bea', tenantId: 'acme' } +const vic: User = { id: 'vic', tenantId: 'acme' } + +const THREAD = 't1' +const NAME = 'test/access-work' + +/** An agent whose run goes on until it is cancelled. */ +const slow = defineAgent({ + name: 'slow', + description: 'Runs until it is cancelled', + run: (ctx) => ctx.chat({ adapter: mockAdapter(untilAborted()).adapter }), +}) + +/** + * Wait until `operation` runs. The receipt of an agent run is not awaited: + * the run answers through its result. + */ +const running = (operation: Operation) => + waitFor(() => expect(operation.status()).toBe('running')) + +/** The op of each input in the log of `THREAD`, and who sent it. */ +async function loggedInputs(log: LogStore, tenantId?: string) { + const state = await loadLogState({ + store: log, + logId: logIdOf({ tenantId, logId: THREAD }), + rootThread: THREAD, + }) + return [...sessionOf(state, THREAD).inputs.values()].map((logged) => ({ + op: logged.input.op, + by: logged.principal?.id, + })) +} + +/** A host with `membersOnly()` access, and a harness with the `slow` agent. */ +function accessHost() { + const persistence = memoryPersistence() + const host = createHarnessHost({ + persistence, + access: { identify: (): User | null => null, authorize: membersOnly() }, + }) + const harness = defineHarness({ + name: NAME, + adapter: mockAdapter(() => text('ok')).adapter, + agents: [slow], + }) + const open = (principal: User) => + host.open(harness, { threadId: THREAD, principal }) + const share = (userId: string, role: 'member' | 'viewer') => + host.sessions.share({ threadId: THREAD, userId, role, principal: olga }) + return { host, harness, open, share, log: persistence.stores.log } +} + +describe('who acts on an agent run', () => { + it('sends to and cancels a run as the user of the run, not as the first user of the session', async () => { + let agents: AgentActions | undefined + const starter = definePlugin({ + name: 'test/starter', + setup: (ctx) => { + agents = ctx.agents + }, + }) + const persistence = memoryPersistence() + const host = createHarnessHost({ persistence }) + const harness = defineHarness({ + name: NAME, + adapter: mockAdapter(() => text('ok')).adapter, + agents: [slow], + plugins: () => [starter], + }) + // Olga opens the session first. + await host.open(harness, { threadId: THREAD, principal: olga }) + if (!agents) throw new Error('No plugin context.') + const run = agents.start('slow', undefined, { principal: bea }) + await running(run) + + await run.send('Look at the tests first.') + await run.cancel() + + await expect(run).rejects.toThrow() + const inputs = await loggedInputs(persistence.stores.log, 'acme') + expect( + inputs.filter( + (input) => input.op === 'agentMessage' || input.op === 'cancel', + ), + ).toEqual([ + { op: 'agentMessage', by: 'bea' }, + { op: 'cancel', by: 'bea' }, + ]) + await host.close() + }) + + it('runs handle.agents..run as the principal that it passes, as start does', async () => { + const persistence = memoryPersistence() + const host = createHarnessHost({ persistence }) + const harness = defineHarness({ + name: NAME, + adapter: mockAdapter(() => text('ok')).adapter, + agents: [slow], + }) + const system = await host.open(harness, { + threadId: THREAD, + principal: { id: 'system' }, + }) + + const run = system.agents.slow.run(undefined, { principal: bea }) + await running(run) + + const inputs = await loggedInputs(persistence.stores.log) + expect(inputs.filter((input) => input.op === 'agent')).toEqual([ + { op: 'agent', by: 'bea' }, + ]) + await run.cancel() + await expect(run).rejects.toThrow() + await host.close() + }) +}) + +describe('operations that a handle returns', () => { + it('check access on cancel and send, as the calls of the handle do', async () => { + const { host, open, share } = accessHost() + const owner = await open(olga) + await share('bea', 'member') + await share('vic', 'viewer') + const member = await open(bea) + const viewer = await open(vic) + const run = owner.agents.slow.start() + await running(run) + + const seen = viewer.agentRun(run.id) + if (!seen) throw new Error('The viewer does not see the run.') + await expect(seen.cancel()).rejects.toBeInstanceOf(AccessDeniedError) + await expect(seen.send('Stop.')).rejects.toBeInstanceOf(AccessDeniedError) + // A member may cancel only the runs that the member started. + await expect(member.operation(run.id)?.cancel()).rejects.toBeInstanceOf( + AccessDeniedError, + ) + expect(run.status()).toBe('running') + + const own = member.agents.slow.start() + await running(own) + expect((await own.cancel()).status).toBe('accepted') + await expect(own).rejects.toThrow() + await run.cancel() + await expect(run).rejects.toThrow() + await host.close() + }) +}) + +describe('a user that loses access', () => { + it('ends a running events() read of a handle and of a view', async () => { + const { host, harness, open, share } = accessHost() + const owner = await open(olga) + await share('bea', 'member') + const member = await open(bea) + const view = await host.observe(harness, { + threadId: THREAD, + principal: bea, + }) + const readAll = async (events: AsyncIterable) => { + for await (const _event of events) { + // Read until the read ends. + } + } + const handleRead = readAll(member.events()) + const viewRead = readAll(view.events()) + await owner.send('Hi.') + + await host.sessions.unshare({ + threadId: THREAD, + userId: 'bea', + principal: olga, + }) + + await expect(handleRead).rejects.toBeInstanceOf(AccessDeniedError) + await expect(viewRead).rejects.toBeInstanceOf(AccessDeniedError) + await expect(view.transcript()).rejects.toBeInstanceOf(AccessDeniedError) + await host.close() + }) + + it('ends an events() read that is first pulled after the revoke', async () => { + const { host, open, share } = accessHost() + await open(olga) + await share('bea', 'member') + const member = await open(bea) + const events = member.events() + + await host.revoke({ threadId: THREAD, principal: bea }) + + const read = (async () => { + for await (const _event of events) { + // Read until the read ends. + } + })() + await expect(read).rejects.toBeInstanceOf(AccessDeniedError) + await host.close() + }) + + it('rejects host.revoke when the work of the user did not stop', async () => { + const persistence = memoryPersistence() + let isDown = false + const log = wrapLogStore(persistence.stores.log, { + offer: (input) => + isDown + ? Promise.reject(new Error('The store is down.')) + : persistence.stores.log.offer(input), + }) + const host = createHarnessHost({ + persistence: { ...persistence, stores: { ...persistence.stores, log } }, + access: { identify: (): User | null => null, authorize: membersOnly() }, + }) + const harness = defineHarness({ + name: NAME, + adapter: mockAdapter(() => text('ok')).adapter, + }) + await host.open(harness, { threadId: THREAD, principal: olga }) + await host.sessions.share({ + threadId: THREAD, + userId: 'bea', + role: 'member', + principal: olga, + }) + const member = await host.open(harness, { + threadId: THREAD, + principal: bea, + }) + await member.schedule({ + at: Date.now() + 3_600_000, + input: { op: 'prompt', message: 'Write the daily summary.' }, + }) + isDown = true + + await expect( + host.revoke({ threadId: THREAD, principal: bea }), + ).rejects.toThrow('their work did not stop') + // The access is gone already. + expect(() => member.transcript()).toThrow(AccessDeniedError) + await host.close() + }) + + it('has their running work cancelled by host.revoke', async () => { + const { host, open, share } = accessHost() + await open(olga) + await share('bea', 'member') + const member = await open(bea) + const run = member.agents.slow.start() + await running(run) + + await host.revoke({ threadId: THREAD, principal: bea }) + + await expect(run).rejects.toThrow() + expect(run.status()).toBe('cancelled') + await host.close() + }) + + it('has their work and their schedules stopped by an unshare on another host', async () => { + const backend = memoryLogBackend() + const shared = memoryPersistence() + const on = (log: LogStore) => ({ + ...shared, + stores: { ...shared.stores, log }, + }) + const access = { + identify: (): User | null => null, + authorize: membersOnly(), + } + const main = mockAdapter(() => text('ran')) + const harness = defineHarness({ + name: NAME, + adapter: main.adapter, + agents: [slow], + }) + const log = memoryLogStore({ backend }) + const hostA = createHarnessHost({ persistence: on(log), access }) + const hostB = createHarnessHost({ + persistence: on(memoryLogStore({ backend })), + access, + }) + await hostA.open(harness, { threadId: THREAD, principal: olga }) + await hostA.sessions.share({ + threadId: THREAD, + userId: 'bea', + role: 'member', + principal: olga, + }) + const member = await hostA.open(harness, { + threadId: THREAD, + principal: bea, + }) + const run = member.agents.slow.start() + await running(run) + const { scheduleId } = await member.schedule({ + at: Date.now() + 300, + input: { op: 'prompt', message: 'Write the daily summary.' }, + }) + + // Host B has no session of the thread open. + await hostB.sessions.unshare({ + threadId: THREAD, + userId: 'bea', + principal: olga, + }) + + await expect(run).rejects.toThrow() + expect(run.status()).toBe('cancelled') + // The time of the schedule passes, and its prompt does not run. + await waitFor(async () => + expect( + ( + await storedThreadState(log, { threadId: THREAD, tenantId: 'acme' }) + ).schedules.has(scheduleId), + ).toBe(false), + ) + expect(JSON.stringify(main.calls)).not.toContain('daily summary') + await hostA.close() + await hostB.close() + }) + + it('has their schedules removed at once by an unshare, not at their time', async () => { + const { host, open, share, log } = accessHost() + const owner = await open(olga) + await share('bea', 'member') + const member = await open(bea) + // A run of the owner keeps the log owned, so the owner applies offers. + const run = owner.agents.slow.start() + await running(run) + const { scheduleId } = await member.schedule({ + at: Date.now() + 3_600_000, + input: { op: 'prompt', message: 'Write the daily summary.' }, + }) + const schedules = async () => + (await storedThreadState(log, { threadId: THREAD, tenantId: 'acme' })) + .schedules + expect((await schedules()).has(scheduleId)).toBe(true) + + await host.sessions.unshare({ + threadId: THREAD, + userId: 'bea', + principal: olga, + }) + + await waitFor(async () => + expect((await schedules()).has(scheduleId)).toBe(false), + ) + await run.cancel() + await expect(run).rejects.toThrow() + await host.close() + }) +}) diff --git a/packages/ai-harness/tests/agent-group.test.ts b/packages/ai-harness/tests/agent-group.test.ts index 6d5104afb6..a710d1a210 100644 --- a/packages/ai-harness/tests/agent-group.test.ts +++ b/packages/ai-harness/tests/agent-group.test.ts @@ -249,7 +249,9 @@ async function stopAndRunAgain( const refuse = vi .spyOn(log, 'append') .mockRejectedValue(new Error('the host is gone')) - await expect(a.session.append([{ type: 'app.ping' }])).rejects.toThrow() + await expect( + a.session.append({ records: [{ type: 'app.ping' }] }), + ).rejects.toThrow() await expect(operation).rejects.toThrow() refuse.mockRestore() diff --git a/packages/ai-harness/tests/agent-limits.test.ts b/packages/ai-harness/tests/agent-limits.test.ts index 9b3df42a4d..8272a1f56e 100644 --- a/packages/ai-harness/tests/agent-limits.test.ts +++ b/packages/ai-harness/tests/agent-limits.test.ts @@ -61,7 +61,9 @@ async function stopHost( const refuse = vi .spyOn(log, 'append') .mockRejectedValue(new Error('the host is gone')) - await expect(session.append([{ type: 'app.ping' }])).rejects.toThrow() + await expect( + session.append({ records: [{ type: 'app.ping' }] }), + ).rejects.toThrow() await expect(running).rejects.toThrow() refuse.mockRestore() } diff --git a/packages/ai-harness/tests/agent-resume.test.ts b/packages/ai-harness/tests/agent-resume.test.ts index f8883a6970..b5e803c90e 100644 --- a/packages/ai-harness/tests/agent-resume.test.ts +++ b/packages/ai-harness/tests/agent-resume.test.ts @@ -94,7 +94,9 @@ async function stopDuringAgent(options: { const refuse = vi .spyOn(log, 'append') .mockRejectedValue(new Error('the host is gone')) - await expect(sessionA.append([{ type: 'app.ping' }])).rejects.toThrow() + await expect( + sessionA.append({ records: [{ type: 'app.ping' }] }), + ).rejects.toThrow() await expect(operation).rejects.toThrow() refuse.mockRestore() diff --git a/packages/ai-harness/tests/agent-runs.test.ts b/packages/ai-harness/tests/agent-runs.test.ts index 64ec3b976d..d59c61fd4f 100644 --- a/packages/ai-harness/tests/agent-runs.test.ts +++ b/packages/ai-harness/tests/agent-runs.test.ts @@ -1,6 +1,6 @@ import { describe, expect, it, vi } from 'vitest' import { z } from 'zod' -import { defineAgent, toolDefinition } from '@tanstack/ai' +import { EventType, defineAgent, toolDefinition } from '@tanstack/ai' import { LogConflictError, memoryLogStore, @@ -25,6 +25,7 @@ import type { HarnessPersistence, HarnessSession, Operation, + SessionEvent, } from '../src' // Coverage runs in CI are slow, so polled checks wait up to 5 s. @@ -155,6 +156,48 @@ describe('the thread of an agent run', () => { }) }) +describe('the receipt of an agent run', () => { + it('resolves when the run starts, with the run as its operation', async () => { + const { host, session } = await open( + writing(mockAdapter([() => text('Done.')]).adapter), + ) + + const run = session.agents.writer.start() + + const receipt = await run.receipt + expect(receipt).toMatchObject({ status: 'accepted', operationId: run.id }) + expect(receipt.inputId).toEqual(expect.any(String)) + await run + await host.close() + }) + + it('resolves rejected for an input that the schema of the agent refuses', async () => { + const checked = defineAgent({ + name: 'checked', + description: 'Needs a topic', + inputSchema: z.object({ topic: z.string() }), + run: (ctx) => ctx.chat({ adapter: mockAdapter([]).adapter }), + }) + const { host, session } = await open( + defineHarness({ + name: 'test/agent-runs', + adapter: mockAdapter([]).adapter, + agents: [checked], + }), + ) + + const run = session.agent('checked')?.start({ topic: 7 }) + if (!run) throw new Error('No agent checked.') + + expect(await run.receipt).toMatchObject({ + status: 'rejected', + reason: 'invalid_input', + }) + await expect(run).rejects.toThrow('Input validation failed') + await host.close() + }) +}) + describe('follow-ups to an agent run', () => { it('runs the agent again after the run ended, on its transcript', async () => { const model = mockAdapter([ @@ -514,6 +557,26 @@ describe('agents that agent code starts', () => { result: 'helped', }, ]) + // The events of the lead and of its child, in one stream. + const events = [] + for await (const entry of session.events()) { + events.push(entry) + const { event } = entry + const isLeadEnd = + entry.operationId === run.id && + event.type === EventType.CUSTOM && + event.name === 'harness.operation.finished' + if (isLeadEnd) break + } + const ofChild = events.filter((entry) => entry.operationId === childId) + const ofLead = events.filter((entry) => entry.operationId === run.id) + expect(ofChild.length).toBeGreaterThan(0) + expect(ofLead.length).toBeGreaterThan(0) + // The RUN_STARTED of a child run names its parent. The lead has none. + const startOf = (entries: Array) => + entries.find((entry) => entry.event.type === EventType.RUN_STARTED)?.event + expect(startOf(ofChild)).toMatchObject({ parentRunId: run.id }) + expect(startOf(ofLead)).not.toHaveProperty('parentRunId') await host.close() }) @@ -668,7 +731,9 @@ function durable() { const refuse = vi .spyOn(log, 'append') .mockRejectedValue(new Error('the host is gone')) - await expect(session.append([{ type: 'app.ping' }])).rejects.toThrow() + await expect( + session.append({ records: [{ type: 'app.ping' }] }), + ).rejects.toThrow() await expect(run).rejects.toThrow() refuse.mockRestore() } diff --git a/packages/ai-harness/tests/agents.test.ts b/packages/ai-harness/tests/agents.test.ts index b1ea907ed2..28b5cff170 100644 --- a/packages/ai-harness/tests/agents.test.ts +++ b/packages/ai-harness/tests/agents.test.ts @@ -1,11 +1,11 @@ import { mkdtemp, rm, writeFile } from 'node:fs/promises' import { tmpdir } from 'node:os' import { join } from 'node:path' -import { afterEach, beforeEach, describe, expect, it } from 'vitest' +import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest' import { z } from 'zod' import { toolDefinition } from '@tanstack/ai' import { memoryPersistence } from '@tanstack/ai-persistence' -import { createHarnessHost, defineHarness } from '../src' +import { createHarnessHost, defineHarness, definePlugin } from '../src' import { permissions } from '../src/first-party' import { agents } from '../src/first-party/agents' import { splitFrontmatter } from '../src/first-party/files' @@ -56,6 +56,25 @@ function toolNames(call: { tools: Array<{ name: string }> }) { return call.tools.map((entry) => entry.name) } +/** The names of the tools a subagent got on one call, also when none. */ +function childTools(call: { tools?: Array<{ name: string }> }) { + return (call.tools ?? []).map((entry) => entry.name) +} + +/** A subagent with its own model, and every tool of the turn. */ +const helper: AgentProfile = { + name: 'helper', + description: 'Helps with a task', + mode: 'subagent', + model: 'child-model', +} + +// Coverage runs in CI are slow, so polled checks wait up to 5 s. +const waitFor = ( + check: () => T | Promise, + options?: { timeout?: number; interval?: number }, +) => vi.waitFor(check, { timeout: 5_000, ...options }) + /** The system prompts of one call, as one string. */ function promptsOf(call: { systemPrompts: unknown }) { return JSON.stringify(call.systemPrompts) @@ -316,6 +335,42 @@ describe('agents', () => { await host.close() }) + it('gives a thread that removes agents() none of its subagents, rules, or commands', async () => { + const { adapter, calls } = mockAdapter([ + () => text('one'), + () => text('two'), + ]) + const careful: AgentProfile = { + name: 'careful', + description: 'Never writes', + mode: 'primary', + permissions: [{ tool: 'write_file', decision: 'deny' }], + } + const { host, session } = await open(adapter, [ + permissions(), + agents({ adapter: () => adapter, agents: [careful], default: 'careful' }), + ]) + + await session.send('one') + expect( + await session.configure({ + settings: { plugins: { remove: ['tanstack/agents'] } }, + }), + ).toMatchObject({ status: 'accepted' }) + await session.send('two') + + expect(toolNames(calls[0])).toEqual(['read_file', 'general', 'explore']) + // No subagent tools, and the deny rule of `careful` is gone. + expect(toolNames(calls[1])).toEqual(['read_file', 'write_file']) + expect(session.commands().map((command) => command.name)).not.toContain( + 'agent', + ) + await expect( + session.command({ name: 'agent', input: 'build' }), + ).rejects.toThrow(/Unknown command: agent/) + await host.close() + }) + it('replaces a built-in agent with one of the same name', async () => { const { adapter, calls } = mockAdapter([() => text('planned')]) const plan: AgentProfile = { @@ -370,6 +425,87 @@ describe('agents', () => { expect(toolNames(calls[0])).toEqual(['read_file', 'write_file']) await host.close() }) + it("gives a subagent the tools of the turn that started it, never another user's", async () => { + // Each turn finds one tool of its user, as a sign-in to a server does. + const secrets = definePlugin({ + name: 'test/secrets', + setup: () => ({ + discoverTools: (turn) => [tool(`secret_${turn.principal?.id}`)], + }), + }) + const lead = mockAdapter([ + () => toolCall('helper', { prompt: 'Help.' }), + () => text('Olga is done.'), + () => text('Bea is done.'), + ]) + const child = mockAdapter(() => text('Helped.')) + const host = createHarnessHost({ persistence: memoryPersistence() }) + const harness = defineHarness({ + name: 'test/agents', + adapter: lead.adapter, + tools: [tool('read_file')], + plugins: () => [ + secrets, + agents({ adapter: () => child.adapter, agents: [helper] }), + ], + }) + const olga = await host.open(harness, { + threadId: 't', + principal: { id: 'olga' }, + }) + const bea = await host.open(harness, { + threadId: 't', + principal: { id: 'bea' }, + }) + + await olga.send('Get help.') + await bea.send('Hi.') + // Code starts the agent after the turn of bea. No turn started it. + await olga.agent('helper')?.start('Again.') + + expect(childTools(child.calls[0])).toEqual(['read_file', 'secret_olga']) + expect(childTools(child.calls[1])).toEqual([]) + await host.close() + }) + + it('fails a subagent run whose turn ran on another host with a clear error', async () => { + const persistence = memoryPersistence() + const lead = mockAdapter([ + () => toolCall('helper', { prompt: 'Help.', background: true }), + () => text('It runs.'), + ]) + const child = mockAdapter(() => text('Helped.')) + const harness = defineHarness({ + name: 'test/agents', + adapter: lead.adapter, + tools: [tool('read_file')], + plugins: () => [ + agents({ adapter: () => child.adapter, agents: [helper] }), + ], + }) + const first = createHarnessHost({ persistence }) + const session = await first.open(harness, { threadId: 't' }) + await session.send('Start a helper.') + await waitFor(() => + expect(session.agentRuns().map((run) => run.status)).toEqual([ + 'completed', + ]), + ) + const [started] = session.agentRuns() + await first.close() + + // Another host has no turn of this thread in its memory. + const second = createHarnessHost({ persistence }) + const again = await second.open(harness, { threadId: 't' }) + const run = again.agentRun(started?.operationId ?? '') + const receipt = await run?.send('Go on.', { mode: 'followUp' }) + + await expect(again.operation(receipt?.operationId ?? '')).rejects.toThrow( + 'only in the memory of the host that ran that turn', + ) + expect(child.calls).toHaveLength(1) + await second.close() + }) }) describe('agents from Markdown files', () => { diff --git a/packages/ai-harness/tests/auth-wait.test.ts b/packages/ai-harness/tests/auth-wait.test.ts index 5a6ecc5a51..3153ad4927 100644 --- a/packages/ai-harness/tests/auth-wait.test.ts +++ b/packages/ai-harness/tests/auth-wait.test.ts @@ -494,7 +494,7 @@ describe('sign-ins in a shared thread', () => { host, harness, openAccess: { principal: bob }, - media: { secret: 'test-secret' }, + media: { secret: 'test-secret-of-at-least-32-bytes!' }, }) const session = await host.open(harness, { threadId: 't', diff --git a/packages/ai-harness/tests/background-tools.test.ts b/packages/ai-harness/tests/background-tools.test.ts index e81c865ec4..b09c8ce153 100644 --- a/packages/ai-harness/tests/background-tools.test.ts +++ b/packages/ai-harness/tests/background-tools.test.ts @@ -1,6 +1,10 @@ import { describe, expect, it, vi } from 'vitest' import { defineAgent, toolDefinition } from '@tanstack/ai' -import { memoryPersistence } from '@tanstack/ai-persistence' +import { + memoryLogBackend, + memoryLogStore, + memoryPersistence, +} from '@tanstack/ai-persistence' import { createHarnessHost, defineHarness } from '../src' import { after, @@ -87,6 +91,96 @@ describe('session.background', () => { await host.close() }) + it('keeps a session with a running job open after its last handle closes', async () => { + const finish = gate() + const woke = gate() + const child = mockAdapter(after(finish.opened, 'Found it')) + const writer = defineAgent({ + name: 'writer', + description: 'Writes notes', + run: (ctx) => ctx.chat({ adapter: child.adapter }), + }) + const model = mockAdapter(lead('writer', { prompt: 'Look around' }, woke)) + const harness = defineHarness({ + name: 'test/background', + adapter: model.adapter, + subagents: { agents: [writer] }, + }) + const host = createHarnessHost({ + persistence: memoryPersistence(), + lifecycle: { idleCloseMs: 0 }, + }) + const session = await host.open(harness, { threadId: 't' }) + + const turn = session.send('Look around for me') + await vi.waitFor(() => expect(child.calls).toHaveLength(1)) + expect(await session.background()).toMatchObject({ status: 'accepted' }) + await turn + // No handle holds the session now, but its job still runs. + await session.close() + expect(host.isOpen(harness, { threadId: 't' })).toBe(true) + + finish.open() + await woke.opened + const note = messageTexts(model.calls[2]).at(-1) ?? '' + expect(note.split('\n')[0]).toBe('Background job call-1 ended.') + // After the job and its note, the session is idle and closes. + await vi.waitFor( + () => expect(host.isOpen(harness, { threadId: 't' })).toBe(false), + { timeout: 5_000 }, + ) + await host.close() + }) + + it('opens on another host while the owner runs a job', async () => { + const backend = memoryLogBackend() + const finish = gate() + const woke = gate() + const child = mockAdapter(after(finish.opened, 'Found it')) + const writer = defineAgent({ + name: 'writer', + description: 'Writes notes', + run: (ctx) => ctx.chat({ adapter: child.adapter }), + }) + const model = mockAdapter(lead('writer', { prompt: 'Look around' }, woke)) + const harnessOf = (adapter: AnyTextAdapter) => + defineHarness({ + name: 'test/background', + adapter, + subagents: { agents: [writer] }, + }) + const hostOf = () => + createHarnessHost({ + persistence: { stores: { log: memoryLogStore({ backend }) } }, + }) + const owner = hostOf() + const session = await owner.open(harnessOf(model.adapter), { + threadId: 't', + }) + const turn = session.send('Look around for me') + await vi.waitFor(() => expect(child.calls).toHaveLength(1)) + expect(await session.background()).toMatchObject({ status: 'accepted' }) + await turn + + // The job runs on the owner, so the other host notes nothing. + const other = hostOf() + const reader = await other.open(harnessOf(mockAdapter([]).adapter), { + threadId: 't', + }) + expect( + (await reader.transcript()).some((message) => + String(message.content).includes('stopped when the host restarted'), + ), + ).toBe(false) + + finish.open() + await woke.opened + const note = messageTexts(model.calls[2]).at(-1) ?? '' + expect(note.split('\n')[0]).toBe('Background job call-1 ended.') + await owner.close() + await other.close() + }) + it('leaves a tool without support running', async () => { const finish = gate() const started = gate() diff --git a/packages/ai-harness/tests/build-stub.test.ts b/packages/ai-harness/tests/build-stub.test.ts index 72b9930274..a2479fd109 100644 --- a/packages/ai-harness/tests/build-stub.test.ts +++ b/packages/ai-harness/tests/build-stub.test.ts @@ -43,6 +43,10 @@ createInterface({ input: process.stdin }).on('line', (line) => { send({ type: 'harness.error', message: 'worker crashed' }) return } + // An error frame of an older input of this worker (its first prompt). + if (message === 'noisy') { + send({ type: 'harness.error', message: 'old input failed', code: 'internal', requestId: 'r1' }) + } turns += 1 const operationId = 'op-' + turns const event = (value) => @@ -285,6 +289,12 @@ describe('artifactText with a built stub', () => { if (chunk.type === EventType.TEXT_MESSAGE_CONTENT) text += chunk.delta } expect(text).toBe('turn 2: x') + // The error of another input does not fail this turn. + text = '' + for await (const chunk of stream('noisy')) { + if (chunk.type === EventType.TEXT_MESSAGE_CONTENT) text += chunk.delta + } + expect(text).toBe('turn 3: noisy') } finally { model.dispose() } diff --git a/packages/ai-harness/tests/client-media.test.ts b/packages/ai-harness/tests/client-media.test.ts index 4f3b698676..6ef6836c6c 100644 --- a/packages/ai-harness/tests/client-media.test.ts +++ b/packages/ai-harness/tests/client-media.test.ts @@ -49,7 +49,7 @@ describe('client media', () => { // A view on a bigger buffer: only the bytes of the view are sent. const bytes = new Uint8Array([0, 1, 2, 3, 9]).subarray(1, 4) - const stored = await client.upload({ + const stored = await client.putMedia({ body: bytes, name: 'cat photo.png', mimeType: 'image/png', @@ -71,7 +71,7 @@ describe('client media', () => { ) await expect( - client.upload({ + client.putMedia({ body: new Blob(['big']), name: 'big.mp4', mimeType: 'video/mp4', @@ -89,7 +89,7 @@ describe('client media', () => { ) await expect( - client.upload({ + client.putMedia({ body: new ArrayBuffer(3), name: 'a.png', mimeType: 'image/png', @@ -101,7 +101,7 @@ describe('client media', () => { const { client } = setup(() => Response.json({ id: 'm1' })) await expect( - client.upload({ + client.putMedia({ body: new Blob(['abc']), name: 'a.png', mimeType: 'image/png', @@ -117,7 +117,7 @@ describe('client media', () => { }), ) - expect(await client.mediaUrl('m1')).toEqual({ + expect(await client.mediaUrl({ id: 'm1' })).toEqual({ url: 'http://local/api/harness/media?threadId=thread-m&id=m1&exp=5000&sig=abc', expiresAt: 5000, }) @@ -133,7 +133,7 @@ describe('client media', () => { it('throws when the media-url answer has no signed path', async () => { const { client } = setup(() => Response.json({ expiresAt: 5000 })) - await expect(client.mediaUrl('m1')).rejects.toThrow( + await expect(client.mediaUrl({ id: 'm1' })).rejects.toThrow( 'Harness media-url failed: the answer has no signed path', ) }) diff --git a/packages/ai-harness/tests/client-parse.test.ts b/packages/ai-harness/tests/client-parse.test.ts new file mode 100644 index 0000000000..aa01232455 --- /dev/null +++ b/packages/ai-harness/tests/client-parse.test.ts @@ -0,0 +1,345 @@ +import { afterEach, describe, expect, it, vi } from 'vitest' +import { z } from 'zod' +import { memoryPersistence } from '@tanstack/ai-persistence' +import { createHarnessHost, defineHarness } from '../src' +import { createHarnessHandler } from '../src/server' +import { HarnessProtocolError, createHarnessClient } from '../src/client' +import { mockAdapter, text } from './helpers' + +/** The part of a /control request body that these tests read. */ +const ControlBody = z.object({ input: z.object({ inputId: z.string() }) }) + +afterEach(() => { + vi.useRealTimers() + vi.restoreAllMocks() +}) + +/** A 200 response whose body sends `chunks` as separate reads. */ +function streamed(chunks: ReadonlyArray) { + const encoder = new TextEncoder() + const body = new ReadableStream({ + start(controller) { + for (const chunk of chunks) controller.enqueue(encoder.encode(chunk)) + controller.close() + }, + }) + return new Response(body, { status: 200 }) +} + +/** A client on `fetch`, for the thread `t`. */ +function clientOn( + fetch: typeof globalThis.fetch, + reconnectDelayMs?: number, + idleTimeoutMs?: number, +) { + return createHarnessClient({ + url: 'http://x/api/harness', + threadId: 't', + fetch, + ...(reconnectDelayMs === undefined ? {} : { reconnectDelayMs }), + ...(idleTimeoutMs === undefined ? {} : { idleTimeoutMs }), + }) +} + +/** Every frame of a host-events body made of `chunks`. */ +async function hostFrames(chunks: ReadonlyArray) { + const client = clientOn(async () => streamed(chunks)) + const frames: Array = [] + for await (const frame of client.hostEvents()) frames.push(frame) + return frames +} + +/** One session event frame with the cursor `cursor`. */ +const sessionFrame = (cursor: string) => + `data: ${JSON.stringify({ + type: 'harness.event', + cursor, + operationId: 'op', + event: { type: 'CUSTOM', name: 'tick', value: cursor, timestamp: 1 }, + })}\n\n` + +describe('SSE parsing', () => { + it('ends lines with CR, LF, or CRLF, also a CRLF split across two reads', async () => { + expect( + await hostFrames([ + 'data: {"n":1}\r\n\r', + '\ndata: {"n":2}\r\rdata: {"n":3}\n\n', + ]), + ).toEqual([{ n: 1 }, { n: 2 }, { n: 3 }]) + }) + + it('joins the data lines of one event, with or without a space', async () => { + expect( + await hostFrames(['data: {"a":\n', 'data:1,\ndata: "b": 2}\n\n']), + ).toEqual([{ a: 1, b: 2 }]) + }) + + it('skips comments and other fields', async () => { + expect( + await hostFrames([ + ': heartbeat\n\n', + 'event: status\nid: 7\n: note\ndata: {"n":1}\n\n', + ]), + ).toEqual([{ n: 1 }]) + }) + + it('keeps an event split across reads until it ends', async () => { + expect(await hostFrames(['da', 'ta: {"n"', ':1}', '\n', '\n'])).toEqual([ + { n: 1 }, + ]) + }) +}) + +describe('events() reconnects', () => { + it('adds one abort listener for all its connections, and removes it', async () => { + let calls = 0 + // Each connection sends one event and closes, so the client reconnects. + const client = clientOn(async () => { + calls += 1 + return streamed([sessionFrame(String(calls))]) + }, 1) + const controller = new AbortController() + const added = vi.spyOn(controller.signal, 'addEventListener') + const removed = vi.spyOn(controller.signal, 'removeEventListener') + + const seen: Array = [] + // A cursor, so the client does not ask for the snapshot head first. + for await (const entry of client.events({ + from: '0', + signal: controller.signal, + })) { + seen.push(entry.cursor) + if (seen.length === 20) break + } + + expect(seen).toHaveLength(20) + expect(added).toHaveBeenCalledTimes(1) + expect(removed).toHaveBeenCalledTimes(1) + expect(removed.mock.calls[0]?.[1]).toBe(added.mock.calls[0]?.[1]) + }) + + it('waits longer after each failure, with jitter, and starts again after a read', async () => { + vi.useFakeTimers() + // The random part takes a quarter of each wait away. + vi.spyOn(Math, 'random').mockReturnValue(0.5) + const times: Array = [] + const client = clientOn(async () => { + times.push(Date.now()) + // The 8th connection reads one event. Every other one fails. + if (times.length === 8) return streamed([sessionFrame('1')]) + throw new Error('network') + }) + const controller = new AbortController() + const iterator = client + .events({ from: '0', signal: controller.signal }) + [Symbol.asyncIterator]() + + const first = iterator.next() + // The 7 waits before the 8th connection. + await vi.advanceTimersByTimeAsync(68_250) + expect((await first).value?.cursor).toBe('1') + const rest = iterator.next() + await vi.advanceTimersByTimeAsync(2_250) + controller.abort() + expect(await rest).toEqual({ done: true, value: undefined }) + + const waits = times.slice(1).map((time, index) => time - times[index]!) + // 1 s doubled each try, capped at 30 s, times 0.75. After the read, + // the wait is 0.75 s again. + expect(waits).toEqual([ + 750, 1500, 3000, 6000, 12_000, 22_500, 22_500, 750, 1500, + ]) + }) + + it('waits idleTimeoutMs before it connects again, for a server with a long heartbeat', async () => { + vi.useFakeTimers() + const urls: Array = [] + // The connection sends nothing. + const client = clientOn( + async (input) => { + urls.push(String(input)) + return new Response(new ReadableStream(), { status: 200 }) + }, + 1, + 120_000, + ) + const controller = new AbortController() + const next = client + .events({ from: '0', signal: controller.signal }) + [Symbol.asyncIterator]() + .next() + + await vi.advanceTimersByTimeAsync(61_000) + expect(urls).toHaveLength(1) + await vi.advanceTimersByTimeAsync(60_000) + expect(urls).toHaveLength(2) + controller.abort() + await next + }) + + it('connects again from its last cursor when a connection sends nothing for 60 s', async () => { + vi.useFakeTimers() + const encoder = new TextEncoder() + const urls: Array = [] + // The first connection sends one event, then nothing: a half-open + // connection. The next ones send one event each. + const client = clientOn(async (input) => { + urls.push(String(input)) + if (urls.length > 1) return streamed([sessionFrame(String(urls.length))]) + const body = new ReadableStream({ + start(controller) { + controller.enqueue(encoder.encode(sessionFrame('1'))) + }, + }) + return new Response(body, { status: 200 }) + }, 1) + const controller = new AbortController() + const iterator = client + .events({ from: '0', signal: controller.signal }) + [Symbol.asyncIterator]() + + expect((await iterator.next()).value?.cursor).toBe('1') + const next = iterator.next() + await vi.advanceTimersByTimeAsync(59_000) + expect(urls).toHaveLength(1) + await vi.advanceTimersByTimeAsync(1_100) + + expect((await next).value?.cursor).toBe('2') + expect(urls[1]).toContain('from=1') + controller.abort() + }) + + it('stops the wait before a reconnect when the signal aborts', async () => { + vi.useFakeTimers() + const client = clientOn(async () => { + throw new Error('network') + }) + const controller = new AbortController() + const next = client + .events({ signal: controller.signal }) + [Symbol.asyncIterator]() + .next() + await vi.advanceTimersByTimeAsync(10) + controller.abort() + + expect(await next).toEqual({ done: true, value: undefined }) + }) +}) + +describe('harness.error frames', () => { + const errorFrame = (code: string) => + `data: ${JSON.stringify({ type: 'harness.error', code, message: `${code} happened`, requestId: 'r' })}\n\n` + + it('events() connects again from its cursor after slow_consumer, and throws any other error', async () => { + const urls: Array = [] + const client = clientOn(async (input) => { + urls.push(String(input)) + return urls.length === 1 + ? streamed([sessionFrame('1'), errorFrame('slow_consumer')]) + : streamed([sessionFrame('2'), errorFrame('internal')]) + }, 1) + const seen: Array = [] + + const reading = (async () => { + for await (const entry of client.events({ from: '0' })) { + seen.push(entry.cursor) + } + })() + + await expect(reading).rejects.toMatchObject({ + name: 'HarnessProtocolError', + code: 'internal', + message: 'internal happened', + }) + await expect(reading).rejects.toBeInstanceOf(HarnessProtocolError) + expect(seen).toEqual(['1', '2']) + expect(urls).toHaveLength(2) + expect(urls[1]).toContain('from=1') + }) + + it('hostEvents() connects again after slow_consumer, and throws any other error', async () => { + let calls = 0 + const client = clientOn(async () => { + calls += 1 + return calls === 1 + ? streamed(['data: {"n":1}\n\n', errorFrame('slow_consumer')]) + : streamed(['data: {"n":2}\n\n', errorFrame('forbidden')]) + }) + const frames: Array = [] + + const reading = (async () => { + for await (const frame of client.hostEvents()) frames.push(frame) + })() + + await expect(reading).rejects.toMatchObject({ code: 'forbidden' }) + expect(frames).toEqual([{ n: 1 }, { n: 2 }]) + }) +}) + +describe('control inputs', () => { + /** A client on a real handler that keeps each /control request body. */ + function recorded() { + const { adapter, calls } = mockAdapter([() => text('once')]) + const harness = defineHarness({ name: 'test/client-parse', adapter }) + const host = createHarnessHost({ persistence: memoryPersistence() }) + const handler = createHarnessHandler({ + host, + harness, + openAccess: { principal: { id: 'u' } }, + }) + const bodies: Array = [] + const client = createHarnessClient({ + url: 'http://local/api/harness', + threadId: 'thread-ids', + fetch: async (input, init) => { + if (typeof init?.body === 'string') bodies.push(init.body) + return handler(new Request(input, init)) + }, + }) + /** Send the body again, as a retry of the same request does. */ + const retry = async (body: string) => { + const response = await handler( + new Request('http://local/api/harness/control', { + method: 'POST', + headers: { 'Content-Type': 'application/json' }, + body, + }), + ) + return response.json() + } + return { host, client, bodies, calls, retry } + } + + /** The `inputId` of the input in a /control request body. */ + const inputIdOf = (body: string | undefined) => + ControlBody.parse(JSON.parse(body ?? 'null')).input.inputId + + it('gives a send an inputId, so a retry of it is the same input', async () => { + const { host, client, bodies, calls, retry } = recorded() + + const receipt = await client.send('hi') + const again = await retry(bodies[0] ?? '') + // Wait until the first turn ends. A second turn would start before. + await vi.waitFor(async () => { + expect(calls.length).toBeGreaterThan(0) + expect((await client.snapshot()).status).toBe('idle') + }) + await host.close() + + expect(inputIdOf(bodies[0])).toMatch( + /^[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}$/, + ) + expect(receipt.inputId).toBe(inputIdOf(bodies[0])) + expect(again).toEqual(receipt) + expect(calls).toHaveLength(1) + }) + + it('keeps the inputId of the input to change in cancelInput', async () => { + const { host, client, bodies } = recorded() + + await client.cancelInput({ inputId: 'waiting-1' }).catch(() => undefined) + await host.close() + + expect(inputIdOf(bodies[0])).toBe('waiting-1') + }) +}) diff --git a/packages/ai-harness/tests/client-reads.test.ts b/packages/ai-harness/tests/client-reads.test.ts index 212caadaeb..08fda1239d 100644 --- a/packages/ai-harness/tests/client-reads.test.ts +++ b/packages/ai-harness/tests/client-reads.test.ts @@ -117,15 +117,15 @@ describe('client reads and actions', () => { ) const [first] = await client.transcript() - const renamed = await client.renameSession({ + const renamed = await client.sessions.rename({ threadId: 'thread-r', title: 'Greeting', }) - const pinned = await client.updateSession({ + const pinned = await client.sessions.update({ threadId: 'thread-r', pinned: true, }) - const fork = await client.forkSession({ + const fork = await client.sessions.fork({ threadId: 'thread-r', at: { through: first?.id ?? '' }, }) @@ -192,13 +192,15 @@ describe('client reads and actions', () => { const statusOf = async (path: string) => (await handler(new Request(`http://local/api/harness/${path}`))).status - const transcript = await client.agentTranscript(childId) + const transcript = await client.agentTranscript({ subagentRunId: childId }) expect(transcript.map((message) => message.content)).toEqual([ 'Take notes', 'Note 1', ]) - expect(await client.agentTranscript('unknown')).toEqual([]) + expect(await client.agentTranscript({ subagentRunId: 'unknown' })).toEqual( + [], + ) expect(await statusOf('agent-transcript?threadId=thread-a')).toBe(400) expect( await statusOf( diff --git a/packages/ai-harness/tests/client.test.ts b/packages/ai-harness/tests/client.test.ts index 17fe2fea62..0cd750aa4b 100644 --- a/packages/ai-harness/tests/client.test.ts +++ b/packages/ai-harness/tests/client.test.ts @@ -52,7 +52,11 @@ describe('createHarnessClient', () => { const controller = new AbortController() const seen: Array = [] - for await (const entry of client.events({ signal: controller.signal })) { + // From the start: the turn can send its events before this read. + for await (const entry of client.events({ + from: '0', + signal: controller.signal, + })) { seen.push(entry) if ( entry.event.type === 'CUSTOM' && @@ -181,7 +185,9 @@ describe('createHarnessClient', () => { ).toContain('billing'), ) expect(calls).toHaveLength(0) - expect(await client.append([{ type: 'app.signal' }])).toMatchObject({ + expect( + await client.append({ records: [{ type: 'app.signal' }] }), + ).toMatchObject({ status: 'rejected', reason: 'not_exposed', }) @@ -246,6 +252,81 @@ describe('createHarnessClient', () => { }) }) +describe('client.events with no from', () => { + it('starts at the head, so it does not send the old events again', async () => { + const { adapter } = mockAdapter([ + () => text('the old answer'), + () => text('the new answer'), + ]) + const harness = defineHarness({ name: 'test/client-head', adapter }) + const host = createHarnessHost({ persistence: memoryPersistence() }) + const handler = createHarnessHandler({ + host, + harness, + openAccess: { principal: { id: 'u' } }, + }) + const client = createHarnessClient({ + url: 'http://local/api/harness', + threadId: 't-head', + fetch: (input, init) => handler(new Request(input, init)), + }) + const session = await host.open(harness, { threadId: 't-head' }) + await session.send('one') + const { cursor } = await client.snapshot() + + const controller = new AbortController() + const seen: Array = [] + const opened = gate() + const reading = (async () => { + const events = client.events({ + signal: controller.signal, + onConnection: (state) => { + if (state === 'open') opened.open() + }, + }) + for await (const entry of events) { + seen.push(entry) + if ( + entry.event.type === 'CUSTOM' && + entry.event.name === 'harness.operation.finished' + ) { + controller.abort() + } + } + })() + // The read asks for the head first, then streams from it. + await opened.opened + await session.send('two') + await reading + + const body = JSON.stringify(seen) + expect(body).toContain('the new answer') + expect(body).not.toContain('the old answer') + expect(seen.every((entry) => Number(entry.cursor) > Number(cursor))).toBe( + true, + ) + await session.close() + await host.close() + }) + + it('throws, with no retry, when the snapshot has no cursor', async () => { + let reads = 0 + const client = createHarnessClient({ + url: 'http://local/api/harness', + threadId: 't', + fetch: async () => { + reads += 1 + return Response.json({ status: 'idle' }) + }, + }) + + await expect(async () => { + for await (const _ of client.events()) break + }).rejects.toThrow('no string "cursor"') + expect(reads).toBe(1) + }) +}) + describe('client.send to an agent run', () => { it('sends a message to an agent run, as the request principal', async () => { const hold = gate() diff --git a/packages/ai-harness/tests/coding-bash.test.ts b/packages/ai-harness/tests/coding-bash.test.ts index 9f028a0e94..99118e88c6 100644 --- a/packages/ai-harness/tests/coding-bash.test.ts +++ b/packages/ai-harness/tests/coding-bash.test.ts @@ -213,6 +213,23 @@ describe('bash', () => { expect(result).toBe(`${LONG_RESULT}\n[Full output saved to ${path}.]`) }) + it('saves the output that scrub gives, so the file keeps no secret', async () => { + const { env, files } = workspace({ + ...LONG, + stdout: 'key=sk-secret\n'.repeat(3000), + }) + await bash( + env, + { command: 'test' }, + { + spillDir: '.agent/bash', + scrub: (text) => text.replaceAll('sk-secret', '[redacted]'), + }, + ) + const [saved] = files.values() + expect(saved).toBe('key=[redacted]\n'.repeat(3000)) + }) + it('saves it with a POSIX path in a sandbox, also on a Windows host', async () => { const { env, files } = workspace(LONG, { shell: 'sh' }) await bash( diff --git a/packages/ai-harness/tests/coding-formatter.test.ts b/packages/ai-harness/tests/coding-formatter.test.ts index 2e6471cc32..0f7233edeb 100644 --- a/packages/ai-harness/tests/coding-formatter.test.ts +++ b/packages/ai-harness/tests/coding-formatter.test.ts @@ -4,7 +4,10 @@ import { describe, expect, it, vi } from 'vitest' import { memoryPersistence } from '@tanstack/ai-persistence' import { createHarnessHost, defineHarness, definePlugin } from '../src' import { workspaceTools } from '../src/first-party/coding' -import { permissions } from '../src/first-party/permissions' +import { + listSavedPermissions, + permissions, +} from '../src/first-party/permissions' import { FormatFailed, formatOnWrite, @@ -377,6 +380,48 @@ describe('formatter', () => { }) }) +describe('formatter with a workspace for each thread', () => { + it('runs in the workspace that bind gave the thread, also with permissions()', async () => { + // A Linux sandbox at /workspace, with a prettier config. + const { backend, runs } = project(prettier, ok, 'sh') + const { adapter } = mockAdapter([ + () => toolCall('write_file', { path: 'a.ts', content: 'let a=1' }, 'c1'), + () => text('done'), + ]) + const host = createHarnessHost({ persistence: memoryPersistence() }) + const session = await host.open( + defineHarness({ + name: 'test/formatter-bind', + adapter, + plugins: () => [ + permissions({ rules: [{ tool: 'formatter:*', decision: 'allow' }] }), + workspaceTools({ + workspace: { + bind: async () => ({ ...backend, root: '/workspace' }), + }, + }), + formatter({}), + ], + }), + { threadId: 't' }, + ) + // Not bypass: the rules check the formatter of the call before it runs. + await session.command({ name: 'mode', input: 'acceptEdits' }) + + await session.send('write it') + expect(session.snapshot().pendingWaits).toEqual([]) + await host.close() + + expect(runs).toEqual([ + { + command: `npx --no-install prettier --write '/workspace/a.ts'`, + cwd: '/workspace', + timeoutMs: 20_000, + }, + ]) + }) +}) + describe('formatter events', () => { it('sends tanstack/formatter:rewrote when the formatter changed the written file', async () => { // The fake prettier puts spaces around `=`. @@ -431,12 +476,14 @@ describe('formatter with permissions()', () => { persistence = memoryPersistence(), files: Record = {}, rules: Array = [], + /** The start of the tool call ids. A thread needs new ids on each host. */ + ids = 'c', ) { const { backend, runs } = project(files) const { adapter, calls } = mockAdapter([ ...paths.map( (path, index) => () => - toolCall('write_file', { path, content: '' }, `c${index}`), + toolCall('write_file', { path, content: '' }, `${ids}${index}`), ), () => text('done'), ]) @@ -545,6 +592,75 @@ describe('formatter with permissions()', () => { await second.host.close() }) + it('asks again after an always answer when the formatter config changed', async () => { + const persistence = memoryPersistence() + const first = await openWith('acceptEdits', ['a.ts'], persistence, { + 'prettier.config.js': 'export default {}', + }) + await first.session.send('write it') + await answer(first.session, { answer: 'always' }) + expect(first.runs).toHaveLength(1) + await first.host.close() + + // The config now loads a plugin: its code runs with the formatter. + const second = await openWith( + 'acceptEdits', + ['b.ts'], + persistence, + { 'prettier.config.js': "export default { plugins: ['./evil.js'] }" }, + [], + 'second-', + ) + await second.session.send('write again') + const message = await answer(second.session, { answer: 'once' }) + expect(message).toContain('Allow write_file') + expect(second.runs.map((run) => run.command)).toEqual([ + `npx --no-install prettier --write ${quoted('b.ts')}`, + ]) + await second.host.close() + }) + + it('asks again after an always answer for a write of a config below the root', async () => { + const persistence = memoryPersistence() + const first = await openWith('acceptEdits', ['a.ts'], persistence, prettier) + await first.session.send('write it') + await answer(first.session, { answer: 'always' }) + await first.host.close() + + // Prettier loads the nested config, and runs its code. The root + // config did not change. + const second = await openWith( + 'acceptEdits', + ['src/prettier.config.js'], + persistence, + prettier, + [], + 'second-', + ) + await second.session.send('write the nested config') + const message = await answer(second.session, { answer: 'once' }) + expect(message).toContain('Allow write_file') + await second.host.close() + }) + + it('saves no rule for the formatter of a nested config on an always answer', async () => { + const persistence = memoryPersistence() + const { host, session } = await openWith( + 'acceptEdits', + ['src/prettier.config.js'], + persistence, + prettier, + ) + await session.send('write the nested config') + + await answer(session, { answer: 'always' }) + + const saved = await listSavedPermissions(persistence.stores, root) + expect(saved.map((rule) => rule.tool)).toContain('formatter:prettier') + expect(saved.some((rule) => rule.tool.includes('@unsaved:'))).toBe(false) + await host.close() + }) + it('runs with no question in bypass', async () => { const { host, session, runs } = await openWith('bypass', [ 'prettier.config.js', diff --git a/packages/ai-harness/tests/coding-jobs.test.ts b/packages/ai-harness/tests/coding-jobs.test.ts index 04edd76afb..be9b63be22 100644 --- a/packages/ai-harness/tests/coding-jobs.test.ts +++ b/packages/ai-harness/tests/coding-jobs.test.ts @@ -5,6 +5,7 @@ import { afterEach, describe, expect, it, vi } from 'vitest' import { hostBackend } from '../src/first-party/coding/backend' import { bashTools } from '../src/first-party/coding/bash' import { createOutputTail } from '../src/first-party/coding/ring-buffer' +import { createWorkspaceTools } from '../src/first-party/coding/workspace' import type { ToolExecutionContext } from '@tanstack/ai' import type { PluginEvent } from '../src/extensions' import type { @@ -380,6 +381,39 @@ describe('bash job ids', () => { }) }) +describe('the note of a background job', () => { + it('wakes the session as the user of the run that started the job', async () => { + const running = runningJob() + const { env } = workspace({ spawn: () => running.job }) + const note = vi.fn(async () => undefined) + const { tools } = createWorkspaceTools( + { root, backend: env.backend, web: false }, + { note }, + ) + const bash = tools.find((tool) => tool.name === 'bash') + const run = { + logId: 'log', + threadId: 't', + runId: 'run-a', + principal: { id: 'user-a' }, + } + + await bash?.execute?.( + { command: 'npm test', background: true }, + { emitCustomEvent: () => undefined, context: { run } }, + ) + running.end(0) + await settle() + + // Another user's turn can run when the job ends. The wake turn is + // still the work of user-a. + expect(note).toHaveBeenCalledWith( + expect.stringContaining('Background job bash-1 ended.'), + { wake: true, principal: { id: 'user-a' } }, + ) + }) +}) + describe('hostBackend.killJob', () => { // Coverage runs in CI are slow, so polled checks wait up to 5 s. const waitFor = (check: () => unknown) => diff --git a/packages/ai-harness/tests/coding-snapshots.test.ts b/packages/ai-harness/tests/coding-snapshots.test.ts index 602aac39c9..c8184d8ed8 100644 --- a/packages/ai-harness/tests/coding-snapshots.test.ts +++ b/packages/ai-harness/tests/coding-snapshots.test.ts @@ -495,6 +495,87 @@ describe.skipIf(!hasGit)('snapshots', () => { }) }) +describe.skipIf(!hasGit)('snapshots with a workspace for each thread', () => { + it('keeps the snapshots in the workspace that bind gave the thread', async () => { + // The workspace of each thread is a folder of its own, not `root`. + const own = join(base, 'thread-a') + await mkdir(own) + await writeFile(join(own, 'a.txt'), 'one') + const { adapter } = mockAdapter([ + () => toolCall('write_file', { path: 'a.txt', content: 'two' }, 'w-1'), + () => text('Done.'), + ]) + const host = createHarnessHost({ persistence: memoryPersistence() }) + const session = await host.open( + defineHarness({ + name: 'test/snapshots-bind', + adapter, + plugins: () => [ + workspaceTools({ + workspace: { bind: async () => ({ ...hostBackend, root: own }) }, + web: false, + }), + snapshots({ dataDir }), + ], + }), + { threadId: THREAD }, + ) + + await session.send('Write a.txt') + expect(await readFile(join(own, 'a.txt'), 'utf8')).toBe('two') + expect(await session.command({ name: 'undo' })).toBe( + 'Undid the last turn. Files restored: 1.', + ) + expect(await readFile(join(own, 'a.txt'), 'utf8')).toBe('one') + await host.close() + }) + + it('unreverts on a new host, which opens the workspace of the thread again from the log', async () => { + const own = join(base, 'thread-a') + await mkdir(own) + await writeFile(join(own, 'a.txt'), 'one') + const persistence = memoryPersistence() + const openOn = (replies: Array) => { + const host = createHarnessHost({ persistence }) + const harness = defineHarness({ + name: 'test/snapshots-bind', + adapter: mockAdapter(replies).adapter, + plugins: () => [ + workspaceTools({ + workspace: { bind: async () => ({ ...hostBackend, root: own }) }, + web: false, + }), + snapshots({ dataDir }), + ], + }) + return { host, opened: host.open(harness, { threadId: THREAD }) } + } + const first = openOn([ + () => toolCall('write_file', { path: 'a.txt', content: 'two' }, 'w-1'), + () => text('First.'), + () => toolCall('write_file', { path: 'a.txt', content: 'three' }, 'w-2'), + () => text('Second.'), + ]) + const session = await first.opened + await session.send('One') + await session.send('Two') + const answer = (await session.transcript()).find( + (message) => message.content === 'First.', + ) + await session.revert({ messageId: answer?.id ?? '' }) + expect(await readFile(join(own, 'a.txt'), 'utf8')).toBe('two') + await first.host.close() + + // No turn ran on this host, so the workspace of the thread is not open. + const second = openOn([]) + const again = await second.opened + + expect(await again.unrevert()).toMatchObject({ status: 'accepted' }) + expect(await readFile(join(own, 'a.txt'), 'utf8')).toBe('three') + await second.host.close() + }) +}) + describe.skipIf(!hasGit)('snapshots of many runs', () => { it("/undo of run A keeps the change that run B made during A's step", async () => { await put('a.txt', 'one') @@ -587,6 +668,27 @@ describe.skipIf(!hasGit)('snapshots of many runs', () => { purgeSnapshots({ root, dataDir }, { threadId: THREAD }), ).resolves.toBeUndefined() }) + + it('host.sessions.delete removes the repository of that thread, also when no session has it open', async () => { + const first = await open({ replies: [() => text('Hi.')] }) + await first.session.close() + const second = await first.host.open( + defineHarness({ + name: 'test/snapshots', + adapter: mockAdapter([]).adapter, + plugins: () => [snapshots({ root, dataDir })], + }), + { threadId: 't2' }, + ) + await second.close() + expect(await readdir(dataDir)).toHaveLength(2) + + await first.host.sessions.delete({ threadId: THREAD }) + + // The repository of the other thread stays. + expect(await readdir(dataDir)).toHaveLength(1) + await first.host.close() + }) }) /** Three turns. Each one changes the files once, with its own tool call. */ @@ -622,7 +724,9 @@ describe.skipIf(!hasGit)('session.revert', () => { const { transcript, firstId } = await runThreeTurns(session) expect(transcript).toHaveLength(12) - expect(await session.revert(firstId)).toMatchObject({ status: 'accepted' }) + expect(await session.revert({ messageId: firstId })).toMatchObject({ + status: 'accepted', + }) expect(await read('a.txt')).toBe('two') expect(await readdir(root)).toEqual(['a.txt']) @@ -634,7 +738,7 @@ describe.skipIf(!hasGit)('session.revert', () => { await put('a.txt', 'one') const { host, session } = await open(threeTurns()) const { transcript, firstId } = await runThreeTurns(session) - await session.revert(firstId) + await session.revert({ messageId: firstId }) expect(await session.unrevert()).toMatchObject({ status: 'accepted' }) @@ -652,7 +756,7 @@ describe.skipIf(!hasGit)('session.revert', () => { replies: [...turns.replies, () => text('Fourth.')], }) const { transcript, firstId } = await runThreeTurns(session) - await session.revert(firstId) + await session.revert({ messageId: firstId }) await session.send('Four') @@ -685,13 +789,15 @@ describe.skipIf(!hasGit)('session.revert', () => { const running = session.send('Again') await waitFor(() => expect(session.snapshot().status).toBe('running')) - expect(await session.revert(message?.id ?? '')).toMatchObject({ + expect( + await session.revert({ messageId: message?.id ?? '' }), + ).toMatchObject({ status: 'rejected', reason: 'busy', }) wait.open() await running - expect(await session.revert('no-such-id')).toMatchObject({ + expect(await session.revert({ messageId: 'no-such-id' })).toMatchObject({ status: 'rejected', reason: 'unknown_message', }) @@ -704,7 +810,7 @@ describe.skipIf(!hasGit)('session.revert', () => { const persistence = memoryPersistence() const first = await open({ ...threeTurns(), persistence }) const { transcript, firstId } = await runThreeTurns(first.session) - await first.session.revert(firstId) + await first.session.revert({ messageId: firstId }) await first.host.close() const second = await open({ replies: [], persistence }) @@ -732,7 +838,7 @@ describe.skipIf(!hasGit)('session.revert keeps the edits of the user', () => { (message) => message.content === 'First.', ) - await session.revert(first?.id ?? '') + await session.revert({ messageId: first?.id ?? '' }) expect(await read('a.txt')).toBe('two') expect(await read('c.txt')).toBe('mine') @@ -747,7 +853,7 @@ describe.skipIf(!hasGit)('session.revert keeps the edits of the user', () => { // The user edits a file that a hidden turn changed. await put('a.txt', 'mine') - await session.revert(firstId) + await session.revert({ messageId: firstId }) expect(await read('a.txt')).toBe('two') await session.unrevert() @@ -780,7 +886,7 @@ describe.skipIf(!hasGit)('session.revert keeps the edits of the user', () => { }) const { firstId } = await runThreeTurns(session) blocking = true - const reverting = session.revert(firstId) + const reverting = session.revert({ messageId: firstId }) await reached.opened const prompting = session.send('Four') @@ -811,7 +917,7 @@ describe.skipIf(!hasGit)('session.revert keeps the edits of the user', () => { replies: [...turns.replies, () => text('Fourth.')], }) const { firstId } = await runThreeTurns(session) - await session.revert(firstId) + await session.revert({ messageId: firstId }) await session.configure({ settings: { cwd: 'sub' } }) await session.send('Four') @@ -831,7 +937,7 @@ describe.skipIf(!hasGit)('session.revert keeps the edits of the user', () => { await mkdir(join(root, 'sub')) const { host, session } = await open(threeTurns()) const { transcript, firstId } = await runThreeTurns(session) - await session.revert(firstId) + await session.revert({ messageId: firstId }) await session.configure({ settings: { cwd: 'sub' } }) await session.unrevert() @@ -849,7 +955,7 @@ describe.skipIf(!hasGit)('session.revert keeps the edits of the user', () => { await put('a.txt', 'one') const { host, session } = await open(threeTurns()) const { transcript, firstId } = await runThreeTurns(session) - await session.revert(firstId) + await session.revert({ messageId: firstId }) const refused = 'A revert stands. Run unrevert, or send a prompt, first.' expect(await session.command({ name: 'undo' })).toBe(refused) @@ -946,7 +1052,7 @@ describe('snapshots with a fake git', () => { // The steps of call-1 to call-505 come after the first answer. Of them, // the revert uses the last 500: from call-6 on. - await session.revert(first.id) + await session.revert({ messageId: first.id }) const restores = commands.filter((command) => command.includes('--name-status'), diff --git a/packages/ai-harness/tests/discover-tools.test.ts b/packages/ai-harness/tests/discover-tools.test.ts index 69cd67fd5d..31a5077934 100644 --- a/packages/ai-harness/tests/discover-tools.test.ts +++ b/packages/ai-harness/tests/discover-tools.test.ts @@ -87,6 +87,110 @@ describe('discoverTools', () => { }) }) +describe('onTurnStart', () => { + it('runs once per turn, before the prompts, and a hook that throws does not stop the turn', async () => { + let turns = 0 + const counter = definePlugin({ + name: 'test/counter', + setup: () => ({ + onTurnStart: (turn) => { + turns += 1 + expect(turn.run.runId).toBe(turn.operationId) + }, + prompts: [() => `Turn ${turns}.`], + }), + }) + const broken = definePlugin({ + name: 'test/broken-start', + setup: () => ({ + onTurnStart: () => { + throw new Error('cannot start') + }, + }), + }) + const { adapter, calls } = mockAdapter([ + () => toolCall('lookup', {}), + () => text('one'), + () => text('two'), + ]) + const host = createHarnessHost({ persistence: memoryPersistence() }) + const session = await host.open( + defineHarness({ + name: 'test/turn-start', + adapter, + tools: [echoTool('lookup', 'found')], + plugins: () => [counter, broken], + }), + { threadId: 't' }, + ) + + expect((await session.send('first')).text).toBe('one') + expect((await session.send('second')).text).toBe('two') + + // Two model calls in the first turn, but one hook call per turn. + expect(turns).toBe(2) + expect(calls[0].systemPrompts).toContain('Turn 1.') + expect(calls[2].systemPrompts).toContain('Turn 2.') + await host.close() + }) +}) + +describe('a tool name that two plugins give', () => { + it('keeps the tool of the earlier plugin and warns once with both names', async () => { + const finder = (name: string, reply: string) => + definePlugin({ + name, + setup: () => ({ discoverTools: () => [echoTool('search', reply)] }), + }) + const { adapter, calls } = mockAdapter([ + () => toolCall('search', {}), + () => text('found'), + () => text('again'), + ]) + const host = createHarnessHost({ persistence: memoryPersistence() }) + const session = await host.open( + defineHarness({ + name: 'test/tool-clash', + adapter, + plugins: () => [ + finder('test/first', 'first'), + finder('test/second', 'second'), + ], + }), + { 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.send('search') + await session.send('search again') + + expect(JSON.stringify(calls[1].messages)).toContain('first') + expect(JSON.stringify(calls[1].messages)).not.toContain('second') + controller.abort() + await reading + const warnings = seen.filter( + (entry) => + entry.event.type === EventType.CUSTOM && + entry.event.name === 'harness.plugin.warning', + ) + expect(warnings.map((entry) => entry.event)).toMatchObject([ + { + value: { + plugin: 'test/second', + message: + 'Two tools are named "search": the tool of plugin test/first is used, and the tool of plugin test/second is not.', + }, + }, + ]) + await host.close() + }) +}) + describe('the default agent loop', () => { it('keeps calling the model past five steps', async () => { const steps = Array.from( diff --git a/packages/ai-harness/tests/durable-background.test.ts b/packages/ai-harness/tests/durable-background.test.ts index 8ed098f448..6801d1fcae 100644 --- a/packages/ai-harness/tests/durable-background.test.ts +++ b/packages/ai-harness/tests/durable-background.test.ts @@ -89,7 +89,7 @@ describe('background work after a crash', () => { ) const turn = crashed.send('Wait') await started.opened - expect(await crashed.background('call-1')).toMatchObject({ + expect(await crashed.background({ toolCallId: 'call-1' })).toMatchObject({ status: 'accepted', }) await turn diff --git a/packages/ai-harness/tests/durable-inputs.test.ts b/packages/ai-harness/tests/durable-inputs.test.ts index 0bb56b56f1..361823df66 100644 --- a/packages/ai-harness/tests/durable-inputs.test.ts +++ b/packages/ai-harness/tests/durable-inputs.test.ts @@ -1060,7 +1060,7 @@ describe('waiting inputs', () => { { inputId: 'third', delivery: 'steer', message: 'third' }, { inputId: 'second', delivery: 'queue', message: 'second' }, ]) - expect(await session.cancelInput('second')).toMatchObject({ + expect(await session.cancelInput({ inputId: 'second' })).toMatchObject({ status: 'accepted', }) expect(await session.settled({ inputId: 'second' })).toMatchObject({ @@ -1093,7 +1093,7 @@ describe('waiting inputs', () => { const turn = session.send('first', { inputId: 'first' }) await waitFor(() => expect(calls).toHaveLength(1)) - expect(await session.cancelInput('first')).toMatchObject({ + expect(await session.cancelInput({ inputId: 'first' })).toMatchObject({ status: 'rejected', reason: 'not_waiting', }) @@ -1166,7 +1166,7 @@ describe('waiting inputs', () => { ), ) - expect(await session.cancelInput('c-1')).toMatchObject({ + expect(await session.cancelInput({ inputId: 'c-1' })).toMatchObject({ status: 'accepted', }) expect(await session.settled({ inputId: 'c-1' })).toMatchObject({ @@ -1362,6 +1362,29 @@ describe('input ids', () => { await second.host.close() }) + it('streams a duplicate after a restart from the start of its run in the log', async () => { + const persistence = durablePersistence() + const first = await openDurable({ + persistence, + replies: [() => text('from the first host')], + }) + await first.session.send('hi', { inputId: 'dup-1' }) + await first.host.close() + + const second = await openDurable({ persistence, replies: [] }) + const replay = second.session.send('hi', { inputId: 'dup-1' }) + await replay + const deltas: Array = [] + for await (const chunk of replay.stream()) { + if (chunk.type === EventType.TEXT_MESSAGE_CONTENT) + deltas.push(chunk.delta) + } + + // The rebuilt operation starts before its run, not at the log head. + expect(deltas.join('')).toBe('from the first host') + await second.host.close() + }) + it('settles an old duplicate after a restart with the text of its own turn', async () => { const persistence = durablePersistence() const release = gate() diff --git a/packages/ai-harness/tests/durable-session.test.ts b/packages/ai-harness/tests/durable-session.test.ts index 1d625876c8..ef5b4651c2 100644 --- a/packages/ai-harness/tests/durable-session.test.ts +++ b/packages/ai-harness/tests/durable-session.test.ts @@ -236,9 +236,9 @@ describe('durable host stores', () => { const durable = await createHarnessHost({ persistence: durablePersistence(), }).open(harness, { threadId: THREAD }) - await expect(durable.append([{ type: 'harness.event' }])).rejects.toThrow( - 'reserved for the harness', - ) + await expect( + durable.append({ records: [{ type: 'harness.event' }] }), + ).rejects.toThrow('reserved for the harness') }) }) @@ -308,7 +308,33 @@ describe('durable session log', () => { await next.close() }) - it('writes host records in the same batch as the waiting events, in order', async () => { + it('logs an append with no inputId as an input, with the id of its receipt', async () => { + const persistence = durablePersistence() + const host = createHarnessHost({ persistence }) + const session = await host.open( + defineHarness({ + name: 'test/append', + adapter: mockAdapter(() => text('ok')).adapter, + }), + { threadId: THREAD }, + ) + + const { inputId } = await session.append({ + records: [{ type: 'app.note' }], + }) + + const state = await loadLogState({ + store: persistence.stores.log, + logId: LOG, + rootThread: THREAD, + }) + expect(sessionOf(state, THREAD).inputs.get(inputId)?.input.op).toBe( + 'append', + ) + await host.close() + }) + + it('writes host records after the waiting events, in order', async () => { const inner = memoryLogStore() const batches: Array> = [] const log = wrapLogStore(inner, { @@ -340,14 +366,20 @@ describe('durable session log', () => { expect(JSON.stringify(batches)).toContain('TEXT_MESSAGE_START'), ) - await session.append([{ type: 'app.note', n: 1 }]) + await session.append({ records: [{ type: 'app.note', n: 1 }] }) - expect(batches.at(-1)).toEqual([ + // The input record of the append takes the waiting events with it, and + // the records come after them. + expect(batches).toContainEqual([ expect.objectContaining({ type: EventType.TEXT_MESSAGE_CONTENT, delta: 'partial', }), + expect.objectContaining({ type: 'harness.input' }), + ]) + expect(batches.at(-2)).toEqual([ { type: 'app.note', n: 1 }, + expect.objectContaining({ type: 'harness.input.applied' }), ]) release.open() await Promise.resolve(turn).catch(() => {}) @@ -361,9 +393,9 @@ describe('durable session log', () => { name: 'check', description: 'Check the build', }).server(async () => { - await session?.append([ - { type: 'app.signal', text: '[signal] tests failed' }, - ]) + await session?.append({ + records: [{ type: 'app.signal', text: '[signal] tests failed' }], + }) await session?.send('more', { busy: 'steer' }).receipt return 'checked' }) @@ -405,9 +437,9 @@ describe('durable session log', () => { name: 'check', description: 'Check the build', }).server(async () => { - await session?.append([ - { type: 'app.signal', text: '[signal] tests failed' }, - ]) + await session?.append({ + records: [{ type: 'app.signal', text: '[signal] tests failed' }], + }) return 'checked' }) const { adapter, calls } = mockAdapter([ @@ -430,9 +462,9 @@ describe('durable session log', () => { await session.send('one') // Between turns: the next turn starts with it. - await session.append([ - { type: 'app.signal', text: '[signal] build failed' }, - ]) + await session.append({ + records: [{ type: 'app.signal', text: '[signal] build failed' }], + }) await session.send('two') expect(messageTexts(calls[1])).toEqual([ @@ -486,9 +518,11 @@ describe('durable session log', () => { await session.send('q1') await session.send('q2') - await session.append([ - { type: 'app.compact', summary: 'Earlier: q1 and a1.', firstKept: 2 }, - ]) + await session.append({ + records: [ + { type: 'app.compact', summary: 'Earlier: q1 and a1.', firstKept: 2 }, + ], + }) await session.send('q3') expect(messageTexts(calls[2])).toEqual([ diff --git a/packages/ai-harness/tests/feed.test.ts b/packages/ai-harness/tests/feed.test.ts index c5b2dc5b37..b1f7eb0484 100644 --- a/packages/ai-harness/tests/feed.test.ts +++ b/packages/ai-harness/tests/feed.test.ts @@ -1,6 +1,7 @@ import { describe, expect, it } from 'vitest' import { EventType } from '@tanstack/ai' import { SessionFeed } from '../src/feed' +import { OperationImpl } from '../src/operation' import type { StreamChunk } from '@tanstack/ai' import type { SessionEvent } from '../src/types' @@ -26,7 +27,7 @@ async function namesOf(events: AsyncIterable) { describe('SessionFeed', () => { it('delivers events published while a reader handles earlier ones', async () => { - const feed = new SessionFeed({ eventBytes: EVENT_BYTES }) + const feed = new SessionFeed({ threadId: 't1', eventBytes: EVENT_BYTES }) const controller = new AbortController() const seen: Array = [] const reading = (async () => { @@ -48,7 +49,7 @@ describe('SessionFeed', () => { }) it('skips filtered events and replays from a cursor', async () => { - const feed = new SessionFeed({ eventBytes: EVENT_BYTES }) + const feed = new SessionFeed({ threadId: 't1', eventBytes: EVENT_BYTES }) feed.publish('a', custom('one')) feed.publish('b', custom('two')) feed.publish('a', custom('three')) @@ -60,30 +61,38 @@ describe('SessionFeed', () => { expect(feed.head()).toBe('3') }) - it('keeps the newest events that fit in eventBytes, and starts an older cursor at the oldest kept', async () => { + it('keeps the newest events that fit in eventBytes, and starts an older cursor with a gap at the oldest kept', async () => { const size = JSON.stringify(custom('e1')).length - const feed = new SessionFeed({ eventBytes: size * 3 }) + const feed = new SessionFeed({ threadId: 't1', eventBytes: size * 3 }) for (const name of ['e1', 'e2', 'e3', 'e4', 'e5']) { feed.publish('op', custom(name)) } feed.close() - expect(await namesOf(feed.read({ from: '0' }))).toEqual(['e3', 'e4', 'e5']) + expect(await namesOf(feed.read({ from: '0' }))).toEqual([ + 'harness.gap', + 'e3', + 'e4', + 'e5', + ]) expect(await namesOf(feed.read({ from: '3' }))).toEqual(['e4', 'e5']) expect(feed.head()).toBe('5') }) it('keeps the newest event when it alone is larger than eventBytes', async () => { - const feed = new SessionFeed({ eventBytes: 1 }) + const feed = new SessionFeed({ threadId: 't1', eventBytes: 1 }) feed.publish('op', custom('e1')) feed.publish('op', custom('e2')) feed.close() - expect(await namesOf(feed.read({ from: '0' }))).toEqual(['e2']) + expect(await namesOf(feed.read({ from: '0' }))).toEqual([ + 'harness.gap', + 'e2', + ]) }) it('ends a read when wake() finds until() true, with no new event', async () => { - const feed = new SessionFeed({ eventBytes: EVENT_BYTES }) + const feed = new SessionFeed({ threadId: 't1', eventBytes: EVENT_BYTES }) let done = false const reading = (async () => { let count = 0 @@ -97,3 +106,25 @@ describe('SessionFeed', () => { expect(await reading).toBe(1) }) }) + +describe('operation stream', () => { + it('starts at the operation start, so an older event of its id does not show', async () => { + const feed = new SessionFeed({ threadId: 't1', eventBytes: EVENT_BYTES }) + // A restart builds an operation again with its stored id. + feed.publish('op-1', custom('before')) + const operation = new OperationImpl( + 'chat', + Object.assign(feed, { listen: () => () => {} }), + () => Promise.reject(new Error('no cancel')), + undefined, + 'op-1', + ) + feed.publish('op-1', custom('after')) + operation.finish('completed', undefined) + + const names: Array = [] + for await (const event of operation.stream()) + names.push(event.type === EventType.CUSTOM ? event.name : '') + expect(names).toEqual(['after']) + }) +}) diff --git a/packages/ai-harness/tests/first-party.test.ts b/packages/ai-harness/tests/first-party.test.ts index 2a1c0d475e..4d39150798 100644 --- a/packages/ai-harness/tests/first-party.test.ts +++ b/packages/ai-harness/tests/first-party.test.ts @@ -187,7 +187,11 @@ describe('todos', () => { await session.send('plan it') expect(await session.command({ name: 'todos' })).toBe('[~] write tests') await session.send('continue') - expect(JSON.stringify(calls[2].systemPrompts)).toContain('[~] write tests') + // A `'context'` prompt: the last system prompt, not a message. + expect(JSON.stringify(calls[2].systemPrompts.at(-1))).toContain( + '[~] write tests', + ) + expect(messageTexts(calls[2]).at(-1)).toBe('continue') const stored = await storedThreadState(persistence.stores.log, { threadId: 't', }) @@ -266,4 +270,38 @@ describe('compact and usage', () => { ) await host.close() }) + + it('counts the summary call of /compact under plugin:tanstack/compact', async () => { + const summarizer = mockAdapter([ + () => + text('We planned a trip.').map((chunk) => + chunk.type === 'RUN_FINISHED' + ? { + ...chunk, + usage: { + promptTokens: 40, + completionTokens: 8, + totalTokens: 48, + }, + } + : chunk, + ), + ]) + const main = mockAdapter([() => text('one'), () => text('two')]) + const { host, session } = await open(main.adapter, [ + compact({ adapter: summarizer.adapter }), + ]) + await session.send('first') + await session.send('second') + + await session.command({ name: 'compact' }) + + expect(session.usage().bySource).toEqual({ + 'plugin:tanstack/compact': expect.objectContaining({ + calls: 1, + totalTokens: 48, + }), + }) + await host.close() + }) }) diff --git a/packages/ai-harness/tests/fork.test.ts b/packages/ai-harness/tests/fork.test.ts index fcaa70567e..1503bc1aea 100644 --- a/packages/ai-harness/tests/fork.test.ts +++ b/packages/ai-harness/tests/fork.test.ts @@ -222,14 +222,14 @@ describe('host.fork', () => { if (!id) throw new Error('The fork lost the media part.') expect(id).not.toBe(record.id) - expect(await fork.getMedia(id)).toMatchObject({ + expect(await fork.getMedia({ id })).toMatchObject({ threadId: FORK, name: 'cat.png', }) expect(await fork.loadMedia({ id })).toEqual(bytes) // Each thread sees only its own media. - expect(await fork.getMedia(record.id)).toBeNull() - expect(await source.getMedia(id)).toBeNull() + expect(await fork.getMedia({ id: record.id })).toBeNull() + expect(await source.getMedia({ id })).toBeNull() await host.close() }) }) diff --git a/packages/ai-harness/tests/gap.test.ts b/packages/ai-harness/tests/gap.test.ts new file mode 100644 index 0000000000..aaac57d05a --- /dev/null +++ b/packages/ai-harness/tests/gap.test.ts @@ -0,0 +1,162 @@ +import { describe, expect, it, vi } from 'vitest' +import { EventType } from '@tanstack/ai' +import { memoryPersistence } from '@tanstack/ai-persistence' +import { HARNESS_EVENTS, createHarnessHost, defineHarness } from '../src' +import { SessionFeed, gapEvent } from '../src/feed' +import { createSessionView } from '../src/view' +import { gate, mockAdapter, text } from './helpers' +import type { StreamChunk } from '@tanstack/ai' +import type { SessionEvent } from '../src/types' +import type { SessionViewSource } from '../src/view' + +// Coverage runs in CI are slow, so polled checks wait up to 5 s. +const waitFor = (check: () => T | Promise) => + vi.waitFor(check, { timeout: 5_000 }) + +const custom = (name: string): StreamChunk => ({ + type: EventType.CUSTOM, + name, + value: {}, + timestamp: 0, +}) + +async function entriesOf(events: AsyncIterable) { + const entries: Array = [] + for await (const entry of events) entries.push(entry) + return entries +} + +describe('harness.gap from the in-memory feed', () => { + it('sends one evicted gap to a cursor whose events the feed dropped, then the kept events', async () => { + const size = JSON.stringify(custom('e1')).length + const feed = new SessionFeed({ threadId: 't1', eventBytes: size * 2 }) + for (const name of ['e1', 'e2', 'e3', 'e4']) + feed.publish('op', custom(name)) + feed.close() + + const entries = await entriesOf(feed.read({ from: '1' })) + + expect(entries.map((entry) => entry.cursor)).toEqual(['2', '3', '4']) + expect(entries[0]).toEqual({ + cursor: '2', + threadId: 't1', + operationId: 'session', + event: expect.objectContaining({ + type: EventType.CUSTOM, + name: HARNESS_EVENTS.gap, + value: { from: '1', resumeAt: '2', reason: 'evicted' }, + }), + }) + }) + + it('sends no gap to a cursor that the feed still has', async () => { + const size = JSON.stringify(custom('e1')).length + const feed = new SessionFeed({ threadId: 't1', eventBytes: size * 2 }) + for (const name of ['e1', 'e2', 'e3', 'e4']) + feed.publish('op', custom(name)) + feed.close() + + const entries = await entriesOf(feed.read({ from: '2' })) + + expect(entries.map((entry) => entry.cursor)).toEqual(['3', '4']) + }) +}) + +describe('the view on an evicted gap', () => { + it('loads the snapshot and the transcript again', async () => { + const harness = defineHarness({ + name: 'test/gap', + adapter: mockAdapter(() => text('Hello back.')).adapter, + }) + const host = createHarnessHost({ persistence: memoryPersistence() }) + const session = await host.open(harness, { threadId: 'gap' }) + await session.send('Hello.') + const transcript = vi.fn(() => session.transcript()) + const snapshot = vi.fn(() => session.snapshot()) + const reading = new AbortController() + const overrides: Partial = { + transcript, + snapshot, + // The source dropped events, and says so first. + events: (options) => + (async function* () { + yield gapEvent( + session.threadId, + '0', + session.snapshot().cursor, + 'evicted', + ) + yield* session.events({ ...options, signal: reading.signal }) + })(), + } + // The session, with the overrides in place of its own methods. + const source = new Proxy(session, { + get: (target, key) => { + if (key in overrides) return overrides[key as keyof SessionViewSource] + const value: unknown = Reflect.get(target, key, target) + return typeof value === 'function' ? value.bind(target) : value + }, + }) + const view = createSessionView(source) + await view.ready + + await waitFor(() => expect(transcript).toHaveBeenCalledTimes(2)) + await waitFor(() => expect(snapshot.mock.calls.length).toBeGreaterThan(1)) + view.dispose() + reading.abort() + await session.close() + await host.close() + }) + + it('applies no backlog event twice after the transcript loads again', async () => { + const harness = defineHarness({ + name: 'test/gap-backlog', + adapter: mockAdapter([() => text('First.'), () => text('Second.')]) + .adapter, + }) + const host = createHarnessHost({ persistence: memoryPersistence() }) + const session = await host.open(harness, { threadId: 'gap-backlog' }) + await session.send('One.') + const before = session.snapshot().cursor + const transcript = vi.fn(() => session.transcript()) + const turnTwo = gate() + const reading = new AbortController() + const overrides: Partial = { + transcript, + // The stream sends a gap, then the backlog of the second turn. The + // backlog comes only after the transcript loaded again. + events: () => + (async function* () { + await turnTwo.opened + yield gapEvent(session.threadId, '0', before, 'evicted') + await waitFor(() => expect(transcript).toHaveBeenCalledTimes(2)) + await new Promise((resolve) => setTimeout(resolve, 20)) + yield* session.events({ from: before, signal: reading.signal }) + })(), + } + const source = new Proxy(session, { + get: (target, key) => { + if (key in overrides) return overrides[key as keyof SessionViewSource] + const value: unknown = Reflect.get(target, key, target) + return typeof value === 'function' ? value.bind(target) : value + }, + }) + const view = createSessionView(source) + await view.ready + await session.send('Two.') + turnTwo.open() + + const replies = () => + view.store + .get() + .messages.filter((message) => message.role === 'assistant') + await waitFor(() => expect(JSON.stringify(replies())).toContain('Second.')) + // The backlog events had time to apply. + await new Promise((resolve) => setTimeout(resolve, 50)) + expect(replies()).toHaveLength(2) + view.dispose() + reading.abort() + await session.close() + await host.close() + }) +}) diff --git a/packages/ai-harness/tests/goal.test.ts b/packages/ai-harness/tests/goal.test.ts index 5e2d4bc36c..89edbdfb5e 100644 --- a/packages/ai-harness/tests/goal.test.ts +++ b/packages/ai-harness/tests/goal.test.ts @@ -4,6 +4,7 @@ import { toolDefinition } from '@tanstack/ai' import { memoryPersistence } from '@tanstack/ai-persistence' import { createHarnessHost, defineHarness, definePlugin } from '../src' import { GoalMet, goal, selectGoal } from '../src/first-party' +import { loadLogState, logIdOf, sessionOf } from '../src/log' import { after, appendToLog, @@ -109,6 +110,9 @@ async function waitForStatus(session: HarnessSession, expected: string) { await waitFor(async () => expect(await status(session)).toBe(expected)) } +/** The goal prompt: a `'context'` prompt, the last system prompt. */ +const GOAL_PROMPT = 'You work toward this goal:' + /** The last message the main model got on `call`. */ function lastMessage(call: unknown) { return messageTexts(call).at(-1) @@ -142,9 +146,11 @@ describe('goal', () => { expect(lastMessage(calls[2])).toBe( 'Keep working on the goal: all tests pass. Last check: One test still fails.', ) - expect(JSON.stringify(calls[0].systemPrompts)).toContain( - 'You work toward this goal: all tests pass.', + // The goal prompt is the last system prompt, not a message. + expect(JSON.stringify(calls[0].systemPrompts.at(-1))).toContain( + `${GOAL_PROMPT} all tests pass.`, ) + expect(JSON.stringify(calls[0].messages)).not.toContain(GOAL_PROMPT) expect(read).toHaveLength(3) expect(read[0]).toContain('Goal: all tests pass') expect(read[0]).toContain('assistant: wrote the test') @@ -154,6 +160,36 @@ describe('goal', () => { await host.close() }) + it('counts each judge call under plugin:tanstack/goal', async () => { + const { adapter } = mockAdapter([]) + const judge: AnyTextAdapter = { + ...adapter, + structuredOutput: async () => ({ + data: { met: true, reason: 'Done.' }, + rawText: '{}', + usage: { promptTokens: 30, completionTokens: 5, totalTokens: 35 }, + }), + } + const { host, session } = await start({ + replies: [() => text('done')], + judge, + }) + + await session.command({ name: 'goal', input: 'finish' }) + await waitForStatus( + session, + 'Goal: finish\nStatus: met, round 1 of 20.\nLast check: Done.', + ) + + expect(session.usage().bySource).toEqual({ + 'plugin:tanstack/goal': expect.objectContaining({ + calls: 1, + totalTokens: 35, + }), + }) + await host.close() + }) + it('keeps going on a harness that rejects prompts while busy', async () => { const { host, session, calls } = await start({ replies: [() => text('try one'), () => text('done')], @@ -242,7 +278,7 @@ describe('goal', () => { await session.command({ name: 'goal', input: 'finish' }) await waitFor(() => expect(session.inputs()).toHaveLength(1)) const [waiting] = session.inputs() - if (waiting) await session.cancelInput(waiting.inputId) + if (waiting) await session.cancelInput({ inputId: waiting.inputId }) hold.open() await question await waitFor(() => expect(session.snapshot().status).toBe('idle')) @@ -483,3 +519,44 @@ describe('goal', () => { expect(selectGoal(state)).toBeNull() }) }) + +describe('goal with two users', () => { + it('runs each round as the user who set the goal, not as the owner of the thread', async () => { + const main = mockAdapter([ + () => text('Round one.'), + () => text('Round two.'), + ]) + const { judge } = scriptedJudge([ + { met: false, reason: 'Not yet.' }, + { met: true, reason: 'Done.' }, + ]) + const persistence = memoryPersistence() + const host = createHarnessHost({ persistence }) + const harness = defineHarness({ + name: 'test/goal', + adapter: main.adapter, + plugins: () => [goal({ judge })], + }) + // Olga opens the thread first, so she owns it. + await host.open(harness, { threadId: 't', principal: { id: 'olga' } }) + const bea = await host.open(harness, { + threadId: 't', + principal: { id: 'bea' }, + }) + + await bea.command({ name: 'goal', input: 'ship it' }) + await waitFor(() => expect(main.calls).toHaveLength(2)) + await bea.whenIdle() + + const state = await loadLogState({ + store: persistence.stores.log, + logId: logIdOf({ logId: 't' }), + rootThread: 't', + }) + const prompts = [...sessionOf(state, 't').inputs.values()] + .filter((logged) => logged.input.op === 'prompt') + .map((logged) => logged.principal?.id) + expect(prompts).toEqual(['bea', 'bea']) + await host.close() + }) +}) diff --git a/packages/ai-harness/tests/harness-text.test.ts b/packages/ai-harness/tests/harness-text.test.ts index 07daaa724c..3da21d9871 100644 --- a/packages/ai-harness/tests/harness-text.test.ts +++ b/packages/ai-harness/tests/harness-text.test.ts @@ -1,14 +1,93 @@ -import { describe, expect, it } from 'vitest' -import { EventType, chat } from '@tanstack/ai' +import { describe, expect, it, vi } from 'vitest' +import { EventType, chat, toolDefinition } from '@tanstack/ai' +import { resolveDebugOption } from '@tanstack/ai/adapter-internals' import { memoryPersistence } from '@tanstack/ai-persistence' import { createHarnessHost, defineHarness, definePlugin, + AccessDeniedError, harnessText, + membersOnly, } from '../src' -import { messageTexts, mockAdapter, text, untilAborted } from './helpers' -import type { StreamChunk } from '@tanstack/ai' +import { innerThreadOf, throughOuterThread } from '../src/host' +import { + messageTexts, + mockAdapter, + text, + toolCall, + untilAborted, +} from './helpers' +import type { AnyTextAdapter, StreamChunk } from '@tanstack/ai' +import type { RemoteHarness } from '../src' + +// ponytail: a copy of the helper in usage.test.ts. Move both to helpers.ts +// when a third file needs it. +/** `chunks` with this usage on RUN_FINISHED, as a provider reports it. */ +function withUsage( + chunks: Array, + usage: { prompt: number; completion: number; cost?: number }, +): Array { + return chunks.map((chunk) => + chunk.type === EventType.RUN_FINISHED + ? { + ...chunk, + usage: { + promptTokens: usage.prompt, + completionTokens: usage.completion, + totalTokens: usage.prompt + usage.completion, + ...(usage.cost !== undefined ? { cost: usage.cost } : {}), + }, + } + : chunk, + ) +} + +/** + * Run one call of `adapter` with the user message "hello", and return its + * chunks. `onChunk` sees each chunk as it comes. + */ +async function callChunks( + adapter: AnyTextAdapter, + options: { + threadId?: string + runId?: string + signal?: AbortSignal + onChunk?: (chunk: StreamChunk) => void + } = {}, +) { + const { signal, onChunk, ...ids } = options + const chunks: Array = [] + for await (const chunk of adapter.chatStream({ + model: adapter.model, + messages: [{ role: 'user', content: 'hello' }], + logger: resolveDebugOption(false), + ...ids, + ...(signal && { request: { signal } }), + })) { + chunks.push(chunk) + onChunk?.(chunk) + } + return chunks +} + +/** A remote harness whose `/run` streams the SSE text `parts`, one read each. */ +function remoteOf(parts: Array): RemoteHarness { + const encoder = new TextEncoder() + return { + url: 'http://remote.test', + fetch: async () => + new Response( + new ReadableStream({ + start(controller) { + for (const part of parts) controller.enqueue(encoder.encode(part)) + controller.close() + }, + }), + { status: 200 }, + ), + } +} describe('harnessText', () => { it('runs a harness as the model of chat() and keeps its own transcript', async () => { @@ -108,8 +187,11 @@ describe('harnessText', () => { persistence: memoryPersistence(), lifecycle: { idleCloseMs: 0 }, }) - const threadId = 'outer-kept:test/inner' - const held = await host.open(studio, { threadId }) + const threadId = innerThreadOf('outer-kept', studio.name) + const held = await host.open( + studio, + throughOuterThread({ threadId }, 'outer-kept'), + ) for await (const _chunk of chat({ adapter: harnessText(studio, { host }), @@ -130,15 +212,10 @@ describe('harnessText', () => { const host = createHarnessHost({ persistence: memoryPersistence() }) const adapter = harnessText(studio, { host }) - const chunks: Array = [] - for await (const chunk of adapter.chatStream({ - model: adapter.model, - messages: [{ role: 'user', content: 'hello' }], + const chunks = await callChunks(adapter, { threadId: 'outer-aborted', - request: { signal: AbortSignal.abort() }, - } as never)) { - chunks.push(chunk) - } + signal: AbortSignal.abort(), + }) expect(chunks.at(-1)?.type).toBe(EventType.RUN_ERROR) await host.close() @@ -151,19 +228,11 @@ describe('harnessText', () => { threadId: 'inner', outcome: { type: 'interrupt', interrupts: [] }, } - const adapter = harnessText({ - url: 'http://remote.test', - fetch: async () => - new Response(`data: ${JSON.stringify(finished)}\n`, { status: 200 }), - }) + const adapter = harnessText( + remoteOf([`data: ${JSON.stringify(finished)}\n`]), + ) - const chunks: Array = [] - for await (const chunk of adapter.chatStream({ - model: adapter.model, - messages: [{ role: 'user', content: 'hello' }], - } as never)) { - chunks.push(chunk) - } + const chunks = await callChunks(adapter) const last = chunks.at(-1) expect(last?.type).toBe(EventType.RUN_FINISHED) @@ -179,22 +248,11 @@ describe('harnessText', () => { messageId: 'm1', delta: 'partial', } - const adapter = harnessText({ - url: 'http://remote.test', - fetch: async () => - new Response(`data: ${JSON.stringify(delta)}\r\n\r\n`, { - status: 200, - }), - }) + const adapter = harnessText( + remoteOf([`data: ${JSON.stringify(delta)}\r\n\r\n`]), + ) - const chunks: Array = [] - for await (const chunk of adapter.chatStream({ - model: adapter.model, - messages: [{ role: 'user', content: 'hello' }], - threadId: 'outer-remote', - } as never)) { - chunks.push(chunk) - } + const chunks = await callChunks(adapter, { threadId: 'outer-remote' }) expect(chunks.map((chunk) => chunk.type)).toEqual([ EventType.RUN_STARTED, @@ -202,4 +260,245 @@ describe('harnessText', () => { EventType.RUN_ERROR, ]) }) + + it('reads SSE events with data on more lines, CR line ends, and a CRLF cut between reads', async () => { + const adapter = harnessText( + remoteOf([ + 'data: {"type":"TEXT_MESSAGE_CONTENT",\r', + '\ndata:"messageId":"m1","delta":"a"}\r\n', + ': a comment\r\n\r\n', + 'data: {"type":"RUN_FINISHED","runId":"i","threadId":"i",', + '"usage":[{"inputTokens":3,"outputTokens":2,"totalTokens":5}]}\r\r', + ]), + ) + + const chunks = await callChunks(adapter, { threadId: 'outer-remote' }) + + expect(chunks).toEqual([ + expect.objectContaining({ type: EventType.RUN_STARTED }), + { type: EventType.TEXT_MESSAGE_CONTENT, messageId: 'm1', delta: 'a' }, + expect.objectContaining({ + type: EventType.RUN_FINISHED, + usage: { promptTokens: 3, completionTokens: 2, totalTokens: 5 }, + metadata: { tanstack: { finishReason: 'stop' } }, + }), + ]) + }) + + it('cancels the remote turn when the call aborts', async () => { + const requests: Array<{ url: string; body: unknown }> = [] + const encoder = new TextEncoder() + const delta = { + type: EventType.TEXT_MESSAGE_CONTENT, + messageId: 'm1', + delta: 'working', + } + const adapter = harnessText({ + url: 'http://remote.test/', + fetch: async (url, init) => { + requests.push({ + url: String(url), + body: typeof init?.body === 'string' ? JSON.parse(init.body) : null, + }) + if (!String(url).endsWith('/run')) return Response.json({}) + // A turn that runs until the fetch aborts. + return new Response( + new ReadableStream({ + start(controller) { + controller.enqueue( + encoder.encode(`data: ${JSON.stringify(delta)}\n\n`), + ) + init?.signal?.addEventListener('abort', () => + controller.error(new Error('The fetch aborted.')), + ) + }, + }), + { status: 200 }, + ) + }, + }) + const abort = new AbortController() + + await expect( + callChunks(adapter, { + threadId: 'outer-remote', + runId: 'run-1', + signal: abort.signal, + onChunk: (chunk) => { + if (chunk.type === EventType.TEXT_MESSAGE_CONTENT) abort.abort() + }, + }), + ).rejects.toThrow('The fetch aborted.') + + await vi.waitFor(() => expect(requests).toHaveLength(2)) + expect(requests[1]).toEqual({ + url: 'http://remote.test/control', + body: { + threadId: 'outer-remote', + input: { op: 'cancel', operationId: 'run-1' }, + }, + }) + }) + + it('opens the inner session as the user of each outer harness turn', async () => { + const inner = mockAdapter([() => text('one'), () => text('two')]) + const seen: Array = [] + const watch = definePlugin({ + name: 'test/watch', + setup: (ctx) => { + seen.push(ctx.session.principal) + return {} + }, + }) + const studio = defineHarness({ + name: 'test/inner', + adapter: inner.adapter, + plugins: () => [watch], + }) + // The inner session closes after each call, so each call opens it. + const innerHost = createHarnessHost({ + persistence: memoryPersistence(), + lifecycle: { idleCloseMs: 0 }, + }) + const lead = defineHarness({ + name: 'test/lead', + adapter: harnessText(studio, { + host: innerHost, + principal: { id: 'fallback' }, + }), + }) + const host = createHarnessHost({ persistence: memoryPersistence() }) + const session = await host.open(lead, { + threadId: 'outer', + principal: { id: 'alice' }, + }) + + await session.send('hello') + await session.send('again', { principal: { id: 'bob' } }) + + expect(seen).toEqual([{ id: 'alice' }, { id: 'bob' }]) + await host.close() + await innerHost.close() + }) + + it('lets every member of a shared outer thread use the inner session, and no one else', async () => { + const inner = mockAdapter([() => text('one'), () => text('two')]) + const studio = defineHarness({ name: 'test/inner', adapter: inner.adapter }) + const host = createHarnessHost({ + persistence: memoryPersistence(), + access: { identify: () => null, authorize: membersOnly() }, + }) + // The inner session runs on the host of the outer thread. + const lead = defineHarness({ + name: 'test/lead', + adapter: harnessText(studio, { host }), + }) + const ada = { id: 'ada' } + const bob = { id: 'bob' } + const cid = { id: 'cid' } + const asAda = await host.open(lead, { threadId: 'a', principal: ada }) + await asAda.send('hello') + await host.sessions.share({ + threadId: 'a', + userId: 'bob', + role: 'member', + principal: ada, + }) + await host.open(lead, { threadId: 'c', principal: cid }) + const asBob = await host.open(lead, { threadId: 'a', principal: bob }) + + await expect(asBob.send('again')).resolves.toMatchObject({ text: 'two' }) + const innerOfA = innerThreadOf('a', studio.name) + const openInner = (outerThread?: string) => { + const openOptions = { threadId: innerOfA, principal: cid } + return host.open( + studio, + outerThread + ? throughOuterThread(openOptions, outerThread) + : openOptions, + ) + } + // Cid reads his own thread `c`, not `a`: no path opens the inner thread of `a`. + await expect(openInner()).rejects.toThrow(/harnessText inner session/) + await expect(openInner('c')).rejects.toThrow(/harnessText inner session/) + await expect(openInner('a')).rejects.toBeInstanceOf(AccessDeniedError) + await host.close() + }) + + it('checks an inner thread linked to another outer thread on its own members', async () => { + const inner = mockAdapter([() => text('one'), () => text('two')]) + const studio = defineHarness({ name: 'test/inner', adapter: inner.adapter }) + const host = createHarnessHost({ + persistence: memoryPersistence(), + access: { identify: () => null, authorize: membersOnly() }, + }) + const lead = defineHarness({ + name: 'test/lead', + adapter: harnessText(studio, { host }), + }) + const ada = { id: 'ada' } + const bob = { id: 'bob' } + const asAda = await host.open(lead, { threadId: 'a', principal: ada }) + await asAda.send('hello') + await host.sessions.share({ + threadId: 'a', + userId: 'bob', + role: 'member', + principal: ada, + }) + // The link of the inner entry names another outer thread now. + await host.sessions.update({ + harness: studio, + threadId: innerThreadOf('a', studio.name), + principal: ada, + update: (entry) => ({ + metadata: { ...entry.metadata, 'harness:outerThread': 'other' }, + }), + }) + const asBob = await host.open(lead, { threadId: 'a', principal: bob }) + + await expect(asBob.send('again')).rejects.toThrow( + 'Access denied: read on thread "a/inner/test/inner".', + ) + await host.close() + }) + + it('gives the outer run the usage of every inner model call', async () => { + const inner = mockAdapter([ + () => withUsage(toolCall('lookup', {}), { prompt: 3, completion: 2 }), + () => withUsage(text('found'), { prompt: 10, completion: 4, cost: 0.5 }), + ]) + const studio = defineHarness({ + name: 'test/inner', + adapter: inner.adapter, + tools: [ + toolDefinition({ + name: 'lookup', + description: 'Look up a fact', + inputSchema: { type: 'object', properties: {} }, + }).server(async () => 'fact'), + ], + }) + const innerHost = createHarnessHost({ persistence: memoryPersistence() }) + const lead = defineHarness({ + name: 'test/lead', + adapter: harnessText(studio, { host: innerHost }), + }) + const host = createHarnessHost({ persistence: memoryPersistence() }) + const session = await host.open(lead, { threadId: 'outer' }) + + expect((await session.send('look it up')).text).toBe('found') + + expect(session.usage().total).toEqual({ + calls: 1, + promptTokens: 13, + completionTokens: 6, + totalTokens: 19, + cachedTokens: 0, + cacheWriteTokens: 0, + cost: 0.5, + }) + await host.close() + await innerHost.close() + }) }) diff --git a/packages/ai-harness/tests/helpers.ts b/packages/ai-harness/tests/helpers.ts index e3b3d32437..284878fc3a 100644 --- a/packages/ai-harness/tests/helpers.ts +++ b/packages/ai-harness/tests/helpers.ts @@ -1,5 +1,6 @@ import { EventType } from '@tanstack/ai' import { loadLogState, logIdOf, sessionOf } from '../src/log' +import { UNCHECKED_TIMING } from '../src/ownership' import type { AnyChatMiddleware, AnyTextAdapter, @@ -9,6 +10,13 @@ import type { } from '@tanstack/ai' import type { LogRecord, LogStore, RunRecord } from '@tanstack/ai-persistence' import type { HostFrame } from '../src/protocol' +import type { + AttributeValue, + Exception, + Span, + SpanContext, + SpanStatus, +} from '@opentelemetry/api' /** * A log store that calls `overrides` for the methods it has, and `inner` for @@ -94,7 +102,11 @@ export async function storedRun( * 200 ms and never renews it in a test, like a process that stopped. After * `ownershipExpired()`, another host can take its logs and recover its work. */ -export const CRASHING_OWNERSHIP = { ttlMs: 200, renewMs: 60_000 } +export const CRASHING_OWNERSHIP = { + ttlMs: 200, + renewMs: 60_000, + [UNCHECKED_TIMING]: true, +} /** Wait until an ownership of {@link CRASHING_OWNERSHIP} expired. */ export const ownershipExpired = () => @@ -421,3 +433,35 @@ export const hideSecret: AnyChatMiddleware = { } }, } + +/** An OpenTelemetry span that calls `on` for the calls that a test reads. */ +export function fakeSpan( + spanContext: SpanContext, + on: { + setAttribute?: (key: string, value: AttributeValue) => void + setStatus?: (status: SpanStatus) => void + recordException?: (exception: Exception) => void + end?: () => void + } = {}, +): Span { + const span: Span = { + spanContext: () => spanContext, + setAttribute: (key, value) => { + on.setAttribute?.(key, value) + return span + }, + setAttributes: () => span, + addEvent: () => span, + addLink: () => span, + addLinks: () => span, + setStatus: (status) => { + on.setStatus?.(status) + return span + }, + updateName: () => span, + end: () => on.end?.(), + isRecording: () => true, + recordException: (exception) => on.recordException?.(exception), + } + return span +} diff --git a/packages/ai-harness/tests/host-event-bus.test.ts b/packages/ai-harness/tests/host-event-bus.test.ts new file mode 100644 index 0000000000..ac14053091 --- /dev/null +++ b/packages/ai-harness/tests/host-event-bus.test.ts @@ -0,0 +1,120 @@ +import { describe, expect, it } from 'vitest' +import { HostEventBus } from '../src/host-event-bus' +import type { HostEvent } from '../src/host' + +/** A status event of `threadId` in `tenantId`. */ +const status = (threadId: string, tenantId?: string): HostEvent => ({ + type: 'status', + threadId, + status: 'running', + at: 1, + ...(tenantId !== undefined ? { tenantId } : {}), +}) + +/** The thread of a status event. */ +const threadOf = (event: HostEvent) => + 'threadId' in event ? event.threadId : undefined + +/** Take `count` events of `read`, then stop it. */ +async function take(read: AsyncIterable, count: number) { + const seen: Array = [] + if (count === 0) return seen + for await (const event of read) { + seen.push(event) + if (seen.length === count) break + } + return seen +} + +describe('the host event bus', () => { + it('gives a tenant reader the events of its tenant only, and an all-tenants reader every event', async () => { + const bus = new HostEventBus() + const ofA = take(bus.read({ tenant: { tenantId: 'a' } }), 1) + const ofDefault = take(bus.read({ tenant: { tenantId: undefined } }), 1) + const ofAll = take(bus.read(), 3) + await Promise.resolve() + + bus.emit(status('t1', 'b')) + bus.emit(status('t2', 'a')) + bus.emit(status('t3')) + + expect((await ofA).map(threadOf)).toEqual(['t2']) + expect((await ofDefault).map(threadOf)).toEqual(['t3']) + expect((await ofAll).map(threadOf)).toEqual(['t1', 't2', 't3']) + expect(bus.readers).toBe(0) + }) + + it('filters before the queue, and starts with the initial events', async () => { + const bus = new HostEventBus() + const asked: Array = [] + const read = take( + bus.read({ + initial: [status('t0')], + filter: (event) => { + const thread = threadOf(event) ?? '' + asked.push(thread) + return thread !== 'skip' + }, + }), + 2, + ) + await Promise.resolve() + bus.emit(status('skip')) + bus.emit(status('t1')) + + expect((await read).map(threadOf)).toEqual(['t0', 't1']) + // The filter runs once per emitted event, not for the initial ones. + expect(asked).toEqual(['skip', 't1']) + }) + + it('ends a reader that falls more than maxQueue events behind, after the events it kept', async () => { + const bus = new HostEventBus() + const iterator = bus.read({ maxQueue: 2 })[Symbol.asyncIterator]() + const first = iterator.next() + await Promise.resolve() + bus.emit(status('t1')) + expect((await first).value?.threadId).toBe('t1') + + // Nobody reads now: two events wait, the third does not fit. + for (const id of ['t2', 't3', 't4']) bus.emit(status(id)) + + expect((await iterator.next()).value?.threadId).toBe('t2') + expect((await iterator.next()).value?.threadId).toBe('t3') + expect(await iterator.next()).toEqual({ done: true, value: undefined }) + expect(bus.readers).toBe(0) + }) + + it('does not count the initial events against maxQueue', async () => { + const bus = new HostEventBus() + const read = bus.read({ + initial: [status('t0'), status('t1')], + maxQueue: 1, + }) + bus.emit(status('t2')) + + expect((await take(read, 3)).map(threadOf)).toEqual(['t0', 't1', 't2']) + }) + + it('keeps an event emitted after the read starts but before the first pull', async () => { + const bus = new HostEventBus() + const read = bus.read({ initial: [status('t0')] }) + bus.emit(status('t1')) + + expect((await take(read, 2)).map(threadOf)).toEqual(['t0', 't1']) + expect(bus.readers).toBe(0) + }) + + it('stops a read when its signal aborts', async () => { + const bus = new HostEventBus() + const controller = new AbortController() + const next = bus + .read({ signal: controller.signal }) + [Symbol.asyncIterator]() + .next() + await Promise.resolve() + controller.abort() + + expect(await next).toEqual({ done: true, value: undefined }) + expect(bus.readers).toBe(0) + }) +}) diff --git a/packages/ai-harness/tests/http-access.test.ts b/packages/ai-harness/tests/http-access.test.ts index 93cc875cc7..b9598bffc0 100644 --- a/packages/ai-harness/tests/http-access.test.ts +++ b/packages/ai-harness/tests/http-access.test.ts @@ -7,6 +7,12 @@ import { isMediaRecord } from '../src/media-entry' import { fakeSocket, mockAdapter, readSse, text } from './helpers' import type { HostEvent } from '../src' +// Coverage runs in CI are slow, so polled checks wait up to 5 s. +const waitFor = ( + check: () => T | Promise, + options?: { timeout?: number; interval?: number }, +) => vi.waitFor(check, { timeout: 5_000, ...options }) + const BASE = 'http://x/api/harness' const olga = { id: 'olga', tenantId: 'acme' } @@ -196,6 +202,50 @@ describe('the handler on a host with access', () => { await host.close() }) + it('watches revokes on the bus of the tenant, not on the every-tenant host.events()', async () => { + const { host, call, prompt, share } = setup() + await prompt('olga', 't1') + await share('bea', 'member') + // The every-tenant reader starts with all statuses and gets every event + // of every tenant, so a watch of each stream must not use it. + const everyTenant = vi.spyOn(host, 'events') + const response = await call('bea', 'events?threadId=t1') + const reading = readSse(response, () => false) + + host.revoke({ threadId: 't1', principal: bea }) + + await expect(reading).resolves.toBeInstanceOf(Array) + expect(everyTenant).not.toHaveBeenCalled() + await host.close() + }) + + it('reads the access of a readable thread once, not on each index change', async () => { + const { host, call, prompt } = setup() + await prompt('olga', 't1') + const get = vi.spyOn(host.sessions, 'get') + const renames: Array> = [] + const reading = readSse( + await call('olga', 'host-events'), + (frame: { type: string } | HostEvent) => { + // The hello: the stream reads the bus now. + if (frame.type === 'harness.hello') { + for (const title of ['a', 'b', 'c']) { + renames.push( + host.sessions.rename({ threadId: 't1', title, principal: olga }), + ) + } + } + return 'entry' in frame && frame.entry.title === 'c' + }, + ) + + await reading + await Promise.all(renames) + const reads = get.mock.calls.filter(([read]) => read.threadId === 't1') + expect(reads).toHaveLength(1) + await host.close() + }) + it('never shows the threads of another tenant in host-events', async () => { const { host, call, prompt } = setup() const reading = readSse( @@ -280,13 +330,23 @@ describe('the socket on a host with access', () => { nobody.receive({ type: 'harness.subscribe', threadId: 't1' }) stranger.receive({ type: 'harness.subscribe', threadId: 't1' }) - await vi.waitFor(() => expect(nobody.closed).toEqual([4401])) + await waitFor(() => expect(nobody.closed).toEqual([4401])) expect(nobody.sent).toEqual([ - { type: 'harness.error', message: 'unauthorized' }, + { + type: 'harness.error', + message: 'unauthorized', + code: 'unauthorized', + requestId: expect.any(String), + }, ]) - await vi.waitFor(() => expect(stranger.closed).toEqual([4403])) + await waitFor(() => expect(stranger.closed).toEqual([4403])) expect(stranger.sent).toEqual([ - { type: 'harness.error', message: 'forbidden' }, + { + type: 'harness.error', + message: 'forbidden', + code: 'forbidden', + requestId: expect.any(String), + }, ]) await host.close() }) @@ -298,14 +358,19 @@ describe('the socket on a host with access', () => { const { socket, sent, closed, receive } = fakeSocket() handleHarnessSocket({ host, harness, socket, request: requestAs('bea') }) receive({ type: 'harness.subscribe', threadId: 't1' }) - await vi.waitFor(() => + await waitFor(() => expect(sent.some((frame) => frame.type === 'harness.hello')).toBe(true), ) host.revoke({ threadId: 't1', principal: bea }) - await vi.waitFor(() => expect(closed).toEqual([4403])) - expect(sent.at(-1)).toEqual({ type: 'harness.error', message: 'forbidden' }) + await waitFor(() => expect(closed).toEqual([4403])) + expect(sent.at(-1)).toEqual({ + type: 'harness.error', + message: 'forbidden', + code: 'forbidden', + requestId: expect.any(String), + }) await host.close() }) }) diff --git a/packages/ai-harness/tests/http-media.test.ts b/packages/ai-harness/tests/http-media.test.ts index 96e9f5e12d..c5d6d674c0 100644 --- a/packages/ai-harness/tests/http-media.test.ts +++ b/packages/ai-harness/tests/http-media.test.ts @@ -14,6 +14,10 @@ const NOW = 1_000_000 const user = { id: 'user-1' } +/** Media signing keys: each one is 32 bytes or more. */ +const SHARED = 'a-shared-secret-of-32-bytes-or-more' +const OLD_KEY = 'the-old-secret-of-32-bytes-or-more' + const allowUser = (request: Request) => request.headers.get('authorization') === 'Bearer good' ? user : null @@ -41,16 +45,25 @@ function setup(options: { media?: MediaOptions; secret?: string } = {}) { }, }) const host = hostWith(allowUser) - /** `secret` signs the media URLs. Without it, a random key per handler. */ + /** + * `secret` or `keys` sign the media URLs. Without them, a random key per + * handler. + */ const handlerWith = ( - handlerOptions: { host?: typeof host; secret?: string } = {}, + handlerOptions: { + host?: typeof host + secret?: string + keys?: Array + } = {}, ) => createHarnessHandler({ host: handlerOptions.host ?? host, harness, ...(handlerOptions.secret ? { media: { secret: handlerOptions.secret } } - : {}), + : handlerOptions.keys + ? { media: { keys: handlerOptions.keys } } + : {}), }) const handler = handlerWith({ secret: options.secret }) const call = (path: string, init: RequestInit = {}) => @@ -150,7 +163,7 @@ describe('media upload and signed URLs', () => { openAccess: true, media: { secret: '' }, }), - ).toThrow(/not empty/) + ).toThrow(/at least 32 bytes/) }) it('stores an upload, signs its URL for 1 hour, and serves the signed URL without auth', async () => { @@ -184,11 +197,11 @@ describe('media upload and signed URLs', () => { it('serves a signed URL on another instance whose identify refuses everyone', async () => { const { host, upload, signedPath, fetchSigned, hostWith, handlerWith } = - setup({ secret: 'shared' }) + setup({ secret: SHARED }) const record = await upload('user-1-a') const { path } = await signedPath(record) const other = hostWith(() => null) - const refuseAll = handlerWith({ host: other, secret: 'shared' }) + const refuseAll = handlerWith({ host: other, secret: SHARED }) const signed = await fetchSigned(path, refuseAll) const unsigned = await fetchSigned( @@ -262,6 +275,40 @@ describe('media upload and signed URLs', () => { await host.close() }) + it('serves a URL that an old key signed on a handler with keys [new, old]', async () => { + const { host, upload, signedPath, fetchSigned, handlerWith } = setup({ + secret: OLD_KEY, + }) + const record = await upload('user-1-a') + const { path } = await signedPath(record) + + const rotated = handlerWith({ keys: [SHARED, OLD_KEY] }) + const newOnly = handlerWith({ keys: [SHARED] }) + + expect((await fetchSigned(path, rotated)).status).toBe(200) + expect((await fetchSigned(path, newOnly)).status).toBe(403) + await host.close() + }) + + it('throws at the start for a signing key under 32 bytes', () => { + const { host, handlerWith } = setup() + + expect(() => handlerWith({ secret: 'shared' })).toThrow(/at least 32/) + void host.close() + }) + + it('refuses a signed URL of a file whose thread was deleted (404)', async () => { + const { host, upload, signedPath, fetchSigned } = setup() + const record = await upload('user-1-a') + const { path } = await signedPath(record) + + await host.sessions.delete({ threadId: 'user-1-a', principal: user }) + const response = await fetchSigned(path) + + expect(response.status).toBe(404) + await host.close() + }) + it('does not sign or serve the file of another thread (404)', async () => { const { host, call, upload } = setup() const record = await upload('user-1-a') @@ -343,6 +390,32 @@ describe('media bytes', () => { await host.close() }) + it('sends at most 8 MB for a Range with no end', async () => { + const { host, call, upload } = setup() + const cap = 8 * 1024 * 1024 + const record = await upload('user-1-a', 'x'.repeat(cap + 10)) + + const capped = await call(`media?threadId=user-1-a&id=${record.id}`, { + headers: { range: 'bytes=2-' }, + }) + expect(capped.status).toBe(206) + expect(capped.headers.get('content-range')).toBe( + `bytes 2-${cap + 1}/${cap + 10}`, + ) + expect(capped.headers.get('content-length')).toBe(String(cap)) + expect((await capped.arrayBuffer()).byteLength).toBe(cap) + + // The rest is under the cap, so it comes whole. + const rest = await call(`media?threadId=user-1-a&id=${record.id}`, { + headers: { range: `bytes=${cap + 2}-` }, + }) + expect(rest.headers.get('content-range')).toBe( + `bytes ${cap + 2}-${cap + 9}/${cap + 10}`, + ) + expect(await rest.text()).toBe('xxxxxxxx') + await host.close() + }) + it('answers a Range past the end with 416', async () => { const { host, call, upload } = setup() const record = await upload('user-1-a') @@ -370,6 +443,7 @@ describe('media bytes', () => { 'content-security-policy': "default-src 'none'; sandbox", 'content-type': 'text/html', 'x-content-type-options': 'nosniff', + 'x-request-id': expect.any(String), }) await host.close() }) @@ -388,12 +462,12 @@ describe('POST run with media parts', () => { 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 }, + { label: 'a harness-media URL', part: mediaPart }, + { label: 'a data part', part: () => data }, ] it.each(cases)( - 'gives the model the bytes of $label in the transcript', + 'gives the model the bytes of $label, and keeps a harness-media URL in the transcript', async ({ part }) => { const { host, call, upload, calls } = setup() const record = await upload('user-1-run') @@ -412,11 +486,110 @@ describe('POST run with media parts', () => { { type: 'text', content: 'What is this?' }, data, ]) + // A data part moved to the media store. Its URL names a new record. expect(transcript[0].content).toEqual([ { type: 'text', content: 'What is this?' }, - part(record), + { + type: 'image', + source: { + type: 'url', + value: expect.stringMatching(/^harness-media:/), + mimeType: 'image/png', + }, + }, ]) await host.close() }, ) + + const refused: Array<{ + label: string + media?: MediaOptions + part: ContentPart + status: number + error: string + }> = [ + { + label: 'bytes over media.maxBytes', + media: { maxBytes: 4 }, + part: data, + status: 413, + error: 'media_too_large', + }, + { + label: 'a kind not in media.kinds', + media: { kinds: ['audio'] }, + part: data, + status: 415, + error: 'media_not_allowed', + }, + { + label: 'an external URL', + part: { + type: 'image', + source: { type: 'url', value: 'https://example.com/cat.png' }, + }, + status: 400, + error: 'media_url_not_allowed', + }, + ] + + it.each(refused)( + 'refuses $label with $status, and starts no turn', + async ({ media, part, status, error }) => { + const { host, call, calls } = setup(media ? { media } : {}) + + const response = await call('run', { + method: 'POST', + body: runBody('user-1-run', [question, part]), + headers: { 'content-type': 'application/json' }, + }) + + expect(response.status).toBe(status) + expect(await response.json()).toEqual({ error }) + expect(calls).toHaveLength(0) + await host.close() + }, + ) +}) + +describe('request bodies over media.maxBodyBytes', () => { + const routes = [ + { + route: 'run', + body: runBody('user-1-run', [{ type: 'text', text: 'hi' }]), + }, + { + route: 'control', + body: JSON.stringify({ + threadId: 'user-1-run', + input: { op: 'prompt', message: 'hi' }, + }), + }, + { + route: 'sessions', + body: JSON.stringify({ + op: 'rename', + threadId: 'user-1-run', + title: 'x', + }), + }, + ] + + it.each(routes)('answers POST $route with 413', async ({ route, body }) => { + const { host, call, calls } = setup({ media: { maxBodyBytes: 64 } }) + + const response = await call(route, { + method: 'POST', + body: body.padEnd(65, ' '), + headers: { 'content-type': 'application/json' }, + }) + + expect(response.status).toBe(413) + expect(await response.json()).toEqual({ + error: 'The body is bigger than the limit of 64 bytes.', + }) + expect(calls).toHaveLength(0) + await host.close() + }) }) diff --git a/packages/ai-harness/tests/idle-budget.test.ts b/packages/ai-harness/tests/idle-budget.test.ts new file mode 100644 index 0000000000..5fca9fde0f --- /dev/null +++ b/packages/ai-harness/tests/idle-budget.test.ts @@ -0,0 +1,76 @@ +import { afterEach, describe, expect, it, vi } from 'vitest' +import { memoryPersistence } from '@tanstack/ai-persistence' +import { createHarnessHost, defineHarness } from '../src' +import { mockAdapter, text } from './helpers' + +const waitFor = (check: () => T | Promise) => + vi.waitFor(check, { timeout: 5_000 }) + +/** Longer than every timer of a host: renew, snapshot, idle close. */ +const IDLE_MS = 10 * 60_000 + +/** + * Memory stores that write the name of each method call to `calls`, as + * `store.method`. + */ +function countedPersistence() { + const persistence = memoryPersistence() + const calls: Array = [] + const counted = (name: string, store: object) => + new Proxy(store, { + get(target, key, receiver) { + const value: unknown = Reflect.get(target, key, receiver) + if (typeof value !== 'function') return value + return (...args: Array) => { + calls.push(`${name}.${String(key)}`) + return value.apply(target, args) + } + }, + }) + const stores = Object.fromEntries( + Object.entries(persistence.stores).map(([name, store]) => [ + name, + typeof store === 'object' ? counted(name, store) : store, + ]), + ) as typeof persistence.stores + return { persistence: { ...persistence, stores }, calls } +} + +afterEach(() => { + vi.useRealTimers() +}) + +describe('an idle session', () => { + it('makes no store calls after its turn, while a handle holds it, and after it closes', async () => { + // Fake timers that also move with real time, so the turn runs as usual + // and the idle time below fires every timer of the host. + vi.useFakeTimers({ shouldAdvanceTime: true }) + const { persistence, calls } = countedPersistence() + const { adapter } = mockAdapter(() => text('Hello back.')) + const harness = defineHarness({ name: 'test/idle-budget', adapter }) + const host = createHarnessHost({ + persistence, + lifecycle: { idleCloseMs: 1_000 }, + }) + const session = await host.open(harness, { threadId: 't1' }) + await session.send('hi') + // The session gives its log back when its work ends. + await waitFor(() => expect(calls).toContain('log.release')) + await vi.advanceTimersByTimeAsync(100) + // The spy sees the writes of the turn. + expect(calls).toContain('log.append') + + const afterTurn = calls.length + await vi.advanceTimersByTimeAsync(IDLE_MS) + expect(calls.slice(afterTurn)).toEqual([]) + expect(host.isOpen(harness, { threadId: 't1' })).toBe(true) + + await session.close() + await vi.advanceTimersByTimeAsync(2_000) + expect(host.isOpen(harness, { threadId: 't1' })).toBe(false) + const afterClose = calls.length + await vi.advanceTimersByTimeAsync(IDLE_MS) + expect(calls.slice(afterClose)).toEqual([]) + await host.close() + }) +}) diff --git a/packages/ai-harness/tests/inline-media.test.ts b/packages/ai-harness/tests/inline-media.test.ts new file mode 100644 index 0000000000..7dcc500024 --- /dev/null +++ b/packages/ai-harness/tests/inline-media.test.ts @@ -0,0 +1,266 @@ +import { describe, expect, it } from 'vitest' +import { memoryPersistence } from '@tanstack/ai-persistence' +import { createHarnessHost, defineHarness } from '../src' +import { readBodyText, storeInlineMedia } from '../src/inline-media' +import { MediaError, createMediaStore } from '../src/media' +import { mediaIdOf } from '../src/media-entry' +import { applyInput } from '../src/server' +import { mockAdapter, text } from './helpers' +import type { ContentPart } from '@tanstack/ai' +import type { MediaOptions } from '../src/media' + +// 'aGVsbG8=' is base64 for "hello", 5 bytes. +const HELLO = 'aGVsbG8=' +const decoder = new TextDecoder() + +const inlineImage = (value = HELLO, mimeType = 'image/png') => { + const part: ContentPart = { + type: 'image', + source: { type: 'data', value, mimeType }, + } + return part +} + +const urlImage = (value: string) => { + const part: ContentPart = { type: 'image', source: { type: 'url', value } } + return part +} + +/** A media store of one thread, and `storeInlineMedia` on it. */ +function storeWith(options?: MediaOptions) { + const store = createMediaStore({ + persistence: memoryPersistence(), + threadId: 't1', + ...(options ? { options } : {}), + }) + const check = (parts: Array) => + storeInlineMedia(parts, { + options, + putMedia: ({ body, mimeType, name }) => + store.put(body, { mimeType, name }), + }) + return { store, check } +} + +/** The text of the bytes that the media part of `checked` points to. */ +async function storedText( + store: ReturnType['store'], + checked: Awaited['check']>>, +) { + if ('refused' in checked) throw new Error(`refused: ${checked.refused}`) + const id = checked.parts.map(mediaIdOf).find((item) => item !== undefined) + if (id === undefined) throw new Error('not a media reference') + return decoder.decode(await store.load(id)) +} + +describe('storeInlineMedia', () => { + it('moves an inline image into the store and sends a reference', async () => { + const { store, check } = storeWith() + + const checked = await check([ + { type: 'text', content: 'What is this?' }, + inlineImage(), + ]) + + expect(checked).toEqual({ + parts: [ + { type: 'text', content: 'What is this?' }, + { + type: 'image', + source: { + type: 'url', + value: expect.stringMatching(/^harness-media:./), + mimeType: 'image/png', + }, + }, + ], + }) + expect(await storedText(store, checked)).toBe('hello') + }) + + it('moves a base64 data URL into the store too', async () => { + const { store, check } = storeWith() + + const checked = await check([urlImage(`data:image/png;base64,${HELLO}`)]) + + expect(await storedText(store, checked)).toBe('hello') + }) + + it('refuses inline bytes over media.maxBytes', async () => { + const { check } = storeWith({ maxBytes: 4 }) + + expect(await check([inlineImage()])).toEqual({ refused: 'media_too_large' }) + }) + + it('refuses a kind that is not in media.kinds', async () => { + const { check } = storeWith({ kinds: ['document'] }) + + expect(await check([inlineImage()])).toEqual({ + refused: 'media_not_allowed', + }) + }) + + it('refuses data that is not base64', async () => { + const { check } = storeWith() + + expect(await check([inlineImage('not base64!')])).toEqual({ + refused: 'media_invalid', + }) + }) + + it('refuses an external URL by default, and keeps it when allowed', async () => { + const cat = urlImage('https://example.com/cat.png') + + expect(await storeWith().check([cat])).toEqual({ + refused: 'media_url_not_allowed', + }) + expect(await storeWith({ externalUrls: true }).check([cat])).toEqual({ + parts: [cat], + }) + const onlyCdn = storeWith({ + externalUrls: (url) => url.hostname === 'cdn.example.com', + }) + expect(await onlyCdn.check([cat])).toEqual({ + refused: 'media_url_not_allowed', + }) + }) + + it('refuses a URL scheme other than http and https, even when URLs are allowed', async () => { + const { check } = storeWith({ externalUrls: true }) + + expect(await check([urlImage('file:///etc/passwd')])).toEqual({ + refused: 'media_url_not_allowed', + }) + }) + + it('stores nothing when a later part is refused', async () => { + const stored: Array = [] + const checked = await storeInlineMedia( + [inlineImage(), urlImage('https://example.com/cat.png')], + { + options: undefined, + putMedia: async (upload) => { + stored.push(upload.mimeType) + throw new Error('not reached') + }, + }, + ) + + expect(checked).toEqual({ refused: 'media_url_not_allowed' }) + expect(stored).toEqual([]) + }) + + it('refuses a provider file handle unless media.allowProviderFiles is set', async () => { + const file: ContentPart = { + type: 'document', + source: { type: 'file', value: 'file-abc' }, + } + + expect(await storeWith().check([file])).toEqual({ + refused: 'media_url_not_allowed', + }) + expect(await storeWith({ allowProviderFiles: true }).check([file])).toEqual( + { parts: [file] }, + ) + }) +}) + +describe('readBodyText', () => { + const post = (body: string) => + new Request('http://x/control', { method: 'POST', body }) + + it('reads a body up to media.maxBodyBytes', async () => { + expect(await readBodyText(post('12345'), { maxBodyBytes: 5 })).toBe('12345') + }) + + it('refuses a body over media.maxBodyBytes with 413', async () => { + const reading = readBodyText(post('123456'), { maxBodyBytes: 5 }) + + await expect(reading).rejects.toBeInstanceOf(MediaError) + await expect(reading).rejects.toMatchObject({ status: 413 }) + }) +}) + +describe('applyInput with media from a client', () => { + async function open(media?: MediaOptions) { + const { adapter, calls } = mockAdapter(() => text('ok')) + const harness = defineHarness({ name: 'test/inline', adapter, media }) + const host = createHarnessHost({ persistence: memoryPersistence() }) + const session = await host.open(harness, { threadId: 't1' }) + return { harness, host, session, calls } + } + + it('stores an inline image, keeps a reference in the transcript, and gives the model the bytes', async () => { + const { harness, host, session, calls } = await open() + + const receipt = await applyInput({ + harness, + session, + input: { op: 'prompt', inputId: 'in-1', message: [inlineImage()] }, + }) + await session.settled({ inputId: 'in-1' }) + + expect(receipt.status).toBe('accepted') + const [saved] = await session.transcript() + const id = Array.isArray(saved?.content) + ? mediaIdOf(saved.content[0]) + : undefined + if (id === undefined) throw new Error('no media reference in transcript') + expect(decoder.decode(await session.loadMedia({ id }))).toBe('hello') + expect(calls[0].messages[0].content).toEqual([ + { + type: 'image', + source: { type: 'data', value: HELLO, mimeType: 'image/png' }, + }, + ]) + await host.close() + }) + + it('treats a retried prompt with inline media as the same input', async () => { + const { harness, host, session, calls } = await open() + const input = { + op: 'prompt' as const, + inputId: 'in-1', + message: [inlineImage()], + } + + const first = await applyInput({ harness, session, input }) + await session.settled({ inputId: 'in-1' }) + const again = await applyInput({ harness, session, input }) + + // Stored again, the media id would differ, and the retry would be a + // conflict. + expect(first.status).toBe('accepted') + expect(again).toEqual(first) + expect(calls).toHaveLength(1) + const other = { ...input, message: [inlineImage(), inlineImage()] } + expect(await applyInput({ harness, session, input: other })).toMatchObject({ + status: 'rejected', + reason: 'conflict', + }) + await host.close() + }) + + it('rejects a prompt with an external URL, and no turn runs', async () => { + const { harness, host, session, calls } = await open() + + const receipt = await applyInput({ + harness, + session, + input: { + op: 'prompt', + inputId: 'in-1', + message: [urlImage('https://example.com/cat.png')], + }, + }) + + expect(receipt).toEqual({ + inputId: 'in-1', + status: 'rejected', + reason: 'media_url_not_allowed', + }) + expect(calls).toEqual([]) + expect(await session.transcript()).toEqual([]) + await host.close() + }) +}) diff --git a/packages/ai-harness/tests/input-lane.test.ts b/packages/ai-harness/tests/input-lane.test.ts index 4f45ea154d..9136923281 100644 --- a/packages/ai-harness/tests/input-lane.test.ts +++ b/packages/ai-harness/tests/input-lane.test.ts @@ -5,12 +5,14 @@ import { memoryLogBackend, memoryLogStore } from '@tanstack/ai-persistence' import { memoryCredentialStore } from '@tanstack/ai/byok/server' import { credentialsFor } from '../src/auth' import { + InputRejectedError, createHarnessHost, defineCommand, defineHarness, definePlugin, logIdOf, } from '../src' +import { createHarnessHandler } from '../src/server' import { CRASHING_OWNERSHIP, after, @@ -30,7 +32,9 @@ import type { AnyHarness, HarnessHostOptions, HarnessSession, + Receipt, } from '../src' +import type { SessionImpl } from '../src/session/harness-session' // Coverage runs in CI are slow, so polled checks wait up to 5 s. const waitFor = ( @@ -173,6 +177,91 @@ describe('an input that reaches a host that does not own the log', () => { await other.host.close() }) + it('ends settled() with InputRejectedError when the owner refuses the offered input', async () => { + const { a, b } = twoStores() + const release = gate() + const owner = await hostOn( + a, + harnessWith(mockAdapter([after(release.opened, 'One.')]).adapter), + ) + const other = await hostOn(b, harnessWith(mockAdapter([]).adapter)) + + const turn = owner.session.send('Go.') + await turn.receipt + const refused = other.session.send('Now!', { + inputId: 'in-b', + busy: 'reject', + }) + refused.then( + () => {}, + () => {}, + ) + expect((await refused.receipt).status).toBe('queued') + await expect( + other.session.settled({ + inputId: 'in-b', + signal: AbortSignal.timeout(3_000), + }), + ).rejects.toBeInstanceOf(InputRejectedError) + + release.open() + await turn + await owner.host.close() + await other.host.close() + }) + + it('refuses an offer whose id the log has with another payload, on the host that offered it', async () => { + const { a, b } = twoStores() + // The owner does not read the lane until the test lets it. + let isLaneOpen = false + const slowLane = wrapLogStore(a, { + offered: (scope) => (isLaneOpen ? a.offered(scope) : Promise.resolve([])), + }) + const release = gate() + const ownerModel = mockAdapter([ + after(release.opened, 'One.'), + () => text('Two.'), + ]) + const owner = await hostOn(slowLane, harnessWith(ownerModel.adapter)) + const other = await hostOn(b, harnessWith(mockAdapter([]).adapter)) + + const turn = owner.session.send('Go.') + await turn.receipt + const offered = other.session.send('From B.', { inputId: 'in-x' }) + offered.then( + () => {}, + () => {}, + ) + expect((await offered.receipt).status).toBe('queued') + const settled = other.session.settled({ + inputId: 'in-x', + signal: AbortSignal.timeout(3_000), + }) + settled.catch(() => {}) + // The owner stores another input with the same id first. + const stored = owner.session.send('From A.', { inputId: 'in-x' }) + expect((await stored.receipt).status).toBe('queued') + isLaneOpen = true + await owner.session.recover() + + await expect(settled).rejects.toMatchObject({ + name: 'InputRejectedError', + receipt: { reason: 'conflict' }, + }) + await expect(Promise.resolve(offered)).rejects.toBeInstanceOf( + InputRejectedError, + ) + expect(await waiting(a)).toEqual([]) + // The input that the log has is not refused: it runs. + release.open() + expect((await owner.session.settled({ inputId: 'in-x' })).outcome).toBe( + 'completed', + ) + expect(ownerModel.calls).toHaveLength(2) + await owner.host.close() + await other.host.close() + }) + it('streams an offered turn under its run id on the host that sent it', async () => { const { a, b } = twoStores() const release = gate() @@ -186,7 +275,10 @@ describe('an input that reaches a host that does not own the log', () => { const first = owner.session.send('Go.') await first.receipt // A `POST run` on the other host: the turn runs on the owner as `run-2`. - const second = other.session.send('And then?', { runId: 'run-2' }) + const second = (other.session as unknown as SessionImpl).sendForRun( + 'And then?', + { runId: 'run-2' }, + ) expect(await second.receipt).toMatchObject({ status: 'queued' }) release.open() const texts: Array = [] @@ -265,6 +357,26 @@ describe('an input that reaches a host that does not own the log', () => { return { credentials, saved } } + /** + * An answer as the owner applies it from the lane: with the `ref` of a + * secret answer, and the input id of the offer. + */ + function laneAnswer( + session: HarnessSession, + options: { + questionId: string + value: unknown + ref: string + controlId: string + }, + ): Promise { + // `applyAnswer` is internal: the input lane calls it. + const internal = session as unknown as { + applyAnswer: (offered: typeof options) => Promise + } + return internal.applyAnswer(options) + } + /** The owner asks for the password. Resolves with its question id. */ async function ask(session: HarnessSession, log: LogStore) { const running = session.command({ name: 'ask' }) @@ -329,11 +441,11 @@ describe('an input that reaches a host that does not own the log', () => { // The same ref again, as the owner applies it from the lane, is // refused: the store no longer has it. - const replay = await owner.session.answer({ + const replay = await laneAnswer(owner.session, { questionId, value: '[secret]', - ref, - inputId: 'replay', + ref: ref!, + controlId: 'replay', }) expect(replay).toMatchObject({ status: 'rejected', @@ -344,11 +456,11 @@ describe('an input that reaches a host that does not own the log', () => { const scope = { userId: 'someone-else' } const credential = { type: 'api_key' as const, value: '"theirs"' } await credentials.set({ scope, id: theirs, credential }) - const wrongUser = await owner.session.answer({ + const wrongUser = await laneAnswer(owner.session, { questionId, value: '[secret]', ref: theirs, - inputId: 'wrong-user', + controlId: 'wrong-user', }) expect(wrongUser).toMatchObject({ reason: 'unknown_secret_answer' }) expect(await credentials.get({ scope, id: theirs })).toEqual(credential) @@ -422,6 +534,42 @@ describe('an input that reaches a host that does not own the log', () => { await other.host.close() }) + it('cancels the turn of the owner when the offer comes while the owner appends', async () => { + const { a, b } = twoStores() + let other: Awaited> | undefined + let cancelIn: string | undefined + // The offer of B lands while the owner's append runs, so the store + // tells the owner about it during its own write. + const ownerLog = wrapLogStore(a, { + append: async (batch) => { + const operationId = cancelIn + if (operationId !== undefined && other) { + cancelIn = undefined + const receipt = await other.session.cancel({ operationId }) + expect(receipt.status).toBe('queued') + } + return a.append(batch) + }, + }) + const ownerModel = mockAdapter([untilAborted()]) + const owner = await hostOn(ownerLog, harnessWith(ownerModel.adapter)) + const otherHost = await hostOn(b, harnessWith(mockAdapter([]).adapter)) + other = otherHost + + const turn = owner.session.send('Go.') + const { operationId } = await turn.receipt + await waitFor(() => expect(ownerModel.calls).toHaveLength(1)) + cancelIn = operationId + // A write of the owner while its turn runs, with no write after it. + await owner.session.configure({ settings: { instructions: 'Be brief.' } }) + expect(cancelIn).toBeUndefined() + + await expect(Promise.resolve(turn)).rejects.toThrow() + expect(turn.status()).toBe('cancelled') + await owner.host.close() + await otherHost.host.close() + }) + it('starts an agent on the owner, and the run of the other host ends with it', async () => { const { a, b } = twoStores() const ran: Array = [] @@ -472,7 +620,9 @@ describe('an input that reaches a host that does not own the log', () => { const other = await hostOn(b, harness) const [first] = await other.session.transcript() - expect(await other.session.revert(first?.id ?? '')).toMatchObject({ + expect( + await other.session.revert({ messageId: first?.id ?? '' }), + ).toMatchObject({ status: 'rejected', reason: 'busy', }) @@ -709,3 +859,81 @@ describe('the input lane while idle', () => { await other.host.close() }) }) + +describe('a cancel from any server', () => { + it('stops the turn of the owner when POST control on another server sends it', async () => { + const { a, b } = twoStores() + const ownerModel = mockAdapter([untilAborted()]) + const owner = await hostOn(a, harnessWith(ownerModel.adapter)) + const harness = harnessWith(mockAdapter([]).adapter) + const other = await hostOn(b, harness) + const handler = createHarnessHandler({ + host: other.host, + harness, + openAccess: true, + }) + + const turn = owner.session.send('Go.') + const { operationId } = await turn.receipt + await waitFor(() => expect(ownerModel.calls).toHaveLength(1)) + const response = await handler( + new Request('http://x/api/harness/control', { + method: 'POST', + headers: { 'content-type': 'application/json' }, + body: JSON.stringify({ + threadId: THREAD, + input: { op: 'cancel', operationId }, + }), + }), + ) + + expect(await response.json()).toMatchObject({ status: 'queued' }) + await expect(Promise.resolve(turn)).rejects.toThrow() + expect(turn.status()).toBe('cancelled') + await owner.host.close() + await other.host.close() + }) + + it('does not stop a turn when the client of POST run leaves', async () => { + const { a } = twoStores() + const release = gate() + const harness = harnessWith( + mockAdapter([after(release.opened, 'Late.')]).adapter, + ) + const owner = await hostOn(a, harness) + const handler = createHarnessHandler({ + host: owner.host, + harness, + openAccess: true, + }) + + const response = await handler( + new Request('http://x/api/harness/run', { + method: 'POST', + headers: { 'content-type': 'application/json' }, + body: JSON.stringify({ + threadId: THREAD, + runId: 'run-1', + messages: [{ id: 'u1', role: 'user', content: 'Go.' }], + tools: [], + context: [], + state: {}, + forwardedProps: {}, + }), + }), + ) + const reader = response.body!.getReader() + await reader.read() + // The client leaves while the turn runs. + await reader.cancel() + release.open() + + await waitFor(() => + expect(owner.session.operation('run-1')?.status()).toBe('completed'), + ) + expect(contents(await owner.session.transcript()).at(-1)?.content).toBe( + 'Late.', + ) + await owner.host.close() + }) +}) diff --git a/packages/ai-harness/tests/lifecycle-events.test.ts b/packages/ai-harness/tests/lifecycle-events.test.ts new file mode 100644 index 0000000000..fb8e7d77e6 --- /dev/null +++ b/packages/ai-harness/tests/lifecycle-events.test.ts @@ -0,0 +1,312 @@ +import { describe, expect, it, vi } from 'vitest' +import { z } from 'zod' +import { EventType, toolDefinition } from '@tanstack/ai' +import { memoryLogStore, memoryPersistence } from '@tanstack/ai-persistence' +import { createHarnessHost, defineHarness, durableTool, logIdOf } from '../src' +import { workspaceTools } from '../src/first-party/coding/workspace' +import { + CRASHING_OWNERSHIP, + after, + appendToLog, + gate, + mockAdapter, + ownershipExpired, + text, + toolCall, + untilAborted, +} from './helpers' +import type { LogStore } from '@tanstack/ai-persistence' +import type { BoundWorkspace } from '../src/first-party/coding/backend' +import type { HarnessEventValues, HarnessSession } from '../src' +import type { Reply } from './helpers' + +// Coverage runs in CI are slow, so polled checks wait up to 5 s. +const waitFor = (check: () => T | Promise) => + vi.waitFor(check, { timeout: 5_000 }) + +const THREAD = 't1' + +const isObject = (value: unknown): value is Record => + typeof value === 'object' && value !== null + +/** The values of the `name` events in the log of `threadId`, in log order. */ +async function valuesOf( + log: LogStore, + threadId: string, + name: TName, +): Promise> { + const entries = await log.read({ logId: logIdOf({ logId: threadId }) }) + return entries.flatMap(({ record }) => { + const { event } = record + return record.type === 'harness.event' && + isObject(event) && + event.type === EventType.CUSTOM && + event.name === name + ? [event.value] + : [] + }) +} + +/** A host on `log` with the session of `threadId` open. */ +async function openOn( + log: LogStore, + replies: Array, + options: { + crashes?: boolean + tools?: Parameters[0]['tools'] + plugins?: Parameters[0]['plugins'] + threadId?: string + } = {}, +) { + const host = createHarnessHost({ + persistence: { stores: { log } }, + ...(options.crashes ? { ownership: CRASHING_OWNERSHIP } : {}), + }) + const session: HarnessSession = await host.open( + defineHarness({ + name: 'test/lifecycle', + adapter: mockAdapter(replies).adapter, + ...(options.tools ? { tools: options.tools } : {}), + ...(options.plugins ? { plugins: options.plugins } : {}), + }), + { threadId: options.threadId ?? THREAD }, + ) + return { host, session } +} + +describe('harness.input.queued', () => { + it('comes for an input that waits behind the running turn, with its delivery', async () => { + const log = memoryLogStore() + const first = gate() + const { host, session } = await openOn(log, [ + after(first.opened, 'one'), + () => text('two'), + () => text('three'), + ]) + const running = session.send('one', { inputId: 'in-1' }) + await waitFor(() => expect(session.snapshot().status).toBe('running')) + + const queued = session.send('two', { inputId: 'in-2' }) + const steered = session.send('three', { inputId: 'in-3', mode: 'steer' }) + await waitFor(async () => + expect(await valuesOf(log, THREAD, 'harness.input.queued')).toEqual([ + { inputId: 'in-2', delivery: 'queue' }, + { inputId: 'in-3', delivery: 'steer' }, + ]), + ) + first.open() + await Promise.all([running, queued, steered]) + await host.close() + }) + + it('does not come for an input that runs at once', async () => { + const log = memoryLogStore() + const { host, session } = await openOn(log, [() => text('one')]) + + await session.send('one') + + expect(await valuesOf(log, THREAD, 'harness.input.queued')).toEqual([]) + await host.close() + }) +}) + +describe('harness.ownership.changed', () => { + it('comes once when the first host acquires the log, not when it acquires it again', async () => { + const log = memoryLogStore() + const { host, session } = await openOn(log, [ + () => text('one'), + () => text('two'), + ]) + + const release = vi.spyOn(log, 'release') + await session.send('one') + // An idle host gives the log up. The next turn acquires it again. + await waitFor(() => expect(release).toHaveBeenCalled()) + await session.send('two') + await host.close() + + expect(await valuesOf(log, THREAD, 'harness.ownership.changed')).toEqual([ + { epoch: 1, reason: 'acquired' }, + ]) + // The host had the log at more than one epoch. + const other = await log.acquire({ + logId: logIdOf({ logId: THREAD }), + harness: 'other-host', + ownerId: 'other-host', + ttlMs: 60_000, + }) + expect(other?.epoch).toBeGreaterThan(2) + }) + + it('says taken_over when a host takes the log with the work of a stopped host', async () => { + const log = memoryLogStore() + const first = await openOn(log, [untilAborted()], { crashes: true }) + void Promise.resolve(first.session.send('one', { inputId: 'in-1' })).catch( + () => {}, + ) + await waitFor(() => expect(first.session.snapshot().status).toBe('running')) + // The first host stops: its ownership of the log expires. + await ownershipExpired() + + const next = await openOn(log, [() => text('one again')]) + await next.session.settled({ inputId: 'in-1' }) + + await waitFor(async () => + expect(await valuesOf(log, THREAD, 'harness.ownership.changed')).toEqual([ + { epoch: 1, reason: 'acquired' }, + { epoch: 2, reason: 'taken_over' }, + ]), + ) + await first.host.close().catch(() => {}) + await next.host.close() + }) + + it('says acquired when the old owner gave the log back with work left', async () => { + const log = memoryLogStore() + const first = await openOn(log, [untilAborted()]) + void Promise.resolve(first.session.send('one', { inputId: 'in-1' })).catch( + () => {}, + ) + await waitFor(() => expect(first.session.snapshot().status).toBe('running')) + // A graceful stop: the host releases the log, and the input waits. + await first.host.close({ recoverable: true }) + + const next = await openOn(log, [() => text('one again')]) + await next.session.settled({ inputId: 'in-1' }) + + await waitFor(async () => + expect(await valuesOf(log, THREAD, 'harness.ownership.changed')).toEqual([ + { epoch: 1, reason: 'acquired' }, + { epoch: 2, reason: 'acquired' }, + ]), + ) + await next.host.close() + }) +}) + +describe('harness.wait.opened and harness.wait.settled', () => { + /** A durable tool that asks for a name, and returns it. */ + const greet = durableTool( + toolDefinition({ + name: 'greet', + description: 'Ask the user for a name', + inputSchema: z.object({}), + }), + async (_args, { step }) => { + const name = await step.ask({ + name: 'name', + message: 'Name?', + schema: z.object({ name: z.string() }), + }) + return `Hello, ${name.name}.` + }, + ) + + it('come when a wait opens in the log, and when its answer ends it', async () => { + const log = memoryLogStore() + const { host, session } = await openOn( + log, + [() => toolCall('greet', {}, 'call-1'), () => text('ok')], + { tools: [greet] }, + ) + const turn = session.send('greet me') + await waitFor(() => expect(session.snapshot().pendingWaits).toHaveLength(1)) + const opened = await valuesOf(log, THREAD, 'harness.wait.opened') + expect(opened).toEqual([ + { + waitId: 'ask:call-1:name', + kind: 'question', + runId: turn.id, + message: 'Name?', + ...(isObject(opened[0]) && opened[0].expiresAt !== undefined + ? { expiresAt: expect.any(Number) } + : {}), + }, + ]) + + await session.answer({ + questionId: 'ask:call-1:name', + value: { name: 'Ada' }, + }) + await turn + + await waitFor(async () => + expect(await valuesOf(log, THREAD, 'harness.wait.settled')).toEqual([ + { waitId: 'ask:call-1:name', outcome: 'answered' }, + ]), + ) + await host.close() + }) +}) + +describe('harness.transferred', () => { + it('comes on the source thread and on the target thread', async () => { + const persistence = memoryPersistence() + const { log } = persistence.stores + const host = createHarnessHost({ persistence }) + const support = defineHarness({ + name: 'test/support', + adapter: mockAdapter(() => text('ok')).adapter, + }) + const billing = defineHarness({ + name: 'test/billing', + adapter: mockAdapter(() => text('billing here')).adapter, + }) + const source = await host.open(support, { threadId: 'src' }) + await source.send('one') + + const { transferId, threadId } = await host.transfer({ + from: { harness: support, threadId: 'src' }, + to: { harness: billing }, + envelope: { note: 'Refund for order 42.' }, + }) + + await waitFor(async () => { + expect(await valuesOf(log, threadId, 'harness.transferred')).toEqual([ + { transferId, from: { threadId: 'src' } }, + ]) + expect(await valuesOf(log, 'src', 'harness.transferred')).toEqual([ + { transferId, to: { threadId } }, + ]) + }) + await host.close() + }) +}) + +describe('harness.stranded', () => { + it('comes when the sandbox that the log binds the thread to is gone', async () => { + const log = memoryLogStore() + await appendToLog(log, { threadId: THREAD }, [ + { + type: 'harness.workspace.bound', + binding: { + kind: 'sandbox', + provider: 'fake', + sandboxId: 'sb-1', + root: '/ws', + }, + }, + ]) + const bind = vi.fn(async (): Promise => { + throw new Error('bind must not run') + }) + const { host, session } = await openOn( + log, + [() => toolCall('read_file', { path: 'a.txt' }, 'c1'), () => text('ok')], + // No reattach: the sandbox of the log cannot come back. + { plugins: () => [workspaceTools({ workspace: { bind }, web: false })] }, + ) + const turn = session.send('read') + await turn + + // Each try of the workspace sends one: before the turn, and at the call. + await waitFor(async () => { + const values = await valuesOf(log, THREAD, 'harness.stranded') + expect(values).not.toHaveLength(0) + for (const value of values) + expect(value).toEqual({ operationId: turn.id }) + }) + expect(bind).not.toHaveBeenCalled() + await host.close() + }) +}) diff --git a/packages/ai-harness/tests/log-stores.test.ts b/packages/ai-harness/tests/log-stores.test.ts index 9d4028405a..92e6ba06a2 100644 --- a/packages/ai-harness/tests/log-stores.test.ts +++ b/packages/ai-harness/tests/log-stores.test.ts @@ -10,6 +10,7 @@ import { HostOwnership } from '../src/ownership' import { mockAdapter, storedThreadState, text, toolCall } from './helpers' import type { LogStore } from '@tanstack/ai-persistence' import type { Reply } from './helpers' +import type { SessionImpl } from '../src/session/harness-session' const THREAD = 't1' @@ -171,11 +172,16 @@ describe('runs and interrupts of a session', () => { it('refuses a client run id that a run in the log has, also on a new host', async () => { const log = memoryLogStore() const first = await openOn(log, [() => text('Hi.')]) - await first.session.send('Hi.', { runId: 'run-1' }) + await (first.session as unknown as SessionImpl).sendForRun('Hi.', { + runId: 'run-1', + }) await first.host.close() const next = await openOn(log, [() => text('never')]) - const again = next.session.send('Hi again.', { runId: 'run-1' }) + const again = (next.session as unknown as SessionImpl).sendForRun( + 'Hi again.', + { runId: 'run-1' }, + ) expect(await again.receipt).toMatchObject({ status: 'rejected', reason: 'conflict', @@ -195,12 +201,20 @@ describe('runs and interrupts of a session', () => { tenantId: 'globex', }) - expect((await acme.session.send('Hi.', { runId: 'run-1' })).text).toBe( - 'Hello, acme.', - ) - expect((await globex.session.send('Hi.', { runId: 'run-1' })).text).toBe( - 'Hello, globex.', - ) + expect( + ( + await (acme.session as unknown as SessionImpl).sendForRun('Hi.', { + runId: 'run-1', + }) + ).text, + ).toBe('Hello, acme.') + expect( + ( + await (globex.session as unknown as SessionImpl).sendForRun('Hi.', { + runId: 'run-1', + }) + ).text, + ).toBe('Hello, globex.') await acme.host.close() await globex.host.close() }) diff --git a/packages/ai-harness/tests/log.test.ts b/packages/ai-harness/tests/log.test.ts index e1a9796df3..489cc95b9f 100644 --- a/packages/ai-harness/tests/log.test.ts +++ b/packages/ai-harness/tests/log.test.ts @@ -11,8 +11,11 @@ import { LogVersionError, SNAPSHOT_VERSION, SharedLog, + commonPrefix, + emptyLogState, emptySharedLogState, engineMessageStore, + foldEntry, loadLogState, logIdOf, logMessageStore, @@ -20,7 +23,7 @@ import { sessionOf, } from '../src/log' import { HostOwnership, LogOwnedElsewhereError } from '../src/ownership' -import { wrapLogStore, writeAsOtherHost } from './helpers' +import { CRASHING_OWNERSHIP, wrapLogStore, writeAsOtherHost } from './helpers' import { HARNESS_EVENTS } from '../src/types' import type { ModelMessage, StreamChunk } from '@tanstack/ai' import type { LogRecord, LogSnapshot, LogStore } from '@tanstack/ai-persistence' @@ -191,6 +194,29 @@ const project: ProjectOptions = { }, } +describe('commonPrefix', () => { + it('counts the same and the equal messages, and does not stringify the same object', () => { + const toJSON = vi.fn(() => ({ role: 'user', content: 'hi' })) + const shared: ModelMessage = { role: 'user', content: 'hi' } + Object.assign(shared, { toJSON }) + const a: Array = [ + shared, + { role: 'assistant', content: 'yes' }, + { role: 'user', content: 'one' }, + ] + const b: Array = [ + shared, + { role: 'assistant', content: 'yes' }, + { role: 'user', content: 'two' }, + ] + + expect(commonPrefix(a, b)).toBe(2) + expect(commonPrefix(a, a.slice(0, 1))).toBe(1) + expect(commonPrefix([], b)).toBe(0) + expect(toJSON).not.toHaveBeenCalled() + }) +}) + describe('transcript fold', () => { it('writes each change as one splice record and folds it back', async () => { const store = memoryLogStore() @@ -329,6 +355,27 @@ describe('fold robustness', () => { }) }) +describe('kept host records', () => { + it('keeps the last records of each type that keep names, and no other', () => { + const state = emptyLogState() + const keep = new Map([['app.step', 2]]) + const records = [ + { type: 'app.step', n: 1 }, + { type: 'app.other' }, + { type: 'app.step', n: 2 }, + { type: 'app.step', n: 3 }, + ] + records.forEach((record, index) => + foldEntry(state, { seq: index + 1, record }, undefined, keep), + ) + + expect(state.keptRecords.map((entry) => entry.record)).toEqual([ + { type: 'app.step', n: 2 }, + { type: 'app.step', n: 3 }, + ]) + }) +}) + describe('tool call records', () => { const result = (toolCallId: string): LogRecord => ({ type: 'harness.tool.result', @@ -661,7 +708,10 @@ describe('write failures', () => { // the test: a host that pauses. const first = new SharedLog({ store, - ownership: hostOwnership(store, { ttlMs: 20, renewMs: 60_000 }), + ownership: hostOwnership(store, { + ...CRASHING_OWNERSHIP, + ttlMs: 20, + }), harness: HARNESS, logId: LOG, rootThread: THREAD, @@ -1012,6 +1062,31 @@ describe('snapshots', () => { expect([...state.thread.config]).toEqual([['level', 'high']]) }) + it('gives usage totals from a snapshot before bySource an empty bySource', async () => { + const snapshot = snapshotOf({}) + const total = { + calls: 1, + promptTokens: 7, + completionTokens: 3, + totalTokens: 10, + cachedTokens: 0, + cacheWriteTokens: 0, + } + const [first] = (snapshot.state as { sessions: Array<[string, any]> }) + .sessions + first![1].usage = { total, byModel: { 'mock/m': total }, bySender: {} } + const store = await seededStore(snapshot) + + const state = await loadSession({ store, project }) + + expect(state.usage).toEqual({ + total, + byModel: { 'mock/m': total }, + bySender: {}, + bySource: {}, + }) + }) + it('folds the workspace binding and the backend ids of jobs, from records and from a snapshot', async () => { const binding = { kind: 'sandbox', @@ -1508,6 +1583,23 @@ describe('engine message store', () => { ]), ).toEqual([user('u1', 'one'), toolCallMessage, toolResult, signal]) }) + + it('skips a second commit of the same model call when nothing changed', async () => { + const { writer, engine } = await engineOnLog() + await engine.loadThread(THREAD) + const list = [user('u1', 'one'), toolCallMessage, toolResult] + await engine.beforeModel(list) + const commit = vi.spyOn(writer, 'commit') + + // The same message objects and the same fold: no commit. + expect(await engine.beforeModel([...list])).toBeUndefined() + expect(commit).not.toHaveBeenCalled() + + // A host record changes the fold: the next call commits and gives it. + await writer.append([{ type: 'app.signal', text: '[signal] build failed' }]) + expect(await engine.beforeModel([...list])).toEqual([...list, signal]) + expect(commit).toHaveBeenCalledTimes(1) + }) }) describe('logIdOf', () => { @@ -1736,7 +1828,10 @@ describe('a log shared by sessions', () => { const onIdle = vi.fn() const log = new SharedLog({ store, - ownership: hostOwnership(store, { ttlMs: 20, renewMs: 60_000 }), + ownership: hostOwnership(store, { + ...CRASHING_OWNERSHIP, + ttlMs: 20, + }), harness: HARNESS, logId: SHARED, rootThread: 'log-1', diff --git a/packages/ai-harness/tests/media-signing.test.ts b/packages/ai-harness/tests/media-signing.test.ts new file mode 100644 index 0000000000..46c819c642 --- /dev/null +++ b/packages/ai-harness/tests/media-signing.test.ts @@ -0,0 +1,82 @@ +import { describe, expect, it } from 'vitest' +import { createMediaSigner } from '../src/media-signing' + +/** Two keys of 32 bytes or more. */ +const OLD_KEY = 'old-key-0123456789-0123456789-0123' +const NEW_KEY = 'new-key-0123456789-0123456789-0123' + +/** A payload in the form of a signed media URL: one field per line. */ +const PAYLOAD = [ + 'test/harness', + 'tenant-1', + 'thread-1', + 'media-1', + 'user-1', + '1900000000000', +].join('\n') + +describe('createMediaSigner', () => { + it('verifies a signature that it made', async () => { + const signer = createMediaSigner({ secret: NEW_KEY }) + const signature = await signer.sign(PAYLOAD) + + // Base64url of a 32-byte HMAC SHA-256, without padding. + expect(signature).toMatch(/^[A-Za-z0-9_-]{43}$/) + expect(await signer.verify(PAYLOAD, signature)).toBe(true) + }) + + it('refuses a payload that changed after it was signed', async () => { + const signer = createMediaSigner({ secret: NEW_KEY }) + const signature = await signer.sign(PAYLOAD) + const tampered = PAYLOAD.replace('user-1', 'user-2') + + expect(await signer.verify(tampered, signature)).toBe(false) + }) + + it('refuses a signature that is not base64url', async () => { + const signer = createMediaSigner({ secret: NEW_KEY }) + + expect(await signer.verify(PAYLOAD, '%%%')).toBe(false) + }) + + it('refuses a signature from a key that it does not have', async () => { + const other = createMediaSigner({ secret: OLD_KEY }) + const signer = createMediaSigner({ secret: NEW_KEY }) + + expect(await signer.verify(PAYLOAD, await other.sign(PAYLOAD))).toBe(false) + }) + + it('signs with the first key and verifies with every key', async () => { + const before = createMediaSigner({ secret: OLD_KEY }) + const after = createMediaSigner({ keys: [NEW_KEY, OLD_KEY] }) + const newOnly = createMediaSigner({ secret: NEW_KEY }) + const oldSignature = await before.sign(PAYLOAD) + const newSignature = await after.sign(PAYLOAD) + + expect(await after.verify(PAYLOAD, oldSignature)).toBe(true) + expect(await newOnly.verify(PAYLOAD, newSignature)).toBe(true) + expect(await before.verify(PAYLOAD, newSignature)).toBe(false) + }) + + it('throws for a key shorter than 32 bytes', () => { + expect(() => createMediaSigner({ secret: 'short-secret' })).toThrow( + 'Media signing key 0 is 12 bytes. Use a key of at least 32 bytes (UTF-8).', + ) + }) + + it('counts key length in UTF-8 bytes, not in characters', () => { + // 16 characters, 32 bytes: each character is 2 bytes in UTF-8. + const twoByteKey = 'é'.repeat(16) + + expect(() => createMediaSigner({ secret: twoByteKey })).not.toThrow() + expect(() => createMediaSigner({ secret: 'é'.repeat(15) })).toThrow( + 'Media signing key 0 is 30 bytes.', + ) + }) + + it('throws for an empty key list', () => { + expect(() => createMediaSigner({ keys: [] })).toThrow( + 'Media signing needs at least one key.', + ) + }) +}) diff --git a/packages/ai-harness/tests/media.test.ts b/packages/ai-harness/tests/media.test.ts index b21e29123c..f965ea7fe7 100644 --- a/packages/ai-harness/tests/media.test.ts +++ b/packages/ai-harness/tests/media.test.ts @@ -217,6 +217,114 @@ describe('createMediaStore', () => { expect(decoder.decode(await store.load(record.id, range))).toBe(expected) }) + it.each([[undefined], [{ offset: 6 }], [{ offset: 0, length: 5 }]])( + 'opens the same bytes as a stream for the range %j', + async (range) => { + const { store } = setup() + const record = await store.put('hello world', { + mimeType: 'text/plain', + name: 'a.txt', + }) + const stream = await store.open(record.id, range) + expect(stream).toBeInstanceOf(ReadableStream) + expect(new Uint8Array(await new Response(stream).arrayBuffer())).toEqual( + await store.load(record.id, range), + ) + }, + ) + + it('keeps the base64 of a file, and drops the oldest over the cap', async () => { + const persistence = memoryPersistence() + const blobGet = vi.spyOn(persistence.stores.blobs, 'get') + // Each 5-byte file is 8 characters of base64: two fit in 16. + const store = createMediaStore({ + persistence, + threadId: 't1', + encodedBytes: 16, + }) + const put = (body: string) => + store.put(body, { mimeType: 'image/png', name: `${body}.png` }) + const [a, b, c] = [ + await put('aaaaa'), + await put('bbbbb'), + await put('ccccc'), + ] + const reads = async (id: string) => { + const before = blobGet.mock.calls.length + await store.dataUrl(id) + return blobGet.mock.calls.length - before + } + + expect(await reads(a.id)).toBe(1) + expect(await reads(a.id)).toBe(0) + expect(await reads(b.id)).toBe(1) + // `c` pushes out `a`, the oldest. + expect(await reads(c.id)).toBe(1) + expect(await reads(b.id)).toBe(0) + expect(await reads(c.id)).toBe(0) + expect(await reads(a.id)).toBe(1) + expect(await store.dataUrl(a.id)).toBe('data:image/png;base64,YWFhYWE=') + }) + + it('keeps one base64 cap for all the threads of a blob store', async () => { + const persistence = memoryPersistence() + const blobGet = vi.spyOn(persistence.stores.blobs, 'get') + // Each 5-byte file is 8 characters of base64: two fit in 16, for both + // threads together. + const one = createMediaStore({ + persistence, + threadId: 't1', + encodedBytes: 16, + }) + const two = createMediaStore({ + persistence, + threadId: 't2', + encodedBytes: 16, + }) + const put = (store: typeof one, body: string) => + store.put(body, { mimeType: 'image/png', name: `${body}.png` }) + const a = await put(one, 'aaaaa') + const b = await put(two, 'bbbbb') + const c = await put(one, 'ccccc') + const reads = async (store: typeof one, id: string) => { + const before = blobGet.mock.calls.length + await store.dataUrl(id) + return blobGet.mock.calls.length - before + } + + expect(await reads(one, a.id)).toBe(1) + expect(await reads(two, b.id)).toBe(1) + // `c` pushes out `a`, the oldest, though `a` is of the other thread. + expect(await reads(one, c.id)).toBe(1) + expect(await reads(two, b.id)).toBe(0) + expect(await reads(one, a.id)).toBe(1) + }) + + it('takes the largest base64 cap that a store of the blob store asks for', async () => { + const persistence = memoryPersistence() + const blobGet = vi.spyOn(persistence.stores.blobs, 'get') + const one = createMediaStore({ + persistence, + threadId: 't1', + encodedBytes: 16, + }) + // A later store asks for more: three files of 8 characters fit in 24. + createMediaStore({ persistence, threadId: 't2', encodedBytes: 24 }) + const put = (body: string) => + one.put(body, { mimeType: 'image/png', name: `${body}.png` }) + const ids = [ + (await put('aaaaa')).id, + (await put('bbbbb')).id, + (await put('ccccc')).id, + ] + for (const id of ids) await one.dataUrl(id) + const before = blobGet.mock.calls.length + + for (const id of ids) await one.dataUrl(id) + + expect(blobGet.mock.calls.length).toBe(before) + }) + it('makes a data URL up to 1 MB only', async () => { const { store } = setup() const small = await store.put(encoder.encode('hello'), { @@ -378,6 +486,32 @@ describe('mediaMiddleware', () => { ]) }) + it('does not load or encode the file again at the next turn', async () => { + const { persistence, store } = setup() + const record = await store.put(encoder.encode('hello'), { + mimeType: 'image/png', + name: 'cat.png', + }) + const blobGet = vi.spyOn(persistence.stores.blobs, 'get') + // A harness makes a new middleware at each turn, on the same store. + const first = await runTurn( + persistence, + mediaMiddleware({ store }), + userMessage(mediaPart(record)), + ) + const second = await runTurn( + persistence, + mediaMiddleware({ store }), + userMessage(mediaPart(record)), + ) + expect(blobGet).toHaveBeenCalledTimes(1) + expect(second.sent).toEqual(first.sent) + expect(second.sent).toContainEqual({ + type: 'image', + source: { type: 'data', value: 'aGVsbG8=', mimeType: 'image/png' }, + }) + }) + it.each([ [ 'audio when transcribe is not set', diff --git a/packages/ai-harness/tests/observability.test.ts b/packages/ai-harness/tests/observability.test.ts new file mode 100644 index 0000000000..d18b840612 --- /dev/null +++ b/packages/ai-harness/tests/observability.test.ts @@ -0,0 +1,130 @@ +import { describe, expect, it } from 'vitest' +import { + noopTracer, + parseTraceparent, + traceparentOf, + withSpan, +} from '../src/observability' +import { fakeSpan } from './helpers' +import type { Tracer } from '@opentelemetry/api' + +const TRACEPARENT = '00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01' + +/** A tracer that records what the harness calls on each span. */ +function recordingTracer() { + const spans: Array<{ name: string; options: unknown; calls: Array }> = + [] + const tracer: Pick = { + startSpan: (name, options) => { + const calls: Array = [] + spans.push({ name, options, calls }) + return fakeSpan( + { traceId: 'a'.repeat(32), spanId: 'b'.repeat(16), traceFlags: 1 }, + { + setAttribute: (key) => calls.push(`setAttribute:${key}`), + setStatus: ({ code, message }) => + calls.push(`setStatus:${code}:${message}`), + recordException: (error) => + calls.push(`recordException:${String(error)}`), + end: () => calls.push('end'), + }, + ) + }, + } + return { tracer, spans } +} + +describe('traceparent', () => { + it('reads a valid value and writes it back the same', () => { + const context = parseTraceparent(TRACEPARENT) + expect(context).toEqual({ + traceId: '4bf92f3577b34da6a3ce929d0e0e4736', + spanId: '00f067aa0ba902b7', + traceFlags: 1, + isRemote: true, + }) + expect(traceparentOf(context!)).toBe(TRACEPARENT) + }) + + it.each([ + ['not a string', 42], + ['the wrong shape', 'abc'], + ['upper case hex', TRACEPARENT.toUpperCase()], + ['version ff', TRACEPARENT.replace(/^00/, 'ff')], + ['an extra field on version 00', `${TRACEPARENT}-x`], + ['a zero trace id', `00-${'0'.repeat(32)}-00f067aa0ba902b7-01`], + [ + 'a zero span id', + `00-4bf92f3577b34da6a3ce929d0e0e4736-${'0'.repeat(16)}-01`, + ], + ])('ignores %s', (_name, value) => { + expect(parseTraceparent(value)).toBeUndefined() + }) + + it('reads the four fields of a later version that has more', () => { + expect( + parseTraceparent(`${TRACEPARENT.replace(/^00/, '01')}-extra`), + ).toMatchObject({ spanId: '00f067aa0ba902b7' }) + }) + + it('writes the flags as two hex characters', () => { + expect( + traceparentOf({ + traceId: 'a'.repeat(32), + spanId: 'b'.repeat(16), + traceFlags: 0, + }), + ).toBe(`00-${'a'.repeat(32)}-${'b'.repeat(16)}-00`) + }) +}) + +describe('withSpan', () => { + it('passes the options to startSpan and ends the span', async () => { + const { tracer, spans } = recordingTracer() + const link = { context: parseTraceparent(TRACEPARENT)! } + const result = await withSpan( + { + tracer, + name: 'harness.operation', + attributes: { 'harness.kind': 'chat' }, + links: [link], + startTime: 5, + }, + () => 'done', + ) + expect(result).toBe('done') + expect(spans).toEqual([ + { + name: 'harness.operation', + options: { + attributes: { 'harness.kind': 'chat' }, + links: [link], + startTime: 5, + }, + calls: ['end'], + }, + ]) + }) + + it('records the error, sets the error status, ends the span, and throws on', async () => { + const { tracer, spans } = recordingTracer() + await expect( + withSpan({ tracer, name: 'harness.recovery' }, async () => { + throw new Error('boom') + }), + ).rejects.toThrow('boom') + expect(spans[0]?.calls).toEqual([ + 'recordException:Error: boom', + 'setStatus:2:boom', + 'end', + ]) + }) +}) + +describe('noopTracer', () => { + it('gives one shared span, so nothing allocates per span', () => { + const first = noopTracer.startSpan('a') + expect(noopTracer.startSpan('b')).toBe(first) + expect(first.spanContext().traceFlags).toBe(0) + }) +}) diff --git a/packages/ai-harness/tests/on-event.test.ts b/packages/ai-harness/tests/on-event.test.ts index 32cddcef6d..230fd35898 100644 --- a/packages/ai-harness/tests/on-event.test.ts +++ b/packages/ai-harness/tests/on-event.test.ts @@ -2,7 +2,11 @@ import { afterEach, describe, expect, it, vi } from 'vitest' import { EventType } from '@tanstack/ai' import { memoryLogStore, memoryPersistence } from '@tanstack/ai-persistence' import { createHarnessHost, defineHarness } from '../src' +import { Redactor } from '../src/auth' +import { tapFeed } from '../src/session/helpers' import { mockAdapter, text } from './helpers' +import type { StreamChunk } from '@tanstack/ai' +import type { EventFeed } from '../src/feed' import type { HarnessPersistence, LiveEvent, SessionEvent } from '../src' // Coverage runs in CI are slow, so polled checks wait up to 5 s. @@ -183,3 +187,39 @@ describe('operation.onEvent', () => { await host.close() }) }) + +describe('a MESSAGES_SNAPSHOT event', () => { + it('goes to the listeners and to the log', () => { + const published: Array = [] + const feed: EventFeed = { + publish: (_operationId, event) => void published.push(event), + head: () => '0', + read: () => (async function* () {})(), + close: () => {}, + wake: () => {}, + } + const heard: Array = [] + const tapped = tapFeed(feed, () => {}, new Set(), new Redactor()) + tapped.listen(({ event }) => void heard.push(event.type)) + const snapshot: StreamChunk = { + type: EventType.MESSAGES_SNAPSHOT, + messages: [], + timestamp: 1, + } + const started: StreamChunk = { + type: EventType.RUN_STARTED, + threadId: 't1', + runId: 'r1', + timestamp: 1, + } + + tapped.publish('r1', snapshot) + tapped.publish('r1', started) + + expect(heard).toEqual([EventType.MESSAGES_SNAPSHOT, EventType.RUN_STARTED]) + expect(published.map((event) => event.type)).toEqual([ + EventType.MESSAGES_SNAPSHOT, + EventType.RUN_STARTED, + ]) + }) +}) diff --git a/packages/ai-harness/tests/operation-settled.test.ts b/packages/ai-harness/tests/operation-settled.test.ts new file mode 100644 index 0000000000..449fdfeb2e --- /dev/null +++ b/packages/ai-harness/tests/operation-settled.test.ts @@ -0,0 +1,73 @@ +import { describe, expect, it } from 'vitest' +import { SessionFeed } from '../src/feed' +import { Redactor } from '../src/auth' +import { OperationImpl } from '../src/operation' +import { tapFeed } from '../src/session/helpers' + +/** Room for every event of a test. */ +const EVENT_BYTES = 1_000_000 + +/** + * A new operation on its own feed. Its cancel fails it as `cancelled`, the + * same as a session does for an operation that did not start. + */ +function createOperation() { + const feed = tapFeed( + new SessionFeed({ threadId: 'thread-1', eventBytes: EVENT_BYTES }), + () => {}, + new Set(), + new Redactor(), + ) + return new OperationImpl('chat', feed, (operation) => { + operation.fail('cancelled', new Error('Cancelled.')) + return Promise.resolve({ inputId: 'input-1', status: 'accepted' }) + }) +} + +describe('a settled operation', () => { + it('stays completed when fail comes after finish', async () => { + const operation = createOperation() + operation.finish('completed', 'first answer') + operation.fail('failed', new Error('Too late.')) + + expect(operation.status()).toBe('completed') + await expect(operation).resolves.toBe('first answer') + }) + + it('stays failed when finish comes after fail', async () => { + const operation = createOperation() + operation.fail('failed', new Error('Model error.')) + operation.finish('completed', 'too late') + + expect(operation.status()).toBe('failed') + await expect(operation).rejects.toThrow('Model error.') + }) + + it('stays cancelled when finish comes after cancel', async () => { + const operation = createOperation() + await operation.cancel() + operation.finish('completed', 'too late') + + expect(operation.status()).toBe('cancelled') + await expect(operation).rejects.toThrow('Cancelled.') + }) + + it('keeps the first result when finish comes twice', async () => { + const operation = createOperation() + operation.finish('interrupted', 'first answer') + operation.finish('completed', 'second answer') + + expect(operation.status()).toBe('interrupted') + await expect(operation).resolves.toBe('first answer') + }) + + it('keeps its final status when setStatus comes after it settled', async () => { + const operation = createOperation() + operation.finish('completed', 'answer') + operation.setStatus('running') + + expect(operation.status()).toBe('completed') + expect(operation.isSettled()).toBe(true) + await expect(operation).resolves.toBe('answer') + }) +}) diff --git a/packages/ai-harness/tests/ownership.test.ts b/packages/ai-harness/tests/ownership.test.ts index 212a718c76..32fa76e498 100644 --- a/packages/ai-harness/tests/ownership.test.ts +++ b/packages/ai-harness/tests/ownership.test.ts @@ -79,7 +79,7 @@ const assistantTexts = ( .map((message) => message.content) describe('ownership of a log by two hosts', () => { - it('refuses the writes of a host while another host owns the log and renews it', async () => { + it('sends the inputs of a host to the owner of the log, and refuses its other writes', async () => { const { logA, logB } = twoHosts() const release = gate() // A renews well before its ownership would expire. @@ -99,9 +99,20 @@ describe('ownership of a log by two hosts', () => { inputId: 'in-b', status: 'queued', }) - await expect(b.session.append([{ type: 'app.from-b' }])).rejects.toThrow( - 'Another host owns the log', - ) + // An append is an input too: it waits for A, and A writes it. + expect( + await b.session.append({ + records: [{ type: 'app.from-b' }], + inputId: 'in-append', + }), + ).toEqual({ inputId: 'in-append', status: 'queued' }) + // A write that is not an input is refused. + await expect( + b.session.schedule({ + at: Date.now() + 60_000, + input: { op: 'prompt', message: 'later' }, + }), + ).rejects.toThrow('Another host owns the log') // B still reads: it gets the answers of A. release.open() @@ -113,10 +124,10 @@ describe('ownership of a log by two hosts', () => { ]), ) expect(b.calls).toHaveLength(0) - const types = (await logA.read({ logId: LOG })).map( - (entry) => entry.record.type, - ) - expect(types).not.toContain('app.from-b') + const types = async () => + (await logA.read({ logId: LOG })).map((entry) => entry.record.type) + await waitFor(async () => expect(await types()).toContain('app.from-b')) + expect(await types()).not.toContain('harness.schedule') await a.host.close() await b.host.close() }) @@ -366,7 +377,7 @@ describe('renewals', () => { const ownership = new HostOwnership({ store: log, ownerId: 'host-1', - timing: { ttlMs: 20, renewMs: 40 }, + timing: { ...CRASHING_OWNERSHIP, ttlMs: 20, renewMs: 40 }, }) const lost = vi.fn() ownership.onLost(LOG, lost) @@ -450,3 +461,40 @@ describe('the run context', () => { await host.close() }) }) + +describe('ownership timing', () => { + it('renews at a third of ttlMs when only ttlMs is set', async () => { + const intervals = vi.spyOn(globalThis, 'setInterval') + const ownership = new HostOwnership({ + store: memoryLogStore(), + ownerId: 'host-1', + timing: { ttlMs: 6_000 }, + }) + await ownership.acquire(LOG, 'test/ownership') + + expect(intervals.mock.calls.map(([, ms]) => ms)).toContain(2_000) + await ownership.close() + intervals.mockRestore() + }) + + it('refuses a renewMs that is not less than ttlMs', () => { + expect(() => + createHarnessHost({ + persistence: memoryPersistence(), + ownership: { ttlMs: 5_000, renewMs: 5_000 }, + }), + ).toThrow(RangeError) + expect(() => + createHarnessHost({ + persistence: memoryPersistence(), + ownership: { renewMs: 30_000 }, + }), + ).toThrow('must be less than ownership.ttlMs') + expect(() => + createHarnessHost({ + persistence: memoryPersistence(), + ownership: { ttlMs: 0 }, + }), + ).toThrow('more than 0') + }) +}) diff --git a/packages/ai-harness/tests/permissions.test.ts b/packages/ai-harness/tests/permissions.test.ts index 7410fac395..631eec1f68 100644 --- a/packages/ai-harness/tests/permissions.test.ts +++ b/packages/ai-harness/tests/permissions.test.ts @@ -11,7 +11,7 @@ import { listSavedPermissions, permissions, } from '../src/first-party/permissions' -import { mockAdapter, text, toolCall } from './helpers' +import { gate, mockAdapter, text, toolCall } from './helpers' import type { MetadataStore } from '@tanstack/ai-persistence' import type { HarnessPersistence, @@ -369,6 +369,43 @@ describe('permissions()', () => { await second.host.close() }) + it('decides on the final arguments: a middleware before it cannot rewrite a call past a deny', async () => { + const tools = fakeTools() + // It runs before permissions(), and its onBeforeToolCall wins first. + const rewrite = definePlugin({ + name: 'test/rewrite', + setup: () => ({ + middleware: [ + { + name: 'test/rewrite', + onBeforeToolCall: (_run, hook) => + hook.toolName === 'read_file' + ? { type: 'transformArgs', args: { path: 'secrets/key' } } + : undefined, + }, + ], + }), + }) + const deny: PermissionRule = { + tool: 'read_file', + resource: 'secrets/**', + decision: 'deny', + } + const { host, session, calls } = await open( + memoryPersistence(), + [rewrite, permissions({ root: ROOT, rules: [deny] }), tools.plugin], + [() => toolCall('read_file', { path: 'src/a.ts' }, 'c1')], + ) + + await session.send('Read a file.') + + expect(tools.ran).toEqual([]) + expect(JSON.stringify(calls[1].messages)).toContain( + 'A permission rule denies this tool.', + ) + await host.close() + }) + it('keeps an always answer for the session when there is no metadata store', async () => { const { stores } = memoryPersistence() const tools = fakeTools() @@ -411,6 +448,41 @@ describe('permissions()', () => { await host.close() }) + it('keeps the mode of a turn until it ends: a change applies at the next turn', async () => { + const tools = fakeTools() + const started = gate() + const release = gate() + const { host, session } = await open( + memoryPersistence(), + [permissions({ root: ROOT }), tools.plugin], + [ + () => + (async function* () { + started.open() + await release.opened + yield* toolCall('bash', { command: 'ls' }, 'c1') + })(), + () => text('listed'), + () => toolCall('bash', { command: 'ls' }, 'c2'), + ], + ) + + const turn = session.send('list it') + await started.opened + await session.setConfig({ key: 'mode', value: 'bypass' }) + expect(session.config().mode?.value).toBe('bypass') + release.open() + // The running turn still has the default mode, so bash asks. + await answer(session, { answer: 'once' }) + await turn + + // The next turn has bypass, so bash runs with no question. + await session.send('list it again') + expect(session.snapshot().pendingQuestions).toHaveLength(0) + expect(tools.ran).toEqual(['bash', 'bash']) + await host.close() + }) + it('gives the model a tool error with the message of a reject', async () => { const tools = fakeTools() const { host, session, calls } = await open( diff --git a/packages/ai-harness/tests/plugin-hotpaths.test.ts b/packages/ai-harness/tests/plugin-hotpaths.test.ts new file mode 100644 index 0000000000..cbb11f4921 --- /dev/null +++ b/packages/ai-harness/tests/plugin-hotpaths.test.ts @@ -0,0 +1,157 @@ +import { mkdtemp, rm, utimes, writeFile } from 'node:fs/promises' +import { tmpdir } from 'node:os' +import { join } from 'node:path' +import { afterEach, describe, expect, it, vi } from 'vitest' +import { memoryPersistence } from '@tanstack/ai-persistence' +import { + createExtensionPoint, + createHarnessHost, + defineCommand, + defineHarness, + definePlugin, +} from '../src' +import { decidePermission } from '../src/first-party/permissions' +import { projectInstructions } from '../src/first-party/node' +import { mockAdapter, text } from './helpers' +import type { HarnessHost, HarnessPlugin } from '../src' +import type { PermissionRule } from '../src/first-party/permissions' +import type * as NodeFs from 'node:fs' + +// The sync calls a turn must not make. They wrap the real ones. +const sync = vi.hoisted(() => ({ + statSync: vi.fn(), + readFileSync: vi.fn(), +})) +vi.mock('node:fs', async (importOriginal) => { + const fs = await importOriginal() + sync.statSync.mockImplementation(fs.statSync) + sync.readFileSync.mockImplementation(fs.readFileSync) + return { ...fs, statSync: sync.statSync, readFileSync: sync.readFileSync } +}) + +const hosts: Array = [] +afterEach(async () => { + vi.restoreAllMocks() + await Promise.all(hosts.splice(0).map((host) => host.close())) +}) + +/** Open a session whose plugin list the test can change before a reload. */ +async function open(plugins: Array) { + const { adapter, calls } = mockAdapter(() => text('ok')) + const host = createHarnessHost({ persistence: memoryPersistence() }) + hosts.push(host) + const session = await host.open( + defineHarness({ + name: 'test/hotpaths', + adapter, + plugins: () => [...plugins], + }), + { threadId: 't' }, + ) + return { session, calls } +} + +describe('ctx.collect', () => { + const Rules = createExtensionPoint<{ tool: string }>('test/hot-rules') + const writer = (name: string, tool: string) => + definePlugin({ + name, + setup: () => ({ contribute: [Rules.item({ tool })] }), + }) + + it('gives the same array at run time until the contributions change', async () => { + let view: ReadonlyArray<{ tool: string }> = [] + const reader = definePlugin({ + name: 'test/reader', + setup: (ctx) => { + view = ctx.collect(Rules) + return { + commands: { + snap: defineCommand({ + description: 'Collect the rules', + run: () => ctx.collect(Rules), + }), + }, + } + }, + }) + const plugins = [reader, writer('test/a', 'bash')] + const { session } = await open(plugins) + + const first = await session.command({ name: 'snap' }) + const second = await session.command({ name: 'snap' }) + expect(second).toBe(first) + expect(first).toEqual([{ tool: 'bash' }]) + // The view from setup filled as the later plugin set up. + expect([...view]).toEqual([{ tool: 'bash' }]) + + plugins.push(writer('test/b', 'write_file')) + await session.reload() + const changed = await session.command({ name: 'snap' }) + expect(changed).not.toBe(first) + expect(changed).toEqual([{ tool: 'bash' }, { tool: 'write_file' }]) + }) +}) + +describe('permission rule patterns', () => { + it('compiles each rule pattern once for many calls', () => { + // Patterns no other test uses, so the cache has none of them yet. + const rules: Array = [ + { tool: 'bash', decision: 'allow', resource: 'hotpath-tool *' }, + { tool: 'edit_file', decision: 'allow', resource: 'hotpath/**/*.ts' }, + ] + const call = () => [ + decidePermission(rules, 'bash', 'default', { + resources: { commands: ['hotpath-tool run'] }, + fallback: 'ask', + }), + decidePermission(rules, 'edit_file', 'default', { + resources: { paths: ['/root/hotpath/src/a.ts'] }, + root: '/root', + fallback: 'ask', + }), + ] + expect(call()).toEqual(['allow', 'allow']) + + const compiled = vi.spyOn(globalThis, 'RegExp') + const results = Array.from({ length: 50 }, call) + const compiledCount = compiled.mock.calls.length + compiled.mockRestore() + + expect(compiledCount).toBe(0) + expect(new Set(results.map((result) => result.join()))).toEqual( + new Set(['allow,allow']), + ) + }) +}) + +describe('projectInstructions file checks', () => { + let dir = '' + afterEach(async () => { + await rm(dir, { recursive: true, force: true }) + }) + + it('finds a changed file with no sync file call in a turn', async () => { + dir = await mkdtemp(join(tmpdir(), 'harness-hotpaths-')) + const file = join(dir, 'AGENTS.md') + await writeFile(file, 'Use tabs.') + const { session, calls } = await open([ + projectInstructions({ root: dir, env: false }), + ]) + sync.statSync.mockClear() + sync.readFileSync.mockClear() + + await session.send('one') + await writeFile(file, 'Use spaces.') + const later = Date.now() / 1000 + 60 + await utimes(file, later, later) + await session.send('two') + + expect(calls[1].systemPrompts).toEqual([ + 'Project instructions from AGENTS.md:\nUse tabs.', + 'Project instructions from AGENTS.md changed. Follow this text instead:\nUse spaces.', + ]) + expect(sync.statSync).not.toHaveBeenCalled() + expect(sync.readFileSync).not.toHaveBeenCalled() + }) +}) diff --git a/packages/ai-harness/tests/plugin-model.test.ts b/packages/ai-harness/tests/plugin-model.test.ts new file mode 100644 index 0000000000..1f38016613 --- /dev/null +++ b/packages/ai-harness/tests/plugin-model.test.ts @@ -0,0 +1,268 @@ +import { describe, expect, it, vi } from 'vitest' +import { EventType } from '@tanstack/ai' +import { memoryLogStore, memoryPersistence } from '@tanstack/ai-persistence' +import { + createHarnessHost, + defineCommand, + defineHarness, + definePlugin, + logIdOf, +} from '../src' +import { gate, mockAdapter, text, untilAborted } from './helpers' +import type { StreamChunk } from '@tanstack/ai' +import type { HarnessPlugin, SessionEvent } from '../src' +import type { Reply } from './helpers' + +const THREAD = 't1' +const alice = { id: 'alice' } + +// Coverage runs in CI are slow, so polled checks wait up to 5 s. +const waitFor = (check: () => T | Promise) => + vi.waitFor(check, { timeout: 5_000 }) + +/** `content` as one model call that reports 7 prompt and 3 completion tokens. */ +const answerWithUsage = + (content: string): Reply => + () => + text(content).map( + (chunk): StreamChunk => + chunk.type === EventType.RUN_FINISHED + ? { + ...chunk, + usage: { promptTokens: 7, completionTokens: 3, totalTokens: 10 }, + } + : chunk, + ) + +/** Open `THREAD` with `plugins`. The main model answers with `main`. */ +async function open(plugins: Array, main: Array) { + const log = memoryLogStore() + const { metadata } = memoryPersistence().stores + const host = createHarnessHost({ persistence: { stores: { log, metadata } } }) + const session = await host.open( + defineHarness({ + name: 'test/plugin-model', + adapter: mockAdapter(main).adapter, + plugins: () => plugins, + }), + { threadId: THREAD, principal: alice }, + ) + return { host, session, log } +} + +describe('ctx.model.chat', () => { + it('counts the usage under plugin: with the label, in the session and in the run', async () => { + const checker = mockAdapter(answerWithUsage('looks fine')) + let runId = '' + const plugin = definePlugin({ + name: 'test/checker', + setup: (ctx) => ({ + commands: { + check: defineCommand({ + description: 'Check', + run: async (_input, { run }) => { + runId = run.runId + const { text: answer } = await ctx.model.chat({ + adapter: checker.adapter, + run, + label: 'check', + options: { messages: [{ role: 'user', content: 'Check it' }] }, + }) + return answer + }, + }), + }, + }), + }) + const { host, session, log } = await open([plugin], []) + const seen: Array = [] + const reader = new AbortController() + const reading = (async () => { + for await (const entry of session.events({ signal: reader.signal })) + seen.push(entry) + })() + + expect(await session.command({ name: 'check' })).toBe('looks fine') + + const usage = session.usage() + expect(usage.total).toMatchObject({ calls: 1, totalTokens: 10 }) + expect(usage.bySource).toEqual({ 'plugin:test/checker': usage.total }) + expect(usage.bySender).toEqual({ alice: usage.total }) + const records = (await log.read({ logId: logIdOf({ logId: THREAD }) })) + .map((entry) => entry.record) + .filter((record) => record.type === 'harness.usage') + expect(records).toEqual([ + expect.objectContaining({ + source: 'plugin:test/checker', + label: 'check', + runId, + principal: alice, + }), + ]) + await waitFor(() => + expect( + seen.find( + (entry) => + entry.event.type === EventType.CUSTOM && + entry.event.name === 'harness.usage', + )?.event, + ).toMatchObject({ + value: { source: 'plugin:test/checker', label: 'check' }, + }), + ) + reader.abort() + await reading.catch(() => {}) + await host.close() + }) + + it('uses the model and the run of the turn by default, and stops when the turn is cancelled', async () => { + let call: Promise | undefined + const plugin = definePlugin({ + name: 'test/side-call', + lifetime: 'turn', + setup: (ctx) => ({ + middleware: [ + { + name: 'test/side-call', + onConfig: (run) => { + if (run.phase !== 'init' || call) return + call = ctx.model.chat({ + options: { messages: [{ role: 'user', content: 'Side' }] }, + }) + // The test reads the rejection below. + call.catch(() => {}) + }, + }, + ], + }), + }) + // The plugin call and the turn both wait until they are aborted. + const { host, session } = await open( + [plugin], + [untilAborted(), untilAborted()], + ) + + const turn = session.send('go') + await waitFor(() => expect(call).toBeDefined()) + await turn.cancel() + + await expect(call).rejects.toThrow() + await expect(turn).rejects.toThrow('Cancelled.') + await host.close() + }) + + it('stops after timeoutMs', async () => { + const slow = mockAdapter(untilAborted()) + const plugin = definePlugin({ + name: 'test/slow', + setup: (ctx) => ({ + commands: { + slow: defineCommand({ + description: 'Slow', + run: async (_input, { run }) => { + const failure = await ctx.model + .chat({ + adapter: slow.adapter, + run, + timeoutMs: 20, + options: { messages: [{ role: 'user', content: 'Wait' }] }, + }) + .then( + () => undefined, + (error: unknown) => error, + ) + return failure instanceof Error ? failure.name : 'no error' + }, + }), + }, + }), + }) + const { host, session } = await open([plugin], []) + + expect(await session.command({ name: 'slow' })).toBe('TimeoutError') + expect(slow.calls).toHaveLength(1) + await host.close() + }) + + it('needs a run in a session plugin', async () => { + let failure: unknown + const plugin = definePlugin({ + name: 'test/no-run', + setup: (ctx) => ({ + commands: { + go: defineCommand({ + description: 'Go', + run: async () => { + failure = await ctx.model + .chat({ options: { messages: [] } }) + .catch((error: unknown) => error) + }, + }), + }, + }), + }) + const { host, session } = await open([plugin], []) + + await session.command({ name: 'go' }) + + expect(String(failure)).toContain( + 'ctx.model.chat of test/no-run needs a run', + ) + await host.close() + }) + + it('records the usage of a call that ends after its run, with no event after the run finished', async () => { + const late = gate() + const checker = mockAdapter(() => + (async function* () { + await late.opened + yield* answerWithUsage('late')({}) + })(), + ) + let call: Promise | undefined + let runId = '' + const plugin = definePlugin({ + name: 'test/late', + setup: (ctx) => ({ + commands: { + go: defineCommand({ + description: 'Go', + // Next to the run, as the title plugin does: not awaited. + run: async (_input, { run }) => { + runId = run.runId + call = ctx.model.chat({ + adapter: checker.adapter, + run, + options: { messages: [{ role: 'user', content: 'Later' }] }, + }) + }, + }), + }, + }), + }) + const { host, session } = await open([plugin], []) + const seen: Array = [] + const reader = new AbortController() + const reading = (async () => { + for await (const entry of session.events({ signal: reader.signal })) + seen.push(entry) + })() + + await session.command({ name: 'go' }) + late.open() + await call + + await waitFor(() => expect(session.usage().total.calls).toBe(1)) + const ofRun = seen.filter((entry) => entry.operationId === runId) + const finished = ofRun.findIndex( + (entry) => + entry.event.type === EventType.CUSTOM && + entry.event.name === 'harness.operation.finished', + ) + expect(finished).toBeGreaterThan(-1) + expect(ofRun.slice(finished + 1)).toEqual([]) + reader.abort() + await reading.catch(() => {}) + await host.close() + }) +}) diff --git a/packages/ai-harness/tests/plugin-owners.test.ts b/packages/ai-harness/tests/plugin-owners.test.ts new file mode 100644 index 0000000000..f553e49028 --- /dev/null +++ b/packages/ai-harness/tests/plugin-owners.test.ts @@ -0,0 +1,238 @@ +import { describe, expect, it, vi } from 'vitest' +import { memoryPersistence } from '@tanstack/ai-persistence' +import { createCapability, defineAgent, toolDefinition } from '@tanstack/ai' +import { AgentRegistry } from '../src/agents' +import { defineCommand } from '../src/commands' +import { configOption } from '../src/config' +import { createExtensionPoint, createPluginEvent } from '../src/extensions' +import { definePlugin, mountPlugins } from '../src/plugins' +import { createHarnessHost, defineHarness } from '../src' +import { mockAdapter, text } from './helpers' +import type { MountedPlugins, PluginSetupContext } from '../src/plugins' + +const env = () => ({ + threadId: 't1', + registry: new AgentRegistry(), + harnessTools: [], + harnessProvides: [], +}) + +const Rules = createExtensionPoint('test/rules') + +const agentNamed = (name: string) => + defineAgent({ name, description: name, run: async () => name }) + +/** A plugin `id` that contributes one item of every kind, each named `id`. */ +function contributor(id: string) { + const capability = createCapability()(`test/${id}-capability`) + const pick = () => undefined + const plugin = definePlugin({ + name: id, + provides: [capability], + setup: ({ provide }) => { + provide(capability, `${id} value`) + return { + tools: [ + toolDefinition({ name: `${id}_tool`, description: id }).server( + () => id, + ), + ], + prompts: [ + `${id} prompt`, + { id: `${id}:context`, text: `${id} context`, placement: 'context' }, + ], + middleware: [{ name: `${id}-middleware` }], + agentMiddleware: [{ name: `${id}-agent-middleware` }], + generationMiddleware: [{ name: `${id}-generation-middleware` }], + agents: [agentNamed(`${id}-agent`)], + subagents: [agentNamed(`${id}-subagent`)], + commands: { + [`${id}-command`]: defineCommand({ description: id, run: () => id }), + }, + config: { + [`${id}-level`]: configOption.select({ + options: ['low', 'high'], + default: 'low', + }), + }, + contribute: [Rules.item(`${id} rule`)], + adapter: pick, + discoverTools: () => [], + prepareTools: ({ tools }) => tools, + } + }, + }) + return { plugin, capability, pick } +} + +/** The name or value of every contribution in `mounted`. */ +function contents(mounted: MountedPlugins) { + const names = (items: ReadonlyArray<{ name?: string }>) => + items.map((item) => item.name) + return { + plugins: mounted.owners.plugins.map((plugin) => plugin.name), + tools: names(mounted.tools), + prompts: mounted.prompts, + contextPrompts: mounted.contextPrompts.map((item) => item.text), + middleware: names(mounted.middleware), + agentMiddleware: names(mounted.agentMiddleware), + generationMiddleware: names(mounted.generationMiddleware), + agents: names(mounted.agents), + subagents: names(mounted.subagents), + commands: [...mounted.commands.keys()], + config: [...mounted.config.keys()], + rules: (mounted.extensions.get(Rules.name) ?? []).map((item) => item.value), + adapters: mounted.adapters, + discoverers: mounted.discoverers.map((item) => item.owner), + preparers: mounted.preparers.map((item) => item.owner), + capabilities: mounted.capabilityBridge?.provides ?? [], + } +} + +describe('MountedPlugins.withoutOwners', () => { + it('removes every kind of contribution of a plugin and keeps the others', async () => { + const a = contributor('a') + const b = contributor('b') + const mounted = await mountPlugins([a.plugin, b.plugin], env()) + + const kept = mounted.withoutOwners(['a']) + + expect(contents(kept)).toEqual({ + plugins: ['b'], + tools: ['b_tool'], + prompts: ['b prompt'], + contextPrompts: ['b context'], + middleware: ['b-middleware'], + agentMiddleware: ['b-agent-middleware'], + generationMiddleware: ['b-generation-middleware'], + agents: ['b-agent'], + subagents: ['b-subagent'], + commands: ['b-command'], + config: ['b-level'], + rules: ['b rule'], + adapters: [b.pick], + discoverers: ['b'], + preparers: ['b'], + capabilities: [b.capability], + }) + expect(kept.values.get(a.capability)).toBeUndefined() + expect(kept.values.get(b.capability)).toBe('b value') + // The mounted set itself keeps both plugins. + expect(contents(mounted).agents).toEqual(['a-agent', 'b-agent']) + expect(mounted.values.get(a.capability)).toBe('a value') + }) + + it('removes an agent that a plugin set after the mount', async () => { + let saved: PluginSetupContext | undefined + const late = definePlugin({ + name: 'late', + setup: (ctx) => { + saved = ctx + }, + }) + const mounted = await mountPlugins([late, contributor('b').plugin], env()) + saved?.agents.set(agentNamed('late-subagent'), { subagent: true }) + + expect(contents(mounted).subagents).toEqual(['b-subagent', 'late-subagent']) + expect(contents(mounted.withoutOwners(['late'])).subagents).toEqual([ + 'b-subagent', + ]) + }) + + it('reads the commands and agents that a kept plugin sets after the copy', async () => { + let saved: PluginSetupContext | undefined + const late = definePlugin({ + name: 'late', + setup: (ctx) => { + saved = ctx + }, + }) + const mounted = await mountPlugins([late, contributor('b').plugin], env()) + const kept = mounted.withoutOwners(['b']) + + saved?.commands.set( + 'deploy', + defineCommand({ description: 'Deploy', run: () => 'deployed' }), + ) + saved?.agents.set(agentNamed('late-subagent'), { subagent: true }) + + expect(contents(kept).commands).toEqual(['deploy']) + expect(contents(kept).subagents).toEqual(['late-subagent']) + saved?.commands.delete('deploy') + expect(contents(kept).commands).toEqual([]) + }) + + it('refuses to remove a plugin whose capability a kept plugin requires', async () => { + const a = contributor('a') + const reader = definePlugin({ + name: 'reader', + requires: [a.capability], + }) + const mounted = await mountPlugins([a.plugin, reader], env()) + + expect(() => mounted.withoutOwners(['a'])).toThrow( + 'Cannot remove plugin a: plugin reader requires its capability "test/a-capability".', + ) + expect(contents(mounted.withoutOwners(['a', 'reader'])).plugins).toEqual([]) + }) +}) + +describe('ctx.on', () => { + it('leaves no abort listener on the plugin signal after off()', async () => { + let saved: PluginSetupContext | undefined + const listener = definePlugin({ + name: 'listener', + setup: (ctx) => { + saved = ctx + }, + }) + await mountPlugins([listener], env()) + if (!saved) throw new Error('setup did not run') + const signal = saved.resources.signal + const add = vi.spyOn(signal, 'addEventListener') + const removeListener = vi.spyOn(signal, 'removeEventListener') + const ping = createPluginEvent('test/ping') + + for (let index = 0; index < 3; index++) { + saved.on(ping, () => {})() + } + + const added = add.mock.calls.map((call) => call[1]) + const removed = removeListener.mock.calls.map((call) => call[1]) + expect(added).toHaveLength(3) + expect(added.filter((item) => !removed.includes(item))).toEqual([]) + }) +}) + +describe('a thread that removes a plugin', () => { + it('does not run an agent of that plugin by name', async () => { + const worker = defineAgent({ + name: 'worker', + description: 'Works', + run: async () => 'worked', + }) + const plugin = definePlugin({ + name: 'test/workers', + setup: () => ({ agents: [worker] }), + }) + const host = createHarnessHost({ persistence: memoryPersistence() }) + const session = await host.open( + defineHarness({ + name: 'test/plugin-owners', + adapter: mockAdapter(() => text('ok')).adapter, + plugins: () => [plugin], + }), + { threadId: 't' }, + ) + const byName = session.agent('worker') + expect(await byName?.run()).toBe('worked') + + await session.configure({ + settings: { plugins: { remove: ['test/workers'] } }, + }) + + expect(session.agent('worker')).toBeUndefined() + await expect(byName?.start()).rejects.toThrow('Unknown agent: worker') + await host.close() + }) +}) diff --git a/packages/ai-harness/tests/prompt-placement.test.ts b/packages/ai-harness/tests/prompt-placement.test.ts new file mode 100644 index 0000000000..1be8e9256f --- /dev/null +++ b/packages/ai-harness/tests/prompt-placement.test.ts @@ -0,0 +1,124 @@ +import { describe, expect, it } from 'vitest' +import { toolDefinition } from '@tanstack/ai' +import { z } from 'zod' +import { memoryPersistence } from '@tanstack/ai-persistence' +import { createHarnessHost, defineHarness } from '../src' +import { createExtensionPoint } from '../src/extensions' +import { definePlugin } from '../src/plugins' +import { mockAdapter, text, toolCall } from './helpers' +import type { SystemPrompt } from '@tanstack/ai' + +/** The text of the last system prompt of a model call. */ +const lastSystem = (call: { systemPrompts?: Array }) => { + const prompt = call.systemPrompts?.at(-1) + return typeof prompt === 'string' ? prompt : prompt?.content +} + +describe('a plugin prompt with placement', () => { + it('keeps the default in the system prompts, and puts a context prompt last, once per call', async () => { + let step = 0 + const tick = toolDefinition({ + name: 'tick', + description: 'Moves the step on', + inputSchema: z.object({}), + }).server(() => { + step += 1 + return 'ticked' + }) + const plugin = definePlugin({ + name: 'test/placement', + setup: () => ({ + tools: [tick], + prompts: [ + 'Be brief.', + { + id: 'test/placement:step', + text: () => `Current step: ${step}`, + placement: 'context', + }, + ], + }), + }) + const { adapter, calls } = mockAdapter([ + () => toolCall('tick', {}, 'c1'), + () => text('done'), + ]) + const harness = defineHarness({ + name: 'test/placement', + adapter, + plugins: () => [plugin], + }) + const host = createHarnessHost({ persistence: memoryPersistence() }) + const session = await host.open(harness, { threadId: 't1' }) + + await session.send('go') + + expect(calls).toHaveLength(2) + const [first, second] = calls + // The default placement stays in the system prompts. The context prompt + // is the last system prompt, not a user message. + expect(first.systemPrompts).toContain('Be brief.') + expect(lastSystem(first)).toBe('Current step: 0') + expect(lastSystem(second)).toBe('Current step: 1') + expect(JSON.stringify(first.messages)).not.toContain('Current step') + // The second call replaces the prompt of the first. The system prompts + // before it are byte for byte the same. + expect(second.systemPrompts).toHaveLength(first.systemPrompts.length) + expect(JSON.stringify(second.systemPrompts.slice(0, -1))).toBe( + JSON.stringify(first.systemPrompts.slice(0, -1)), + ) + // The log does not keep the context prompt. + expect(JSON.stringify(await session.transcript())).not.toContain( + 'Current step', + ) + await host.close() + }) + + it('drops the context prompt and the extension items of a plugin that the thread removes', async () => { + const Rules = createExtensionPoint('test/placement-rules') + const notes = definePlugin({ + name: 'test/notes', + setup: () => ({ + prompts: [ + { id: 'test/notes:on', text: 'Notes on.', placement: 'context' }, + ], + contribute: [Rules.item('notes rule')], + }), + }) + // Reads the rules at each model call, after setup: the cached list. + const reader = definePlugin({ + name: 'test/reader', + setup: (ctx) => ({ + prompts: [ + { + id: 'test/reader:rules', + text: () => `Rules: ${ctx.collect(Rules).join(', ') || 'none'}`, + placement: 'context', + }, + ], + }), + }) + const { adapter, calls } = mockAdapter(() => text('ok')) + const harness = defineHarness({ + name: 'test/placement-remove', + adapter, + plugins: () => [notes, reader], + }) + const host = createHarnessHost({ persistence: memoryPersistence() }) + const session = await host.open(harness, { threadId: 't1' }) + + await session.send('One.') + await session.configure({ + settings: { plugins: { remove: ['test/notes'] } }, + }) + await session.send('Two.') + + expect(calls).toHaveLength(2) + expect(lastSystem(calls[0])).toBe('Notes on.\n\nRules: notes rule') + // The removed plugin adds no context prompt, and `collect` does not give + // its items from the list that it cached before the change. + expect(lastSystem(calls[1])).toBe('Rules: none') + expect(JSON.stringify(calls[1].systemPrompts)).not.toContain('Notes on.') + await host.close() + }) +}) diff --git a/packages/ai-harness/tests/protocol.test.ts b/packages/ai-harness/tests/protocol.test.ts index ae7f843671..36dbaea0e1 100644 --- a/packages/ai-harness/tests/protocol.test.ts +++ b/packages/ai-harness/tests/protocol.test.ts @@ -2,8 +2,15 @@ import { describe, expect, it, vi } from 'vitest' import { z } from 'zod' import { defineAgent } from '@tanstack/ai' import { memoryPersistence } from '@tanstack/ai-persistence' -import { createHarnessHost, defineHarness } from '../src' import { + createHarnessHost, + defineCommand, + defineHarness, + definePlugin, + logIdOf, +} from '../src' +import { + applyInput, createHarnessHandler, handleHarnessSocket, parseControlFrame, @@ -18,9 +25,21 @@ import { text, wakeMessages, } from './helpers' +import { + HARNESS_PROTOCOL_VERSION, + requestIdOf, + versionError, +} from '../src/protocol' +import type { Receipt } from '../src' import type { HostFrame } from '../src/server' import type { Reply } from './helpers' +// Coverage runs in CI are slow, so polled checks wait up to 5 s. +const waitFor = ( + check: () => T | Promise, + options?: { timeout?: number; interval?: number }, +) => vi.waitFor(check, { timeout: 5_000, ...options }) + const pricer = defineAgent({ name: 'pricer', description: 'Prices a vendor', @@ -162,7 +181,15 @@ describe('HTTP handler', () => { ), isTurnEnd, ) - expect(first.every((frame) => frame.type === 'harness.event')).toBe(true) + // The hello comes first, then the events. + expect(first[0]).toEqual({ + type: 'harness.hello', + v: 2, + threadId: 'user-1-t', + }) + expect( + first.slice(1).every((frame) => frame.type === 'harness.event'), + ).toBe(true) const last = first.at(-1) const cursor = last?.type === 'harness.event' ? last.cursor : '' @@ -176,11 +203,13 @@ describe('HTTP handler', () => { isTurnEnd, ) expect( - second.every( - (frame) => - frame.type === 'harness.event' && - Number(frame.cursor) > Number(cursor), - ), + second + .slice(1) + .every( + (frame) => + frame.type === 'harness.event' && + Number(frame.cursor) > Number(cursor), + ), ).toBe(true) expect(JSON.stringify(second)).toContain('again') await host.close() @@ -195,7 +224,7 @@ describe('HTTP handler', () => { expect(first).toMatchObject({ inputId: 'c-1', status: 'accepted' }) expect(retry).toEqual(first) - await vi.waitFor(() => expect(calls).toHaveLength(1)) + await waitFor(() => expect(calls).toHaveLength(1)) expect( await control({ op: 'prompt', message: 'another', inputId: 'c-1' }), ).toMatchObject({ status: 'rejected', reason: 'conflict' }) @@ -205,6 +234,88 @@ describe('HTTP handler', () => { await host.close() }) + it('runs a resent command or agent start with the same inputId once', async () => { + const runs = { command: 0, agent: 0 } + const counter = defineAgent({ + name: 'counter', + description: 'Counts its runs', + run: async () => ++runs.agent, + }) + const harness = defineHarness({ + name: 'test/protocol-resend', + adapter: mockAdapter([]).adapter, + agents: [counter], + plugins: () => [ + definePlugin({ + name: 'test/count', + setup: () => ({ + commands: { + count: defineCommand({ + description: 'Counts its runs', + run: () => ++runs.command, + }), + }, + }), + }), + ], + expose: { agents: ['counter'], commands: ['count'] }, + }) + const host = createHarnessHost({ persistence: memoryPersistence() }) + const session = await host.open(harness, { threadId: 'resend' }) + + // The operation of a receipt ends, also as a duplicate. + const settled = (receipt: Receipt) => + receipt.operationId === undefined + ? undefined + : session.operation(receipt.operationId)?.then( + () => {}, + () => {}, + ) + const command = { op: 'command', name: 'count', inputId: 'c-1' } as const + const first = await applyInput({ harness, session, input: command }) + expect(first).toMatchObject({ inputId: 'c-1', status: 'accepted' }) + await vi.waitFor(() => expect(runs.command).toBe(1), { timeout: 5000 }) + await settled(await applyInput({ harness, session, input: command })) + + const agent = { op: 'agent', agent: 'counter', inputId: 'a-1' } as const + expect(await applyInput({ harness, session, input: agent })).toMatchObject({ + inputId: 'a-1', + status: 'accepted', + }) + await vi.waitFor(() => expect(runs.agent).toBe(1), { timeout: 5000 }) + await settled(await applyInput({ harness, session, input: agent })) + + expect(runs).toEqual({ command: 1, agent: 1 }) + await host.close() + }) + + it('writes the records of a resent append with the same inputId once', async () => { + const persistence = memoryPersistence() + const harness = defineHarness({ + name: 'test/protocol-append', + adapter: mockAdapter([]).adapter, + expose: { records: ['app.signal'] }, + }) + const host = createHarnessHost({ persistence }) + const session = await host.open(harness, { threadId: 'append' }) + const input = { + op: 'append' as const, + records: [{ type: 'app.signal', text: 'The build failed.' }], + inputId: 'ap-1', + } + + const first = await applyInput({ harness, session, input }) + const again = await applyInput({ harness, session, input }) + + expect(first).toEqual({ inputId: 'ap-1', status: 'accepted' }) + expect(again).toEqual(first) + const signals = ( + await persistence.stores.log.read({ logId: logIdOf({ logId: 'append' }) }) + ).filter(({ record }) => record.type === 'app.signal') + expect(signals).toHaveLength(1) + await host.close() + }) + it('runs exposed agents from a client and refuses the others', async () => { const { handler, host } = setup() const control = controlOf(handler, 'user-1-a') @@ -290,7 +401,7 @@ describe('HTTP handler', () => { ).json() await control({ op: 'prompt', message: 'hi' }) // The model call runs, so a steer can join only at a later call. - await vi.waitFor(() => expect(calls).toHaveLength(1)) + await waitFor(() => expect(calls).toHaveLength(1)) await control({ op: 'followUp', message: 'later', inputId: 'f-1' }) expect( @@ -317,12 +428,12 @@ describe('HTTP handler', () => { await control({ op: 'followUp', message: 'I am real', wake: true }) - await vi.waitFor(() => expect(calls).toHaveLength(1)) + await waitFor(() => expect(calls).toHaveLength(1)) const session = await host.open(harness, { threadId: 'user-1-wake', principal: user, }) - await vi.waitFor(() => expect(session.snapshot().status).toBe('idle')) + await waitFor(() => expect(session.snapshot().status).toBe('idle')) expect(wakeMessages(await session.transcript())).toEqual([]) await host.close() }) @@ -357,6 +468,50 @@ describe('parseHarnessInput', () => { ).toThrow("Invalid input: setDelivery needs 'steer' or 'queue'.") }) + it('refuses a message part without the fields the harness reads', () => { + // An image with no source, then a data source with no MIME type. + for (const part of [ + { type: 'image' }, + { type: 'image', source: { type: 'data', value: 'aGVsbG8=' } }, + ]) { + expect(() => + parseHarnessInput({ op: 'prompt', message: [part] }), + ).toThrow( + 'Invalid input: prompt needs a message: text, or content parts.', + ) + } + expect(() => + parseHarnessInput({ + op: 'continue', + ephemeral: [{ role: 'user', content: [{ type: 'audio' }] }], + }), + ).toThrow('Invalid input: ephemeral must be an array of user messages.') + }) + + it('reports the WebSocket in capabilities only when the app serves it', async () => { + const { host, harness } = setup() + const transportOf = async (socket?: boolean) => { + const handler = createHarnessHandler({ + host, + harness, + ...(socket === undefined ? {} : { socket }), + }) + const response = await handler( + new Request('http://x/api/harness/capabilities', { headers: auth }), + ) + const body: unknown = await response.json() + return body + } + + expect(await transportOf()).toMatchObject({ + transport: { streaming: true, websocket: false }, + }) + expect(await transportOf(true)).toMatchObject({ + transport: { streaming: true, websocket: true }, + }) + await host.close() + }) + it('checks the ephemeral messages of prompt and continue', () => { const refused = 'Invalid input: ephemeral must be an array of user messages.' @@ -396,10 +551,10 @@ describe('WebSocket', () => { requestId: 'r0', input: { op: 'prompt', message: 'x' }, }) - await vi.waitFor(() => expect(sent.at(-1)?.type).toBe('harness.error')) + await waitFor(() => expect(sent.at(-1)?.type).toBe('harness.error')) receive({ type: 'harness.subscribe', threadId: 'user-1-ws' }) - await vi.waitFor(() => + await waitFor(() => expect(sent.some((frame) => frame.type === 'harness.hello')).toBe(true), ) receive({ @@ -408,7 +563,7 @@ describe('WebSocket', () => { input: { op: 'prompt', message: 'hi' }, }) - await vi.waitFor(() => + await waitFor(() => expect( sent.some( (frame) => @@ -425,7 +580,7 @@ describe('WebSocket', () => { status: 'accepted', }) receive({ type: 'harness.snapshot' }) - await vi.waitFor(() => + await waitFor(() => expect(sent.at(-1)).toMatchObject({ type: 'harness.snapshot', snapshot: { status: 'idle' }, @@ -436,6 +591,53 @@ describe('WebSocket', () => { }) }) +describe('WebSocket frame order', () => { + it('applies the inputs of one socket in the order they came', async () => { + const { host, harness } = setup() + const { socket, sent, receive } = fakeSocket() + handleHarnessSocket({ + host, + harness, + socket, + request: new Request('http://x/api/harness/socket', { headers: auth }), + }) + receive({ type: 'harness.subscribe', threadId: 'user-1-order' }) + await vi.waitFor( + () => + expect(sent.some((frame) => frame.type === 'harness.hello')).toBe(true), + { timeout: 5000 }, + ) + + // The first input stores its image before it reaches the session. The + // second is text only, so it is faster on its own. + const image = { + type: 'image', + source: { type: 'data', value: 'aGVsbG8=', mimeType: 'image/png' }, + } + receive({ + type: 'harness.input', + requestId: 'first', + input: { op: 'prompt', inputId: 'first', message: [image] }, + }) + receive({ + type: 'harness.input', + requestId: 'second', + input: { op: 'prompt', inputId: 'second', message: 'B' }, + }) + + const receipts = () => + sent.flatMap((frame) => + frame.type === 'harness.receipt' ? [frame.requestId] : [], + ) + await vi.waitFor(() => expect(receipts()).toHaveLength(2), { + timeout: 5000, + }) + expect(receipts()).toEqual(['first', 'second']) + socket.close() + await host.close() + }) +}) + describe('parseControlFrame', () => { it('rejects unknown frames', () => { expect(() => parseControlFrame('{"type":"nope"}')).toThrow('unknown type') @@ -443,4 +645,94 @@ describe('parseControlFrame', () => { parseControlFrame('{"type":"harness.input","requestId":"1","input":{}}'), ).toThrow('Invalid input') }) + + it('keeps the v of a subscribe, and refuses a v that is not a number', () => { + expect( + parseControlFrame('{"type":"harness.subscribe","threadId":"t","v":2}'), + ).toEqual({ type: 'harness.subscribe', threadId: 't', v: 2 }) + expect(() => + parseControlFrame('{"type":"harness.subscribe","threadId":"t","v":"2"}'), + ).toThrow('v must be a number') + }) + + it('keeps a valid traceparent of an input, and drops an invalid one', () => { + const traceparent = + '00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01' + const frame = (value: string) => + parseControlFrame( + JSON.stringify({ + type: 'harness.input', + requestId: 'r1', + input: { op: 'prompt', message: 'hi' }, + traceparent: value, + }), + ) + expect(frame(traceparent)).toMatchObject({ traceparent }) + expect(frame('not-a-traceparent')).not.toHaveProperty('traceparent') + }) +}) + +describe('versionError', () => { + it('takes the current version, and no version', () => { + expect(HARNESS_PROTOCOL_VERSION).toBe(2) + expect(versionError(undefined)).toBeUndefined() + expect(versionError(2)).toBeUndefined() + }) + + it('refuses another version', () => { + expect(versionError(1)).toEqual({ + message: + 'Protocol version 1 is not supported. This host speaks version 2.', + code: 'unsupported_version', + }) + }) +}) + +describe('requestIdOf', () => { + it('reads the requestId of a frame that fails the full parse', () => { + const badInput = '{"type":"harness.input","requestId":"r9","input":{}}' + expect(() => parseControlFrame(badInput)).toThrow('Invalid input') + expect(requestIdOf(badInput)).toBe('r9') + }) + + it('never throws', () => { + expect(requestIdOf('{not json')).toBeUndefined() + expect(requestIdOf('null')).toBeUndefined() + expect(requestIdOf('{"requestId":7}')).toBeUndefined() + }) +}) + +describe('traceparent on an input', () => { + const traceparent = '00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01' + + it('applyInput passes the traceparent to the send of a message', async () => { + const { host, harness } = setup() + const session = await host.open(harness, { + threadId: 'user-1-trace', + principal: user, + }) + const send = vi.spyOn(session, 'send') + + await applyInput({ + harness, + session, + input: { op: 'prompt', message: 'hi' }, + principal: user, + traceparent, + }) + await applyInput({ + harness, + session, + input: { op: 'steer', message: 'more' }, + principal: user, + traceparent, + }) + + expect(send.mock.calls.map(([, options]) => options?.traceparent)).toEqual([ + traceparent, + traceparent, + ]) + await session.close() + await host.close() + }) }) diff --git a/packages/ai-harness/tests/redaction.test.ts b/packages/ai-harness/tests/redaction.test.ts index 3e1e115727..0961ffc350 100644 --- a/packages/ai-harness/tests/redaction.test.ts +++ b/packages/ai-harness/tests/redaction.test.ts @@ -1,3 +1,6 @@ +import { mkdtemp, readFile, readdir, rm } from 'node:fs/promises' +import { tmpdir } from 'node:os' +import { join } from 'node:path' import { memoryCredentialStore } from '@tanstack/ai/byok/server' import { describe, expect, it } from 'vitest' import { z } from 'zod' @@ -15,10 +18,13 @@ import { import { Redactor, redactMiddleware } from '../src/auth' import { createHarnessHandler } from '../src/server' import { scrubbedBy } from '../src/session/harness-session' +import { hostBackend } from '../src/first-party/coding/backend' +import { workspaceTools } from '../src/first-party/coding/workspace' import { mockAdapter, text, toolCall } from './helpers' import type { AnyTool, ChatMiddlewareConfig } from '@tanstack/ai' import type { Reply } from './helpers' -import type { HarnessRunContext } from '../src' +import type { WorkspaceBackend } from '../src/first-party/coding/backend' +import type { HarnessPlugin, HarnessRunContext } from '../src' const THREAD = 't1' const LOG = logIdOf({ logId: THREAD }) @@ -38,10 +44,12 @@ const leaked = createPluginEvent<{ key: string }>('test/leaked') async function open({ fail, tools = [], + plugins = [], replies = [() => toolCall('leak', {}), () => text('done')], }: { fail: boolean tools?: Array + plugins?: Array replies?: Array }) { const credentials = memoryCredentialStore() @@ -80,7 +88,7 @@ async function open({ name: 'test/redaction', adapter: model.adapter, tools, - plugins: () => [leaky], + plugins: () => [leaky, ...plugins], }) const session = await host.open(harness, { threadId: THREAD, @@ -130,6 +138,46 @@ describe('the redaction sink of a session', () => { ) }) +describe('a bash spill', () => { + it('keeps no key of the session in the file of the full output', async () => { + const root = await mkdtemp(join(tmpdir(), 'tanstack-spill-')) + try { + // Long output with the key: the model gets only its end, and the + // file gets all of it. + const backend: WorkspaceBackend = { + ...hostBackend, + exec: async () => ({ + exitCode: 0, + stdout: `${KEY}\n`.repeat(3000), + stderr: '', + }), + } + const { host, session } = await open({ + fail: false, + plugins: [ + workspaceTools({ root, backend, spillDir: 'spill', web: false }), + ], + replies: [ + () => toolCall('leak', {}, 'c0'), + () => toolCall('bash', { command: 'print-key' }, 'c1'), + () => text('done'), + ], + }) + + await session.send('Leak it twice.') + + const [file] = await readdir(join(root, 'spill')) + if (file === undefined) throw new Error('No spill file.') + const saved = await readFile(join(root, 'spill', file), 'utf8') + expect(saved).toContain('[redacted]') + expect(saved).not.toContain(KEY_PART) + await host.close() + } finally { + await rm(root, { recursive: true, force: true }) + } + }) +}) + describe('the edges of a session', () => { it.each([ { fail: false, what: 'result' }, diff --git a/packages/ai-harness/tests/remote-hooks.test.ts b/packages/ai-harness/tests/remote-hooks.test.ts new file mode 100644 index 0000000000..d0ac41fa14 --- /dev/null +++ b/packages/ai-harness/tests/remote-hooks.test.ts @@ -0,0 +1,144 @@ +import { describe, expect, it, vi } from 'vitest' +import { z } from 'zod' +import { EventType, toolDefinition } from '@tanstack/ai' +import { memoryLogBackend, memoryLogStore } from '@tanstack/ai-persistence' +import { + createHarnessHost, + defineHarness, + definePlugin, + durableTool, +} from '../src' +import { mockAdapter, text, toolCall } from './helpers' +import type { AnyTextAdapter } from '@tanstack/ai' +import type { LogStore } from '@tanstack/ai-persistence' +import type { LiveEvent, SessionEvent } from '../src' + +// Coverage runs in CI are slow, so polled checks wait up to 5 s. +const waitFor = (check: () => T | Promise) => + vi.waitFor(check, { timeout: 5_000 }) + +const THREAD = 't1' + +const probe = toolDefinition({ + name: 'probe', + description: 'Asks a question.', + inputSchema: z.object({}), +}) + +/** A harness whose tool asks a question. Both hosts run the same one. */ +function harnessWith(adapter: AnyTextAdapter) { + const asker = definePlugin({ + name: 'test/asker', + setup: () => ({ + tools: [ + durableTool(probe, async (_args, { step }) => ({ + sure: await step.ask({ name: 'sure', message: 'Sure?' }), + })), + ], + }), + }) + return defineHarness({ + name: 'test/remote-hooks', + adapter, + plugins: () => [asker], + }) +} + +/** One model call that answers with usage, so the log has `harness.usage`. */ +const withUsage = (content: string) => + text(content).map((chunk) => + chunk.type === EventType.RUN_FINISHED + ? { + ...chunk, + usage: { promptTokens: 1, completionTokens: 2, totalTokens: 3 }, + } + : chunk, + ) + +/** A host on `log` (one server), with the session of `THREAD` open. */ +async function hostOn(log: LogStore, adapter: AnyTextAdapter) { + const host = createHarnessHost({ persistence: { stores: { log } } }) + const session = await host.open(harnessWith(adapter), { threadId: THREAD }) + const statuses: Array = [] + const reader = new AbortController() + void (async () => { + for await (const event of host.events({ signal: reader.signal })) { + if (event.type === 'status' && event.threadId === THREAD) + statuses.push(event.status) + } + })() + const close = async () => { + reader.abort() + await session.close() + await host.close() + } + return { host, session, statuses, close } +} + +const isFinished = (event: LiveEvent['event']) => + event.type === EventType.CUSTOM && event.name === 'harness.operation.finished' + +/** The `CUSTOM` names of `events`, with each name once per event. */ +const customNames = (events: ReadonlyArray) => + events.flatMap(({ event }) => + event.type === EventType.CUSTOM ? [event.name] : [], + ) + +describe('hooks on a host that does not own the log', () => { + it('fire once for each event that the owner writes, and the status follows', async () => { + const backend = memoryLogBackend() + const owner = await hostOn( + memoryLogStore({ backend }), + mockAdapter([() => toolCall('probe', {}), () => withUsage('Done.')]) + .adapter, + ) + const other = await hostOn( + memoryLogStore({ backend }), + mockAdapter([]).adapter, + ) + const from = owner.session.snapshot().cursor + const seen: Array = [] + const own: Array = [] + const questions: Array = [] + const usages: Array = [] + other.session.onEvent((entry) => void seen.push(entry)) + owner.session.onEvent((entry) => void own.push(entry)) + other.session.onQuestion(({ question }) => questions.push(question.message)) + other.session.onUsage(({ usage }) => usages.push(usage.totalTokens)) + + const turn = owner.session.send('Go.') + await waitFor(() => + expect(owner.session.snapshot().pendingQuestions).toHaveLength(1), + ) + await waitFor(() => expect(questions).toEqual(['Sure?'])) + await waitFor(() => expect(other.statuses.at(-1)).toBe('waiting')) + const [question] = owner.session.snapshot().pendingQuestions + await owner.session.answer({ + questionId: question?.questionId ?? '', + value: 'yes', + }) + await turn + await waitFor(() => + expect(seen.some(({ event }) => isFinished(event))).toBe(true), + ) + + // The other host saw each logged event once, in log order. + const logged: Array = [] + for await (const entry of other.session.events({ from })) { + logged.push(entry) + if (isFinished(entry.event)) break + } + expect(seen).toEqual( + logged.map(({ operationId, event }) => ({ operationId, event })), + ) + expect(questions).toEqual(['Sure?']) + expect(usages).toEqual([3]) + const statuses = ['idle', 'running', 'waiting', 'running', 'idle'] + expect(other.statuses).toEqual(statuses) + expect(owner.statuses).toEqual(statuses) + // The owner saw its own events at publish, and not again from the log. + expect(customNames(own)).toEqual(customNames(seen)) + await other.close() + await owner.close() + }) +}) diff --git a/packages/ai-harness/tests/remote-worker.test.ts b/packages/ai-harness/tests/remote-worker.test.ts index 5b3142e8b2..4339caeb6b 100644 --- a/packages/ai-harness/tests/remote-worker.test.ts +++ b/packages/ai-harness/tests/remote-worker.test.ts @@ -8,6 +8,12 @@ import { runHarnessWorker } from '../src/worker' import { mockAdapter, text } from './helpers' import type { StreamChunk } from '@tanstack/ai' +// Coverage runs in CI are slow, so polled checks wait up to 5 s. +const waitFor = ( + check: () => T | Promise, + options?: { timeout?: number; interval?: number }, +) => vi.waitFor(check, { timeout: 5_000, ...options }) + async function collect(stream: AsyncIterable) { const chunks: Array = [] for await (const chunk of stream) chunks.push(chunk) @@ -163,29 +169,37 @@ describe('runHarnessWorker edges', () => { send('') send({ type: 'harness.snapshot' }) - await vi.waitFor(() => + await waitFor(() => expect(frames.at(-1)).toEqual({ type: 'harness.error', message: 'Send harness.subscribe first.', + code: 'invalid_frame', + requestId: expect.any(String), }), ) send('{not json') - await vi.waitFor(() => expect(frames).toHaveLength(2)) - expect(frames[1].type).toBe('harness.error') + await waitFor(() => expect(frames).toHaveLength(2)) + expect(frames[1]).toMatchObject({ + type: 'harness.error', + code: 'invalid_frame', + requestId: expect.any(String), + }) send({ type: 'harness.subscribe', threadId: 'w', from: '0' }) - await vi.waitFor(() => + await waitFor(() => expect(frames.some((frame) => frame.type === 'harness.hello')).toBe(true), ) send({ type: 'harness.subscribe', threadId: 'w' }) - await vi.waitFor(() => + await waitFor(() => expect(frames.at(-1)).toEqual({ type: 'harness.error', message: 'Already subscribed.', + code: 'invalid_frame', + requestId: expect.any(String), }), ) send({ type: 'harness.snapshot' }) - await vi.waitFor(() => + await waitFor(() => expect(frames.at(-1)).toMatchObject({ type: 'harness.snapshot', snapshot: { threadId: 'w', status: 'idle' }, @@ -194,4 +208,60 @@ describe('runHarnessWorker edges', () => { input.end() await running }) + + it('names the requestId of a frame with a bad input', async () => { + const { adapter } = mockAdapter([]) + const input = new PassThrough() + const frames: Array = [] + const running = runHarnessWorker( + defineHarness({ name: 'test/worker-request-id', adapter }), + { + input, + output: { write: (line: string) => frames.push(JSON.parse(line)) }, + persistence: memoryPersistence(), + }, + ) + input.write( + `${JSON.stringify({ type: 'harness.input', requestId: 'r7', input: {} })} +`, + ) + await waitFor(() => + expect(frames.at(-1)).toMatchObject({ + type: 'harness.error', + code: 'invalid_frame', + requestId: 'r7', + }), + ) + input.end() + await running + }) + + it('refuses a protocol version that it does not speak, and stops', async () => { + const { adapter } = mockAdapter([]) + const input = new PassThrough() + const frames: Array = [] + const running = runHarnessWorker( + defineHarness({ name: 'test/worker-version', adapter }), + { + input, + output: { write: (line: string) => frames.push(JSON.parse(line)) }, + persistence: memoryPersistence(), + }, + ) + input.write( + `${JSON.stringify({ type: 'harness.subscribe', threadId: 'w', v: 1, requestId: 's1' })} +`, + ) + // The worker ends without stdin closing. + await running + expect(frames).toEqual([ + { + type: 'harness.error', + message: + 'Protocol version 1 is not supported. This host speaks version 2.', + code: 'unsupported_version', + requestId: 's1', + }, + ]) + }) }) diff --git a/packages/ai-harness/tests/resume-pending.test.ts b/packages/ai-harness/tests/resume-pending.test.ts index e46a01e8f0..4d988c100f 100644 --- a/packages/ai-harness/tests/resume-pending.test.ts +++ b/packages/ai-harness/tests/resume-pending.test.ts @@ -1,7 +1,12 @@ import { describe, expect, it, vi } from 'vitest' import { defineAgent } from '@tanstack/ai' -import { UNOWNED_EPOCH, memoryLogStore } from '@tanstack/ai-persistence' +import { + UNOWNED_EPOCH, + memoryLogBackend, + memoryLogStore, +} from '@tanstack/ai-persistence' import { createHarnessHost, defineHarness, definePlugin, logIdOf } from '../src' +import { OPEN_RETRY_MS } from '../src/ownership' import { after, canOtherHostTake, @@ -357,7 +362,10 @@ describe('host.resumePending', () => { }) it('goes on after a thread that fails to open, and names the threads it opened', async () => { - const log = memoryLogStore() + let skew = 0 + const log = memoryLogStore({ + backend: memoryLogBackend({ now: () => Date.now() + skew }), + }) // The failing thread expired first, so it comes first. await stopOwner(log, logIdOf({ logId: 't-bad' })) await seedStoppedTurn(log) @@ -393,6 +401,15 @@ describe('host.resumePending', () => { }), ) expect(calls).toHaveLength(1) + // The log of the thread that failed goes back with a wake, so a later + // sweep lists it again. + const listed = async () => + (await log.listOrphaned({ limit: 10, harnesses: [NAME] })).map( + (orphan) => orphan.logId, + ) + expect(await listed()).not.toContain(logIdOf({ logId: 't-bad' })) + skew = OPEN_RETRY_MS + expect(await listed()).toContain(logIdOf({ logId: 't-bad' })) await host.close() }) diff --git a/packages/ai-harness/tests/retention.test.ts b/packages/ai-harness/tests/retention.test.ts index 7f70304654..f461b6d7af 100644 --- a/packages/ai-harness/tests/retention.test.ts +++ b/packages/ai-harness/tests/retention.test.ts @@ -4,6 +4,7 @@ import { memoryPersistence } from '@tanstack/ai-persistence' import { createHarnessHost, defineHarness } from '../src' import { mockAdapter, text } from './helpers' import type { HarnessPersistence, RetentionLimits } from '../src' +import type { SessionImpl } from '../src/session/harness-session' /** * A session with `retention`. The model answers `done` to every call, and @@ -109,7 +110,9 @@ describe('retention', () => { settledOperations: 2, settledAgentRuns: 1, }) - await session.send('Hi', { runId: 'run-0' }) + await (session as unknown as SessionImpl).sendForRun('Hi', { + runId: 'run-0', + }) for (let index = 0; index < 2; index += 1) await session.send('Hi') expect(session.operation('run-0')).toBe(undefined) // Both windows keep 3 operations. @@ -117,7 +120,9 @@ describe('retention', () => { expect(privateSetSize(session, 'retiredIds')).toBe(3) // The log keeps every run, so the id stays taken. - const reused = session.send('Hi', { runId: 'run-0' }) + const reused = (session as unknown as SessionImpl).sendForRun('Hi', { + runId: 'run-0', + }) expect(await reused.receipt).toMatchObject({ status: 'rejected', reason: 'conflict', diff --git a/packages/ai-harness/tests/ring.test.ts b/packages/ai-harness/tests/ring.test.ts new file mode 100644 index 0000000000..2f7fd84615 --- /dev/null +++ b/packages/ai-harness/tests/ring.test.ts @@ -0,0 +1,116 @@ +import { describe, expect, it } from 'vitest' +import { EventType } from '@tanstack/ai' +import { EventRing } from '../src/ring' +import type { SessionEvent } from '../src/types' + +function event(cursor: number): SessionEvent { + return { + cursor: String(cursor), + threadId: 't', + operationId: 'op-1', + event: { type: EventType.RUN_STARTED, threadId: 't', runId: 'r' }, + } +} + +/** Push events with the cursors and threads given, each `size` bytes. */ +function ringOf( + maxBytes: number, + items: Array<[cursor: number, thread: string]>, + size = 10, +) { + const ring = new EventRing({ maxBytes }) + for (const [cursor, thread] of items) { + ring.push({ thread, entry: event(cursor), size }) + } + return ring +} + +const cursors = (read: { entries: Array }) => + read.entries.map((entry) => Number(entry.cursor)) + +describe('EventRing', () => { + it('reads the events after a cursor, also with gaps between cursors', () => { + const ring = ringOf(1000, [ + [2, 'a'], + [5, 'a'], + [9, 'a'], + ]) + expect(cursors(ring.after(0))).toEqual([2, 5, 9]) + expect(cursors(ring.after(2))).toEqual([5, 9]) + expect(cursors(ring.after(3))).toEqual([5, 9]) + expect(cursors(ring.after(9))).toEqual([]) + expect(ring.after(0).gap).toBe(false) + }) + + it('drops the oldest events over the byte cap', () => { + const ring = ringOf(25, [ + [1, 'a'], + [2, 'a'], + [3, 'a'], + [4, 'a'], + ]) + expect(cursors(ring.after(0))).toEqual([3, 4]) + expect(ring.bytes).toBe(20) + expect(ring.length).toBe(2) + expect(ring.from).toBe(2) + }) + + it('keeps the newest event when it alone is over the cap', () => { + const ring = ringOf(5, [ + [1, 'a'], + [2, 'a'], + ]) + expect(cursors(ring.after(0))).toEqual([2]) + expect(ring.bytes).toBe(10) + }) + + it('reports a gap for a cursor older than the oldest kept event', () => { + const ring = ringOf(20, [ + [1, 'a'], + [2, 'a'], + [3, 'a'], + ]) + // Event 1 is gone: a reader at 0 missed it, a reader at 1 did not. + expect(ring.after(0)).toMatchObject({ gap: true }) + expect(ring.after(1)).toMatchObject({ gap: false }) + expect(cursors(ring.after(1))).toEqual([2, 3]) + }) + + it('reports a gap before the start cursor', () => { + const ring = new EventRing({ maxBytes: 100, from: 7 }) + ring.push({ thread: 'a', entry: event(8), size: 1 }) + expect(ring.after(6).gap).toBe(true) + expect(ring.after(7)).toEqual({ entries: [event(8)], gap: false }) + }) + + it('reads the events of one thread only, and drops them in order', () => { + const ring = ringOf(40, [ + [1, 'a'], + [2, 'b'], + [3, 'a'], + [4, 'b'], + [5, 'a'], + [6, 'b'], + ]) + // Events 1 and 2 are gone. + expect(cursors(ring.after(0, 'a'))).toEqual([3, 5]) + expect(cursors(ring.after(3, 'a'))).toEqual([5]) + expect(cursors(ring.after(0, 'b'))).toEqual([4, 6]) + expect(cursors(ring.after(0, 'c'))).toEqual([]) + }) + + it('keeps the order through many pushes and drops', () => { + const ring = new EventRing({ maxBytes: 30 }) + for (let cursor = 1; cursor <= 1000; cursor++) { + ring.push({ + thread: cursor % 2 ? 'a' : 'b', + entry: event(cursor), + size: 10, + }) + } + expect(cursors(ring.after(0))).toEqual([998, 999, 1000]) + expect(cursors(ring.after(0, 'a'))).toEqual([999]) + expect(cursors(ring.after(0, 'b'))).toEqual([998, 1000]) + expect(ring.from).toBe(997) + }) +}) diff --git a/packages/ai-harness/tests/routing.test.ts b/packages/ai-harness/tests/routing.test.ts index ad752e0bb9..282c689793 100644 --- a/packages/ai-harness/tests/routing.test.ts +++ b/packages/ai-harness/tests/routing.test.ts @@ -805,3 +805,54 @@ describe('routing and ownership', () => { await host.close() }) }) + +describe('routing.router model', () => { + it('counts the model calls of context.model.chat under router, in the turn', async () => { + const billing = textAgent('billing', 'Billing here') + const { host, session, log, main } = await openRouting({ + router: async ({ model, messages }) => { + const { text: answer } = await model.chat({ + options: { messages }, + label: 'pick', + }) + return answer === 'billing' ? 'billing' : undefined + }, + agents: [billing.agent], + replies: [ + () => + text('billing').map((chunk) => + chunk.type === EventType.RUN_FINISHED + ? { + ...chunk, + usage: { + promptTokens: 7, + completionTokens: 3, + totalTokens: 10, + }, + } + : chunk, + ), + ], + }) + + const operation = session.send('my invoice is wrong') + const turn = await operation + + expect(turn.text).toBe('billing:\nBilling here') + expect(messageTexts(main.calls[0])).toEqual(['my invoice is wrong']) + const usage = session.usage() + expect(usage.total).toMatchObject({ calls: 1, totalTokens: 10 }) + expect(usage.bySource).toEqual({ router: usage.total }) + const records = (await log.read({ logId: logIdOf({ logId: THREAD }) })) + .map((entry) => entry.record) + .filter((record) => record.type === 'harness.usage') + expect(records).toEqual([ + expect.objectContaining({ + source: 'router', + label: 'pick', + runId: operation.id, + }), + ]) + await host.close() + }) +}) diff --git a/packages/ai-harness/tests/run-route-hosts.test.ts b/packages/ai-harness/tests/run-route-hosts.test.ts index e0fde6b65f..92fa24ff76 100644 --- a/packages/ai-harness/tests/run-route-hosts.test.ts +++ b/packages/ai-harness/tests/run-route-hosts.test.ts @@ -6,20 +6,32 @@ import { } from '@tanstack/ai-persistence' import { createHarnessHost, defineHarness } from '../src' import { createHarnessHandler } from '../src/server' -import { gate, mockAdapter, text } from './helpers' +import { appendToLog, gate, mockAdapter, text, wrapLogStore } from './helpers' import type { Reply } from './helpers' +import type { SnapshotTiming } from '../src/log' /** * Two server instances over the same stores, with no sticky routing: a run * that starts on host A must be joinable and resumable from host B. The two * log stores are two connections to one database. */ -function twoHosts() { +function twoHosts(options: { snapshot?: SnapshotTiming } = {}) { const backend = memoryLogBackend() const { metadata, sessions } = memoryPersistence().stores - const persistenceOf = () => ({ - stores: { log: memoryLogStore({ backend }), metadata, sessions }, - }) + /** The `after` of each log read of host B. */ + const readsOfB: Array = [] + const persistenceOf = (reads?: Array) => { + const inner = memoryLogStore({ backend }) + const log = reads + ? wrapLogStore(inner, { + read: (query) => { + reads.push(query.after ?? 0) + return inner.read(query) + }, + }) + : inner + return { stores: { log, metadata, sessions } } + } const release = gate() // RUN_STARTED, then the rest after `release`. By RUN_STARTED, host A holds // the run lease, so host B only reads. @@ -32,10 +44,13 @@ function twoHosts() { })() const { adapter, calls } = mockAdapter([reply]) const harness = defineHarness({ name: 'test/run-hosts', adapter }) - const handlerOf = () => { + const handlerOf = (reads?: Array) => { const host = createHarnessHost({ - persistence: persistenceOf(), - log: { coalesceMs: 0 }, + persistence: persistenceOf(reads), + log: { + coalesceMs: 0, + ...(options.snapshot ? { snapshot: options.snapshot } : {}), + }, access: { identify: ({ request }) => ({ id: request.headers.get('x-user') ?? 'user-1', @@ -48,14 +63,14 @@ function twoHosts() { const handler = createHarnessHandler({ host, harness, - media: { secret: 'shared' }, + media: { secret: 'a-shared-secret-of-32-bytes-or-more' }, }) const call = (path: string, init: RequestInit = {}) => handler(new Request(`http://h.test/api/${path}`, init)) return { host, call } } const a = handlerOf() - const b = handlerOf() + const b = handlerOf(readsOfB) const post = (on: typeof a, headers: Record = {}) => on.call('run', { method: 'POST', @@ -69,7 +84,7 @@ function twoHosts() { }), }) const close = () => Promise.all([a.host.close(), b.host.close()]) - return { a, b, post, release, calls, close } + return { a, b, post, release, calls, close, backend, readsOfB } } /** SSE events with their `id:` lines. */ @@ -128,6 +143,36 @@ describe('a run on a durable host, across instances', () => { await close() }) + it('reads the log on host B from the start of the run, not from the start of the log', async () => { + const { a, b, post, release, close, backend, readsOfB } = twoHosts({ + // Host A writes a snapshot at each write, so host B folds from it. + snapshot: { afterBytes: 1 }, + }) + // Records before the run, so the run does not start at the log start. + const history = Array.from({ length: 20 }, (_, n) => ({ + type: 'app.history', + n, + })) + await appendToLog( + memoryLogStore({ backend }), + { threadId: 'user-1-t' }, + history, + ) + const started = await readThenCut(await post(a)) + expect(started.at(-1)?.chunk.type).toBe('RUN_STARTED') + + const joined = await b.call('run?runId=run-a&threadId=user-1-t&offset=-1') + expect(joined.status).toBe(200) + release.open() + const events = withIds(await joined.text()) + expect(events.slice(0, started.length)).toEqual(started) + expect(events.at(-1)?.chunk.name).toBe('harness.operation.finished') + // No read of host B walks the history before the run. + expect(readsOfB.length).toBeGreaterThan(0) + expect(Math.min(...readsOfB)).toBeGreaterThanOrEqual(history.length) + await close() + }) + it('resumes a cut join on host B after Last-Event-ID, with no repeated event', async () => { const { a, b, post, release, close } = twoHosts() release.open() diff --git a/packages/ai-harness/tests/run-route.test.ts b/packages/ai-harness/tests/run-route.test.ts index 5f3e746102..f5d17fcfc4 100644 --- a/packages/ai-harness/tests/run-route.test.ts +++ b/packages/ai-harness/tests/run-route.test.ts @@ -588,6 +588,35 @@ describe('join a running turn', () => { await host.close() }) + it('replays the messages snapshot of a run that paused at an interrupt', async () => { + const { host, call, run } = setup( + approvalHarness([ + { toolCalls: [{ name: 'deploy', input: { env: 'prod' } }] }, + ]), + ) + const posted = sse( + await ( + await run({ + runId: 'run-a', + messages: [{ id: 'm-1', role: 'user', content: 'Deploy.' }], + }) + ).text(), + ) + expect( + posted.find((chunk) => chunk.type === 'RUN_FINISHED')?.outcome?.type, + ).toBe('interrupt') + + // A rejoin reads the log. ai-client rebuilds the paused run from the + // snapshot, so the log must keep it, before the interrupt. + const joined = withIds( + await (await call('run?runId=run-a&threadId=t-1&offset=-1')).text(), + ).map((event) => event.chunk.type) + const snapshotAt = joined.indexOf('MESSAGES_SNAPSHOT') + expect(snapshotAt).toBeGreaterThan(-1) + expect(snapshotAt).toBeLessThan(joined.lastIndexOf('RUN_FINISHED')) + await host.close() + }) + it('refuses a join of a thread the principal may not use', async () => { let allowed = true const { host, call, answer, posted } = await startAsking(() => allowed) diff --git a/packages/ai-harness/tests/run-tags.test.ts b/packages/ai-harness/tests/run-tags.test.ts new file mode 100644 index 0000000000..f51d063c98 --- /dev/null +++ b/packages/ai-harness/tests/run-tags.test.ts @@ -0,0 +1,207 @@ +import { describe, expect, it, vi } from 'vitest' +import { EventType, defineAgent } from '@tanstack/ai' +import { UNOWNED_EPOCH, memoryPersistence } from '@tanstack/ai-persistence' +import { createHarnessHost, defineHarness, logIdOf } from '../src' +import { createHarnessClient } from '../src/client' +import { createHarnessHandler, handleHarnessSocket } from '../src/server' +import { fakeSocket, mockAdapter, readSse } from './helpers' +import type { SessionEvent } from '../src/types' + +const BASE = 'http://x/api/harness' +const THREAD = 't1' + +// Coverage runs in CI are slow, so polled checks wait up to 5 s. +const waitFor = (check: () => T | Promise) => + vi.waitFor(check, { timeout: 5_000 }) + +/** What the tree reads of an event or a frame. */ +type Tagged = Pick + +const helper = defineAgent({ + name: 'helper', + description: 'Helps', + run: async () => 'helped', +}) + +/** Starts a helper of its own, so the tree has three levels. */ +const nester = defineAgent({ + name: 'nester', + description: 'Starts a helper', + run: async (ctx) => await ctx.harness.agents.run('helper'), +}) + +/** Runs a helper and a nester in a parallel group. */ +const lead = defineAgent({ + name: 'lead', + description: 'Runs a group', + run: async (ctx) => { + const results = await ctx.harness.agents.group({}, (g) => [ + g.start('helper'), + g.start('nester'), + ]) + return results.join(' ') + }, +}) + +const harness = defineHarness({ + name: 'test/run-tags', + adapter: mockAdapter([]).adapter, + agents: [lead, nester, helper], + // The helper of the nester is two levels below the lead. + subagents: { agents: [], limits: { maxDepth: 3 } }, +}) + +/** Is `entry` the end of the operation `operationId`? */ +const isEndOf = ( + entry: { type?: string; operationId?: string; event?: any }, + operationId: string, +) => + entry.operationId === operationId && + entry.event?.type === EventType.CUSTOM && + entry.event.name === 'harness.operation.finished' + +/** + * The run tree that the events give: the parent of each operation, from the + * `parentRunId` of its AG-UI RUN_STARTED. + */ +function treeOf(events: ReadonlyArray) { + const tree = new Map() + for (const { operationId, event } of events) { + if (event.type === EventType.RUN_STARTED) { + tree.set(operationId, event.parentRunId) + } + } + return tree +} + +/** Run the lead on a new host. The tree that the session lists. */ +async function runLead() { + const persistence = memoryPersistence() + const host = createHarnessHost({ persistence }) + const session = await host.open(harness, { threadId: THREAD }) + const run = session.agents.lead.start() + await expect(run).resolves.toBe('helped helped') + const expected = new Map( + session.agentRuns().map((item) => [item.runId, item.parentRunId] as const), + ) + return { host, session, run, expected } +} + +/** The tree of the agent runs in `events`, which has the session events too. */ +function agentTreeOf( + events: ReadonlyArray, + expected: Map, +) { + const tree = treeOf(events) + for (const id of tree.keys()) if (!expected.has(id)) tree.delete(id) + return tree +} + +describe('run tags on session events', () => { + it('gives the whole tree of a parallel group of agents in the root events', async () => { + const { host, session, run, expected } = await runLead() + + const events: Array = [] + for await (const entry of session.events()) { + events.push(entry) + if (isEndOf(entry, run.id)) break + } + + // The lead, its helper and nester, and the helper of the nester. + expect(expected.size).toBe(4) + const helpers = [...expected].filter(([, parent]) => parent === run.id) + expect(helpers).toHaveLength(2) + expect([...expected.values()].filter((p) => p === undefined)).toEqual([ + undefined, + ]) + expect(agentTreeOf(events, expected)).toEqual(expected) + for (const entry of events) expect(entry.threadId).toBe(THREAD) + await host.close() + }) + + it('sends the tags on GET events, on the socket, and the client keeps them', async () => { + const { host, run, expected } = await runLead() + const handler = createHarnessHandler({ host, harness, openAccess: true }) + + const sse = await readSse( + await handler(new Request(`${BASE}/events?threadId=${THREAD}`)), + (frame) => isEndOf(frame, run.id), + ) + + const { socket, sent, receive } = fakeSocket() + handleHarnessSocket({ + host, + harness, + socket, + request: new Request(`${BASE}/socket`), + openAccess: true, + }) + receive({ type: 'harness.subscribe', threadId: THREAD }) + await waitFor(() => + expect(sent.some((frame) => isEndOf(frame, run.id))).toBe(true), + ) + + const client = createHarnessClient({ + url: BASE, + threadId: THREAD, + fetch: (input, init) => handler(new Request(input, init)), + }) + const reading = new AbortController() + const kept: Array = [] + // The run ended: read from the start, not from the snapshot cursor. + for await (const entry of client.events({ + from: '0', + signal: reading.signal, + })) { + kept.push(entry) + if (isEndOf(entry, run.id)) break + } + reading.abort() + + for (const frames of [sse, sent]) { + const events = frames.filter((frame) => frame.type === 'harness.event') + expect(agentTreeOf(events, expected)).toEqual(expected) + for (const frame of events) expect(frame.threadId).toBe(THREAD) + } + expect(agentTreeOf(kept, expected)).toEqual(expected) + for (const entry of kept) expect(entry.threadId).toBe(THREAD) + socket.close() + await host.close() + }) + + it('reads an event record of an older version with no parent', async () => { + const persistence = memoryPersistence() + await persistence.stores.log.append({ + logId: logIdOf({ logId: 'old' }), + seq: 1, + epoch: UNOWNED_EPOCH, + records: [ + { + type: 'harness.event', + operationId: 'op-old', + event: { + type: EventType.CUSTOM, + name: 'test.old', + value: null, + timestamp: 1, + }, + }, + ], + }) + const host = createHarnessHost({ persistence }) + const view = await host.observe(harness, { threadId: 'old' }) + + const [first] = await (async () => { + for await (const entry of view.events()) return [entry] + return [] + })() + + expect(first).toEqual({ + cursor: '1', + threadId: 'old', + operationId: 'op-old', + event: expect.objectContaining({ name: 'test.old' }), + }) + await host.close() + }) +}) diff --git a/packages/ai-harness/tests/session-media.test.ts b/packages/ai-harness/tests/session-media.test.ts index a39321a0a8..0b271b8256 100644 --- a/packages/ai-harness/tests/session-media.test.ts +++ b/packages/ai-harness/tests/session-media.test.ts @@ -149,7 +149,7 @@ describe('session media store', () => { }) expect(record).toEqual(storedImage()) - expect(await session.getMedia(record.id)).toEqual(record) + expect(await session.getMedia({ id: record.id })).toEqual(record) expect(decoder.decode(await session.loadMedia({ id: record.id }))).toBe( 'hello', ) @@ -201,7 +201,7 @@ describe('session media store', () => { name: 'a.png', }) - expect(await other.getMedia(record.id)).toBeNull() + expect(await other.getMedia({ id: record.id })).toBeNull() await host.close() }) @@ -213,10 +213,10 @@ describe('session media store', () => { name: 'cat.png', }) - expect(await session.mediaUrl(record.id)).toEqual({ + expect(await session.mediaUrl({ id: record.id })).toEqual({ url: 'data:image/png;base64,aGVsbG8=', }) - expect(await session.mediaUrl('nope')).toStrictEqual({}) + expect(await session.mediaUrl({ id: 'nope' })).toStrictEqual({}) await host.close() }) }) diff --git a/packages/ai-harness/tests/session-plugins.test.ts b/packages/ai-harness/tests/session-plugins.test.ts index 8caf491c0b..36de9ecc2a 100644 --- a/packages/ai-harness/tests/session-plugins.test.ts +++ b/packages/ai-harness/tests/session-plugins.test.ts @@ -17,8 +17,16 @@ import { } from '../src' import { compact, modelPicker, todos, usage } from '../src/first-party' import { fileCommands, projectInstructions } from '../src/first-party/node' -import { appendToLog, gate, mockAdapter, text, toolCall } from './helpers' -import type { PluginState, SessionEvent } from '../src' +import { loadLogState, logIdOf, sessionOf } from '../src/log' +import { + after, + appendToLog, + gate, + mockAdapter, + text, + toolCall, +} from './helpers' +import type { PluginSessionApi, PluginState, SessionEvent } from '../src' // Coverage runs in CI are slow, so polled checks wait up to 5 s. const waitFor = ( @@ -878,3 +886,66 @@ describe('session events', () => { await host.close() }) }) + +describe('the turns that a session plugin starts', () => { + it('run as the owner, or as the principal that the plugin passes, not as the sender of the running turn', async () => { + let api: PluginSessionApi | undefined + const notes = definePlugin({ + name: 'test/notes', + setup: (ctx) => { + api = ctx.session + }, + }) + const release = gate() + const { adapter, calls } = mockAdapter([ + after(release.opened, 'Bea is done.'), + () => text('One.'), + () => text('Two.'), + () => text('Three.'), + ]) + const persistence = memoryPersistence() + const host = createHarnessHost({ persistence }) + const harness = defineHarness({ + name: 'test/notes', + adapter, + plugins: () => [notes], + }) + // Olga opens the thread first, so she owns it. + await host.open(harness, { threadId: 't', principal: { id: 'olga' } }) + const bea = await host.open(harness, { + threadId: 't', + principal: { id: 'bea' }, + }) + const turn = bea.send('Start a job.') + await waitFor(() => expect(turn.status()).toBe('running')) + if (!api) throw new Error('No plugin context.') + + // A job of olga ends while the turn of bea runs. + await api.note('The job of olga ended.', { wake: true }) + await api.note('The job of carl ended.', { + wake: true, + principal: { id: 'carl' }, + }) + api.send('Go on, carl.', { principal: { id: 'carl' } }) + release.open() + await turn + await waitFor(() => expect(calls).toHaveLength(4)) + + const state = await loadLogState({ + store: persistence.stores.log, + logId: logIdOf({ logId: 't' }), + rootThread: 't', + }) + const senders = [...sessionOf(state, 't').inputs.values()].map((logged) => [ + logged.input.op, + logged.principal?.id, + ]) + expect(senders).toEqual([ + ['prompt', 'bea'], + ['followUp', 'olga'], + ['followUp', 'carl'], + ['prompt', 'carl'], + ]) + await host.close() + }) +}) diff --git a/packages/ai-harness/tests/session-surface.test.ts b/packages/ai-harness/tests/session-surface.test.ts index 1a44b3263c..c4942dcd81 100644 --- a/packages/ai-harness/tests/session-surface.test.ts +++ b/packages/ai-harness/tests/session-surface.test.ts @@ -91,6 +91,27 @@ describe('the HarnessSession type', () => { runId: 'run-1', } void answers + // @ts-expect-error the handler sends a `POST run` as its run id. + void session.sendForRun + const sent: Parameters[1] = { + // @ts-expect-error `send` takes no run id. + runId: 'run-1', + } + void sent + const config: Parameters[0] = { + key: 'model', + value: 'fast', + // @ts-expect-error only an offered input has its own control id. + inputId: 'in-1', + } + void config + const answer: Parameters[0] = { + questionId: 'q-1', + value: 'yes', + // @ts-expect-error only an offered answer has a secret ref. + ref: 'harness:answer:q-1:in-1', + } + void answer await host.close() }) }) diff --git a/packages/ai-harness/tests/shared-log.test.ts b/packages/ai-harness/tests/shared-log.test.ts index e5ccc05bfd..cc6bf103f2 100644 --- a/packages/ai-harness/tests/shared-log.test.ts +++ b/packages/ai-harness/tests/shared-log.test.ts @@ -49,8 +49,19 @@ const plain = defineHarness({ describe('the log of a session tree', () => { it('lands a child record and a root record in one append, or neither', async () => { const inner = memoryLogStore() + // The app record types of each store append. + const batches: Array> = [] // A plain object, so the test can swap `append` below. - const log = wrapLogStore(inner, {}) + const log = wrapLogStore(inner, { + append: async (batch) => { + await inner.append(batch) + batches.push( + batch.records + .map((record) => record.type) + .filter((type) => type.startsWith('app.')), + ) + }, + }) const host = hostFor(durablePersistence(log)) const root = await host.open(plain, { threadId: ROOT }) const pair: Array = [ @@ -58,19 +69,26 @@ describe('the log of a session tree', () => { { type: 'app.child_linked', child: 'task-1' }, ] - await root.append(pair) - const types = async () => - (await inner.read({ logId: STORE_LOG })).map((entry) => entry.record.type) - expect(await types()).toEqual(['app.child_created', 'app.child_linked']) + await root.append({ records: pair }) + // The append is an input, so the log has its input records too. The app + // records are in one store append. + const appTypes = async () => + (await inner.read({ logId: STORE_LOG })) + .map((entry) => entry.record.type) + .filter((type) => type.startsWith('app.')) + expect(await appTypes()).toEqual(['app.child_created', 'app.child_linked']) + expect(batches.filter((types) => types.length > 0)).toEqual([ + ['app.child_created', 'app.child_linked'], + ]) // The writer tries twice, so both tries fail. const original = log.append log.append = async () => { throw new Error('store down') } - await expect(root.append(pair)).rejects.toThrow('store down') + await expect(root.append({ records: pair })).rejects.toThrow('store down') log.append = original - expect(await types()).toEqual(['app.child_created', 'app.child_linked']) + expect(await appTypes()).toEqual(['app.child_created', 'app.child_linked']) await host.close().catch(() => {}) }) @@ -79,11 +97,17 @@ describe('the log of a session tree', () => { const host = hostFor(persistence) const root = await host.open(tracked, { threadId: ROOT }) - await root.append([{ type: 'app.state_write', name: 'plan', value: 'a' }]) - await root.append([ - { type: 'app.state_write', thread: 'task-1', name: 'plan', value: 'b' }, - ]) - await root.append([{ type: 'app.state_write', name: 'done', value: true }]) + await root.append({ + records: [{ type: 'app.state_write', name: 'plan', value: 'a' }], + }) + await root.append({ + records: [ + { type: 'app.state_write', thread: 'task-1', name: 'plan', value: 'b' }, + ], + }) + await root.append({ + records: [{ type: 'app.state_write', name: 'done', value: true }], + }) const folded = { plan: 'b', done: true } expect(await root.logState()).toEqual(folded) @@ -104,7 +128,9 @@ describe('the log of a session tree', () => { it('gives undefined from session.logState() for a harness without reduce', async () => { const host = hostFor(durablePersistence()) const root = await host.open(plain, { threadId: ROOT }) - await root.append([{ type: 'app.state_write', name: 'plan', value: 'a' }]) + await root.append({ + records: [{ type: 'app.state_write', name: 'plan', value: 'a' }], + }) expect(await root.logState()).toBeUndefined() expectTypeOf(root.logState).returns.resolves.toEqualTypeOf() @@ -126,10 +152,12 @@ describe('the log of a session tree', () => { }, }) const root = await host.open(harness, { threadId: ROOT }) - await root.append([ - { type: 'app.signal', thread: 'task-1', text: 'for the child' }, - { type: 'app.signal', text: 'for the root' }, - ]) + await root.append({ + records: [ + { type: 'app.signal', thread: 'task-1', text: 'for the child' }, + { type: 'app.signal', text: 'for the root' }, + ], + }) expect((await root.transcript()).map((m) => m.content)).toEqual([ 'for the root', @@ -141,7 +169,7 @@ describe('the log of a session tree', () => { const persistence = durablePersistence() const host = hostFor(persistence) const root = await host.open(plain, { threadId: ROOT }) - await root.append([{ type: 'app.signal' }]) + await root.append({ records: [{ type: 'app.signal' }] }) // The session has no work, so the host gave the log up. Another host // takes it and writes. @@ -154,7 +182,7 @@ describe('the log of a session tree', () => { expect(host.isOpen(plain, { threadId: ROOT })).toBe(true) // The next write acquires the log again, after the record of the other // host. - await root.append([{ type: 'app.root' }]) + await root.append({ records: [{ type: 'app.root' }] }) const records = await persistence.stores.log.read({ logId: STORE_LOG }) const types = records .map((entry) => entry.record.type) @@ -168,7 +196,7 @@ describe('the log of a session tree', () => { const root = await host.open(plain, { threadId: ROOT }) await expect( - root.append([{ type: 'app.signal', thread: 7 }]), + root.append({ records: [{ type: 'app.signal', thread: 7 }] }), ).rejects.toThrow('thread') await host.close() }) diff --git a/packages/ai-harness/tests/snapshots.test.ts b/packages/ai-harness/tests/snapshots.test.ts index ea02afeafe..99be8126bb 100644 --- a/packages/ai-harness/tests/snapshots.test.ts +++ b/packages/ai-harness/tests/snapshots.test.ts @@ -1,13 +1,25 @@ import { describe, expect, it, vi } from 'vitest' import { EventType } from '@tanstack/ai' -import { memoryPersistence } from '@tanstack/ai-persistence' -import { HARNESS_EVENTS, createHarnessHost, defineHarness } from '../src' +import { UNOWNED_EPOCH, memoryPersistence } from '@tanstack/ai-persistence' +import { + HARNESS_EVENTS, + createHarnessHost, + defineHarness, + definePlugin, +} from '../src' +import { RevertFiles } from '../src/plugins' import { createHarnessHandler } from '../src/server' import { createHarnessClient } from '../src/client' import { createSessionView } from '../src/view' import { logIdOf } from '../src/log' import { mockAdapter, text, wrapLogStore } from './helpers' -import type { HarnessHost, HarnessHostOptions, SessionEvent } from '../src' +import type { + AnyHarness, + HarnessHost, + HarnessHostOptions, + PluginSessionApi, + SessionEvent, +} from '../src' // Coverage runs in CI are slow, so polled checks wait up to 5 s. const waitFor = ( @@ -25,16 +37,16 @@ const harness = defineHarness({ }) /** One turn on `THREAD`, then the session closes. */ -async function chat(host: HarnessHost) { - const session = await host.open(harness, { threadId: THREAD }) +async function chat(host: HarnessHost, target: AnyHarness = harness) { + const session = await host.open(target, { threadId: THREAD }) await session.send('Hello.') await session.close() } /** Write one turn on a host with `options`, which then stops. */ -async function seed(options: HarnessHostOptions) { +async function seed(options: HarnessHostOptions, target?: AnyHarness) { const host = createHarnessHost(options) - await chat(host) + await chat(host, target) await host.close() } @@ -107,6 +119,7 @@ describe('log snapshots on a host', () => { expect(seen).toEqual([ { cursor: head, + threadId: THREAD, operationId: 'session', event: expect.objectContaining({ type: EventType.CUSTOM, @@ -134,4 +147,111 @@ describe('log snapshots on a host', () => { view.dispose() await host.close() }) + + it('gives a plugin its kept records after a truncate, oldest first', async () => { + const persistence = memoryPersistence() + const records = [ + { type: 'acme/steps:step', n: 1 }, + { type: 'acme/other', n: 2 }, + { type: 'acme/steps:step', n: 3 }, + ] + await persistence.stores.log.append({ + logId: LOG, + seq: 1, + epoch: UNOWNED_EPOCH, + records, + }) + let session: PluginSessionApi | undefined + const steps = definePlugin({ + name: 'acme/steps', + keepRecords: { 'acme/steps:step': 10 }, + setup: (ctx) => { + session = ctx.session + }, + }) + const withSteps = defineHarness({ + name: 'test/snapshots', + adapter: mockAdapter(() => text('Hello back.')).adapter, + plugins: () => [steps], + }) + const options: HarnessHostOptions = { + persistence, + log: { snapshot: { afterBytes: 1 } }, + lifecycle: { retention: { logRecords: 'since-snapshot' } }, + } + await seed(options, withSteps) + await expect( + persistence.stores.log.read({ logId: LOG, after: 0 }), + ).rejects.toThrow() + + const host = createHarnessHost(options) + const handle = await host.open(withSteps, { threadId: THREAD }) + + expect(await session?.keptRecords('acme/steps:step')).toEqual([ + records[0], + records[2], + ]) + await handle.close() + await host.close() + }) + + it('reverts after a truncate with the file steps from before the cut', async () => { + const persistence = memoryPersistence() + const step = { + type: 'tanstack/snapshots:step', + from: 'a', + to: 'b', + files: ['notes.md'], + toolCallIds: ['call-1'], + } + // A step record of the snapshots plugin, before the first turn. + await persistence.stores.log.append({ + logId: LOG, + seq: 1, + epoch: UNOWNED_EPOCH, + records: [step], + }) + let steps: Array> = [] + const files = definePlugin({ + name: 'test/files', + provides: [RevertFiles], + // The fold keeps the steps, so the snapshot has them. + keepRecords: { [step.type]: 10 }, + setup: ({ provide }) => + provide(RevertFiles, { + revert: async (_hidden, records) => { + steps = await records(step.type) + return undefined + }, + unrevert: async () => {}, + }), + }) + const withFiles = defineHarness({ + name: 'test/snapshots', + adapter: mockAdapter(() => text('Hello back.')).adapter, + plugins: () => [files], + }) + const options: HarnessHostOptions = { + persistence, + log: { snapshot: { afterBytes: 1 } }, + lifecycle: { retention: { logRecords: 'since-snapshot' } }, + } + await seed(options, withFiles) + const snapshot = await snapshotOf(persistence) + if (!snapshot) throw new Error('The first host wrote a snapshot.') + await expect( + persistence.stores.log.read({ logId: LOG, after: 0 }), + ).rejects.toThrow() + + const host = createHarnessHost(options) + const session = await host.open(withFiles, { threadId: THREAD }) + const [first] = await session.transcript() + + expect(await session.revert({ messageId: first?.id ?? '' })).toMatchObject({ + status: 'accepted', + }) + expect(steps).toEqual([step]) + await session.close() + await host.close() + }) }) diff --git a/packages/ai-harness/tests/sse.test.ts b/packages/ai-harness/tests/sse.test.ts new file mode 100644 index 0000000000..9db7e6a885 --- /dev/null +++ b/packages/ai-harness/tests/sse.test.ts @@ -0,0 +1,320 @@ +import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest' +import { + HEARTBEAT_MS, + SLOW_CONSUMER_CLOSE_CODE, + SLOW_CONSUMER_MS, + SOCKET_BUFFER_BYTES, + SOCKET_RECHECK_MS, + STREAM_DEFAULTS, + socketSender, + sseStream, + streamValues, + withHeartbeat, +} from '../src/sse' +import type { WebSocketLike } from '@tanstack/ai' +import type { SseFrame } from '../src/sse' + +const decoder = new TextDecoder() +const encoder = new TextEncoder() +const textOf = (value: Uint8Array | undefined) => decoder.decode(value) + +beforeEach(() => { + vi.useFakeTimers() +}) +afterEach(() => { + vi.useRealTimers() +}) + +/** Frames that never come. It ends when the stream aborts its signal. */ +async function* idle(signal: AbortSignal): AsyncGenerator { + await new Promise((resolve) => signal.addEventListener('abort', resolve)) +} + +/** Endless frames of about 100 kB. `pulled` counts the reads. */ +function endless() { + const state = { pulled: 0, stopped: false } + const data = JSON.stringify('x'.repeat(100_000)) + async function* frames(signal: AbortSignal): AsyncGenerator { + try { + while (!signal.aborted) { + state.pulled += 1 + yield { data } + } + } finally { + state.stopped = true + } + } + return { state, frames } +} + +describe('streamValues', () => { + it('keeps the default of each value that is not set', () => { + expect(streamValues('test')).toEqual(STREAM_DEFAULTS) + expect(streamValues('test', { heartbeatMs: 5 })).toEqual({ + ...STREAM_DEFAULTS, + heartbeatMs: 5, + }) + }) + + it.each([ + ['heartbeatMs', 0], + ['slowConsumerMs', -1], + ['sseQueueBytes', Number.NaN], + ['socketBufferBytes', Number.POSITIVE_INFINITY], + ])('throws for a %s of %s, and names the field', (key, value) => { + expect(() => streamValues('caller', { [key]: value })).toThrow( + `caller needs a stream.${key} that is a finite number above 0.`, + ) + }) +}) + +describe('sseStream', () => { + it('sends a ping comment every 20 s, and clears its timers on cancel', async () => { + const reader = sseStream({ frames: idle, requestId: 'r1' }).getReader() + await vi.advanceTimersByTimeAsync(HEARTBEAT_MS) + expect(textOf((await reader.read()).value)).toBe(': ping\n\n') + await reader.cancel() + expect(vi.getTimerCount()).toBe(0) + }) + + it('reads the log only while the client reads: the queue holds 1 MB', async () => { + const { state, frames } = endless() + const reader = sseStream({ frames, requestId: 'r1' }).getReader() + await vi.advanceTimersByTimeAsync(0) + // About 1 MB of 100 kB frames, then it waits for the client. + expect(state.pulled).toBeGreaterThan(5) + expect(state.pulled).toBeLessThanOrEqual(12) + const before = state.pulled + await reader.read() + await vi.advanceTimersByTimeAsync(0) + expect(state.pulled).toBe(before + 1) + await reader.cancel() + }) + + it('closes a client that reads nothing for 30 s with a slow_consumer frame', async () => { + const { state, frames } = endless() + const stream = sseStream({ frames, requestId: 'r1' }) + await vi.advanceTimersByTimeAsync(SLOW_CONSUMER_MS - 1) + expect(state.stopped).toBe(false) + await vi.advanceTimersByTimeAsync(1) + expect(state.stopped).toBe(true) + const text = await new Response(stream).text() + expect(text.endsWith('\n\n')).toBe(true) + const last = text.trimEnd().split('\n\n').at(-1) + expect(JSON.parse(last!.replace(/^data: /, ''))).toEqual({ + type: 'harness.error', + message: 'The client reads too slowly.', + code: 'slow_consumer', + requestId: 'r1', + }) + expect(vi.getTimerCount()).toBe(0) + }) + + it('uses stream.heartbeatMs for the ping', async () => { + const reader = sseStream({ + frames: idle, + requestId: 'r1', + stream: streamValues('test', { heartbeatMs: 1000 }), + }).getReader() + await vi.advanceTimersByTimeAsync(1000) + expect(textOf((await reader.read()).value)).toBe(': ping\n\n') + await reader.cancel() + }) + + it('uses stream.sseQueueBytes for the queue, and stream.slowConsumerMs for the close', async () => { + const { state, frames } = endless() + const stream = sseStream({ + frames, + requestId: 'r1', + stream: streamValues('test', { + sseQueueBytes: 100_000, + slowConsumerMs: 1000, + }), + }) + await vi.advanceTimersByTimeAsync(0) + // One 100 kB frame fills the queue. + expect(state.pulled).toBeLessThanOrEqual(2) + await vi.advanceTimersByTimeAsync(999) + expect(state.stopped).toBe(false) + await vi.advanceTimersByTimeAsync(1) + expect(state.stopped).toBe(true) + const text = await new Response(stream).text() + expect(text).toContain('"code":"slow_consumer"') + }) + + it('sends an internal error frame with the requestId when the frames fail', async () => { + async function* broken(): AsyncGenerator { + yield { id: '1', data: '{"a":1}' } + throw new Error('store down') + } + const text = await new Response( + sseStream({ frames: broken, requestId: 'req-9' }), + ).text() + expect(text).toBe( + 'id: 1\ndata: {"a":1}\n\n' + + 'data: {"type":"harness.error","message":"store down","code":"internal","requestId":"req-9"}\n\n', + ) + expect(vi.getTimerCount()).toBe(0) + }) + + it('ends when the signal aborts, and stops the frames', async () => { + const watch = new AbortController() + let readerSignal: AbortSignal | undefined + const stream = sseStream({ + frames: (signal) => { + readerSignal = signal + return idle(signal) + }, + requestId: 'r1', + signal: watch.signal, + }) + watch.abort() + expect(await new Response(stream).text()).toBe('') + expect(readerSignal?.aborted).toBe(true) + }) +}) + +describe('withHeartbeat', () => { + it('adds a ping between two events, never inside one', async () => { + let source!: ReadableStreamDefaultController + const reader = withHeartbeat( + new ReadableStream({ start: (c) => void (source = c) }), + ).getReader() + source.enqueue(encoder.encode('data: 1')) + expect(textOf((await reader.read()).value)).toBe('data: 1') + await vi.advanceTimersByTimeAsync(HEARTBEAT_MS) + source.enqueue(encoder.encode('\n\n')) + expect(textOf((await reader.read()).value)).toBe('\n\n') + await vi.advanceTimersByTimeAsync(HEARTBEAT_MS) + expect(textOf((await reader.read()).value)).toBe(': ping\n\n') + source.close() + expect((await reader.read()).done).toBe(true) + expect(vi.getTimerCount()).toBe(0) + }) + + it('uses heartbeatMs for the ping', async () => { + const reader = withHeartbeat(new ReadableStream(), { + heartbeatMs: 1000, + }).getReader() + await vi.advanceTimersByTimeAsync(1000) + expect(textOf((await reader.read()).value)).toBe(': ping\n\n') + await reader.cancel() + }) +}) + +function fakeSocket() { + const handlers: Record void>> = {} + const socket = { + sent: [] as Array, + closedWith: undefined as number | undefined, + bufferedAmount: 0, + send(data: string) { + socket.sent.push(data) + }, + close(code?: number) { + socket.closedWith = code + handlers.close?.forEach((handler) => handler()) + }, + addEventListener(type: string, handler: (event?: any) => void) { + ;(handlers[type] ??= []).push(handler) + }, + } + const typed: WebSocketLike = socket + return { socket, typed } +} + +describe('socketSender', () => { + it('waits while the socket buffer is full, then sends in order', async () => { + const { socket, typed } = fakeSocket() + const send = socketSender({ socket: typed, requestId: 's1' }) + socket.bufferedAmount = SOCKET_BUFFER_BYTES + 1 + const first = send('{"n":1}') + const second = send({ type: 'harness.ping' }) + await vi.advanceTimersByTimeAsync(SOCKET_RECHECK_MS * 4) + expect(socket.sent).toEqual([]) + socket.bufferedAmount = 0 + await vi.advanceTimersByTimeAsync(SOCKET_RECHECK_MS) + await expect(first).resolves.toBe(true) + await expect(second).resolves.toBe(true) + expect(socket.sent).toEqual(['{"n":1}', '{"type":"harness.ping"}']) + socket.close() + }) + + it('closes a socket that stays full for 30 s, with a slow_consumer frame', async () => { + const { socket, typed } = fakeSocket() + const send = socketSender({ socket: typed, requestId: 's1' }) + socket.bufferedAmount = SOCKET_BUFFER_BYTES + 1 + const sent = send('{"n":1}') + await vi.advanceTimersByTimeAsync(SLOW_CONSUMER_MS) + await expect(sent).resolves.toBe(false) + expect(socket.sent.map((data) => JSON.parse(data))).toEqual([ + { + type: 'harness.error', + message: 'The client reads too slowly.', + code: 'slow_consumer', + requestId: 's1', + }, + ]) + expect(socket.closedWith).toBe(SLOW_CONSUMER_CLOSE_CODE) + expect(vi.getTimerCount()).toBe(0) + }) + + it('sends a ping every 20 s, and stops its timers when the socket closes', async () => { + const { socket, typed } = fakeSocket() + const send = socketSender({ socket: typed, requestId: 's1' }) + await vi.advanceTimersByTimeAsync(HEARTBEAT_MS) + expect(socket.sent).toEqual(['{"type":"harness.ping"}']) + socket.close() + expect(vi.getTimerCount()).toBe(0) + await expect(send('{"n":1}')).resolves.toBe(false) + }) + + it('uses stream.socketBufferBytes for a full buffer, and stream.slowConsumerMs for the close', async () => { + const { socket, typed } = fakeSocket() + const send = socketSender({ + socket: typed, + requestId: 's1', + stream: streamValues('test', { + socketBufferBytes: 1000, + slowConsumerMs: 1000, + }), + }) + // Under the default of 8 MiB, but above the 1000 bytes of the option. + socket.bufferedAmount = 1001 + const sent = send('{"n":1}') + await vi.advanceTimersByTimeAsync(SOCKET_RECHECK_MS) + expect(socket.sent).toEqual([]) + await vi.advanceTimersByTimeAsync(1000) + await expect(sent).resolves.toBe(false) + expect(socket.closedWith).toBe(SLOW_CONSUMER_CLOSE_CODE) + }) + + it('uses stream.heartbeatMs for the ping', async () => { + const { socket, typed } = fakeSocket() + socketSender({ + socket: typed, + requestId: 's1', + stream: streamValues('test', { heartbeatMs: 1000 }), + }) + await vi.advanceTimersByTimeAsync(1000) + expect(socket.sent).toEqual(['{"type":"harness.ping"}']) + socket.close() + }) + + it('ends a wait when the signal aborts', async () => { + const { socket, typed } = fakeSocket() + const stop = new AbortController() + const send = socketSender({ + socket: typed, + requestId: 's1', + signal: stop.signal, + }) + socket.bufferedAmount = SOCKET_BUFFER_BYTES + 1 + const sent = send('{"n":1}') + await vi.advanceTimersByTimeAsync(SOCKET_RECHECK_MS) + stop.abort() + await expect(sent).resolves.toBe(false) + expect(socket.sent).toEqual([]) + expect(vi.getTimerCount()).toBe(0) + }) +}) diff --git a/packages/ai-harness/tests/status-feed.test.ts b/packages/ai-harness/tests/status-feed.test.ts index aeb5e7f951..8cedbf843d 100644 --- a/packages/ai-harness/tests/status-feed.test.ts +++ b/packages/ai-harness/tests/status-feed.test.ts @@ -1,16 +1,33 @@ import { describe, expect, it, vi } from 'vitest' import { z } from 'zod' import { toolDefinition } from '@tanstack/ai' -import { memoryPersistence } from '@tanstack/ai-persistence' +import { + memoryLogBackend, + memoryLogStore, + memoryPersistence, +} from '@tanstack/ai-persistence' import { createHarnessHost, defineHarness, definePlugin, durableTool, } from '../src' -import { mockAdapter, text, toolCall } from './helpers' +import { + CRASHING_OWNERSHIP, + mockAdapter, + ownershipExpired, + text, + toolCall, + untilAborted, +} from './helpers' import type { HostEvent } from '../src' +// Coverage runs in CI are slow, so polled checks wait up to 5 s. +const waitFor = ( + check: () => T | Promise, + options?: { timeout?: number; interval?: number }, +) => vi.waitFor(check, { timeout: 5_000, ...options }) + /** Read a stream of host events into `seen` until `stop()`. */ function watch(read: (signal: AbortSignal) => AsyncIterable) { const seen: Array = [] @@ -75,15 +92,13 @@ describe('host.events', () => { const b = await host.open(plainHarness(), { threadId: 'b' }) const turn = a.send('go') - await vi.waitFor(() => - expect(a.snapshot().pendingQuestions).toHaveLength(1), - ) + await waitFor(() => expect(a.snapshot().pendingQuestions).toHaveLength(1)) await b.send('hi') const [question] = a.snapshot().pendingQuestions await a.answer({ questionId: question?.questionId ?? '', value: 'yes' }) await turn - await vi.waitFor(() => + await waitFor(() => expect(statusesOf(feed.seen, 'a')).toEqual([ 'idle', 'running', @@ -102,13 +117,11 @@ describe('host.events', () => { const a = await host.open(askingHarness(), { threadId: 'a' }) await host.open(plainHarness(), { threadId: 'b' }) void a.send('go') - await vi.waitFor(() => - expect(a.snapshot().pendingQuestions).toHaveLength(1), - ) + await waitFor(() => expect(a.snapshot().pendingQuestions).toHaveLength(1)) const feed = watch((signal) => host.events({ signal })) - await vi.waitFor(() => expect(feed.seen).toHaveLength(2)) + await waitFor(() => expect(feed.seen).toHaveLength(2)) expect(feed.seen).toEqual([ { type: 'status', @@ -130,7 +143,7 @@ describe('host.events', () => { await host.sessions.rename({ threadId: 'b', title: 'Greetings' }) await host.sessions.delete({ threadId: 'b' }) - await vi.waitFor(() => expect(feed.seen).toHaveLength(3)) + await waitFor(() => expect(feed.seen).toHaveLength(3)) expect(feed.seen).toMatchObject([ { type: 'status', threadId: 'b', status: 'idle' }, { type: 'session', threadId: 'b', entry: { title: 'Greetings' } }, @@ -152,8 +165,74 @@ describe('host.events', () => { expect(await events.next()).toEqual({ done: true, value: undefined }) // `stop` resolves only when the waiting read ends. const feed = watch((signal) => host.events({ signal })) - await vi.waitFor(() => expect(feed.seen).toHaveLength(1)) + await waitFor(() => expect(feed.seen).toHaveLength(1)) await feed.stop() await host.close() }) }) + +describe('host.events on a host that does not own the log', () => { + it('shows running, then idle, for a turn that another host runs', async () => { + const backend = memoryLogBackend() + const hostOn = () => + createHarnessHost({ + persistence: { stores: { log: memoryLogStore({ backend }) } }, + }) + const owner = hostOn() + const other = hostOn() + const harness = plainHarness() + const running = await owner.open(harness, { threadId: 'r' }) + const reading = await other.open(harness, { threadId: 'r' }) + const feed = watch((signal) => other.events({ signal })) + + await running.send('hi') + + // Coverage runs in CI are slow, so the check waits up to 5 s. + await waitFor( + () => + expect(statusesOf(feed.seen, 'r')).toEqual(['idle', 'running', 'idle']), + { timeout: 5_000 }, + ) + await feed.stop() + await reading.close() + await running.close() + await other.close() + await owner.close() + }) + + it('shows idle after it takes the log of a host that stopped in a turn', async () => { + const backend = memoryLogBackend() + const harnessOf = (replies: Parameters[0]) => + defineHarness({ + name: 'test/status-takeover', + adapter: mockAdapter(replies).adapter, + }) + const stopped = createHarnessHost({ + persistence: { stores: { log: memoryLogStore({ backend }) } }, + ownership: CRASHING_OWNERSHIP, + }) + const other = createHarnessHost({ + persistence: { stores: { log: memoryLogStore({ backend }) } }, + }) + const first = await stopped.open(harnessOf([untilAborted()]), { + threadId: 's', + }) + const second = await other.open(harnessOf([() => text('two')]), { + threadId: 's', + }) + const feed = watch((signal) => other.events({ signal })) + void Promise.resolve(first.send('one', { inputId: 'in-1' })).catch(() => {}) + await waitFor(() => expect(statusesOf(feed.seen, 's')).toContain('running')) + // The first host stops: its ownership of the log expires. + await ownershipExpired() + + // The next send takes the log from the stopped host, and runs. + await second.send('two') + + await waitFor(() => expect(statusesOf(feed.seen, 's').at(-1)).toBe('idle')) + await feed.stop() + await second.close() + await stopped.close().catch(() => {}) + await other.close() + }) +}) diff --git a/packages/ai-harness/tests/subagent-children.test.ts b/packages/ai-harness/tests/subagent-children.test.ts index bae8b2e45f..cc0b855532 100644 --- a/packages/ai-harness/tests/subagent-children.test.ts +++ b/packages/ai-harness/tests/subagent-children.test.ts @@ -190,14 +190,16 @@ describe('subagent children of a harness session', () => { // A fork copies the agent result, but the child stays a child of `t`. const fork = await host.fork(harness, { threadId: 't', newThreadId: 'f' }) - const transcript = await session.agentTranscript(id) + const transcript = await session.agentTranscript({ subagentRunId: id }) expect( transcript.map((message) => [message.role, message.content]), ).toEqual([ ['user', 'Take notes'], ['assistant', 'Note 1'], ]) - expect(await session.agentTranscript('unknown')).toEqual([]) + expect(await session.agentTranscript({ subagentRunId: 'unknown' })).toEqual( + [], + ) const other = await host.open( defineHarness({ name: 'test/subagent-children', @@ -207,8 +209,8 @@ describe('subagent children of a harness session', () => { threadId: 'u', }, ) - expect(await other.agentTranscript(id)).toEqual([]) - expect(await fork.agentTranscript(id)).toEqual([]) + expect(await other.agentTranscript({ subagentRunId: id })).toEqual([]) + expect(await fork.agentTranscript({ subagentRunId: id })).toEqual([]) await host.close() }) @@ -287,10 +289,9 @@ describe('subagent children of a harness session', () => { const rootTranscript = contents(await rootView.transcript()) expect(rootTranscript).toContain('Notes, please') expect(rootTranscript).not.toContain('Note 1') - expect(contents(await rootView.agentTranscript(id))).toEqual([ - 'Take notes', - 'Note 1', - ]) + expect( + contents(await rootView.agentTranscript({ subagentRunId: id })), + ).toEqual(['Take notes', 'Note 1']) await next.close() }) }) diff --git a/packages/ai-harness/tests/thread-settings.test.ts b/packages/ai-harness/tests/thread-settings.test.ts index 26c9abcad5..5e5f5e27fb 100644 --- a/packages/ai-harness/tests/thread-settings.test.ts +++ b/packages/ai-harness/tests/thread-settings.test.ts @@ -1,15 +1,20 @@ -import { describe, expect, it } from 'vitest' +import { describe, expect, it, vi } from 'vitest' import { z } from 'zod' -import { EventType, toolDefinition } from '@tanstack/ai' -import { memoryLogStore, memoryPersistence } from '@tanstack/ai-persistence' +import { EventType, createCapability, toolDefinition } from '@tanstack/ai' +import { + memoryLogBackend, + memoryLogStore, + memoryPersistence, +} from '@tanstack/ai-persistence' import { HARNESS_EVENTS, configOption, createHarnessHost, defineHarness, definePlugin, + logIdOf, } from '../src' -import { mockAdapter, text } from './helpers' +import { canOtherHostTake, mockAdapter, text } from './helpers' import type { AnyTool, ModelMessage } from '@tanstack/ai' import type { HarnessConfig, HarnessPersistence } from '../src' import type { Reply } from './helpers' @@ -258,6 +263,40 @@ describe('thread settings', () => { await host.close() }) + it('refuses to remove a plugin whose capability a turn plugin or a harness middleware requires', async () => { + const Notes = createCapability()('acme/notes-capability') + const { adapter } = mockAdapter(() => text('ok')) + const { host, session } = await open({ + adapter, + middleware: [{ name: 'audit', requires: [Notes] }], + plugins: () => [ + definePlugin({ + name: 'acme/notes', + provides: [Notes], + setup: (ctx) => void ctx.provide(Notes, 'notes'), + }), + definePlugin({ + name: 'acme/search', + lifetime: 'turn', + requires: [Notes], + }), + ], + }) + const remove = (names: Array) => + session.configure({ settings: { plugins: { remove: names } } }) + + expect(await remove(['acme/notes'])).toMatchObject({ + status: 'rejected', + reason: expect.stringContaining('plugin acme/search requires'), + }) + expect(await remove(['acme/notes', 'acme/search'])).toMatchObject({ + status: 'rejected', + reason: expect.stringContaining('harness middleware audit requires'), + }) + expect(session.settings()).toEqual({}) + await host.close() + }) + it('keeps the settings after a restart', async () => { const persistence = durablePersistence() const { base, strong } = models([() => text('strong')]) @@ -375,3 +414,48 @@ describe('plugin config next to the settings', () => { await host.close() }) }) + +describe('thread settings on two hosts', () => { + it('keeps the change of the other host, and a turn uses it', async () => { + const backend = memoryLogBackend() + const model = mockAdapter(() => text('ok')) + const harness = defineHarness({ + name: 'test/settings', + adapter: model.adapter, + }) + const hostOn = () => + createHarnessHost({ + persistence: { stores: { log: memoryLogStore({ backend }) } }, + }) + const hostA = hostOn() + const hostB = hostOn() + const a = await hostA.open(harness, { threadId: THREAD }) + const b = await hostB.open(harness, { threadId: THREAD }) + + await a.configure({ settings: { reasoning: 'high' } }) + // A gives the log back, so B owns it for its change. + const store = memoryLogStore({ backend }) + await vi.waitFor( + async () => + expect(await canOtherHostTake(store, logIdOf({ logId: THREAD }))).toBe( + true, + ), + { timeout: 5_000 }, + ) + expect( + await b.configure({ settings: { instructions: 'Answer in French.' } }), + ).toMatchObject({ status: 'accepted' }) + + // Each host reads the change of the other one from the log. + const both = { reasoning: 'high', instructions: 'Answer in French.' } + for (const session of [a, b]) { + await vi.waitFor(() => expect(session.settings()).toEqual(both), { + timeout: 5_000, + }) + } + await b.send('Hello.') + expect(model.calls.map((call) => call.reasoning?.level)).toEqual(['high']) + await hostA.close() + await hostB.close() + }) +}) diff --git a/packages/ai-harness/tests/title-usage.test.ts b/packages/ai-harness/tests/title-usage.test.ts index 5f04113bf8..055f24d578 100644 --- a/packages/ai-harness/tests/title-usage.test.ts +++ b/packages/ai-harness/tests/title-usage.test.ts @@ -125,6 +125,37 @@ describe('title()', () => { expect(titler.calls).toHaveLength(0) await host.close() }) + + it('counts the title call under plugin:tanstack/title', async () => { + const titler = mockAdapter(() => + text('Trip to Rome').map( + (chunk): StreamChunk => + chunk.type === EventType.RUN_FINISHED + ? { + ...chunk, + usage: { + promptTokens: 9, + completionTokens: 3, + totalTokens: 12, + }, + } + : chunk, + ), + ) + const { host, session } = await open([title({ adapter: titler.adapter })]) + + await session.send('Plan a trip to Rome for me') + + await waitFor(() => + expect(session.usage().bySource).toEqual({ + 'plugin:tanstack/title': expect.objectContaining({ + calls: 1, + totalTokens: 12, + }), + }), + ) + await host.close() + }) }) describe('usage()', () => { diff --git a/packages/ai-harness/tests/tool-op.test.ts b/packages/ai-harness/tests/tool-op.test.ts index a5d6c63107..38de6a2f73 100644 --- a/packages/ai-harness/tests/tool-op.test.ts +++ b/packages/ai-harness/tests/tool-op.test.ts @@ -5,7 +5,7 @@ import { memoryPersistence } from '@tanstack/ai-persistence' import { createHarnessHost, defineHarness } from '../src' import { applyInput, capabilitiesOf } from '../src/server' import { permissions } from '../src/first-party/permissions' -import { mockAdapter, text, toolCall } from './helpers' +import { gate, messageTexts, mockAdapter, text, toolCall } from './helpers' import type { HarnessPlugin } from '../src' import type { StreamChunk } from '@tanstack/ai' @@ -249,6 +249,49 @@ describe('{ op: "tool" }', () => { await host.close() }) + it('runs a steer sent during a tool turn as its own turn', async () => { + const started = gate() + const release = gate() + const slow = toolDefinition({ + name: 'slow_sync', + description: 'Waits for the test', + inputSchema: z.object({}), + }).server(async () => { + started.open() + await release.opened + return { synced: true } + }) + const { adapter, calls } = mockAdapter([() => text('answered')]) + const harness = defineHarness({ + name: 'test/tools-steer', + adapter, + tools: [slow], + expose: { tools: ['slow_sync'] }, + }) + const host = createHarnessHost({ persistence: memoryPersistence() }) + const session = await host.open(harness, { threadId: 't-tool-steer' }) + + const operation = session.tool({ name: 'slow_sync', args: {} }) + await started.opened + const receipt = await session.send('are you there?', { mode: 'steer' }) + .receipt + release.open() + await operation + + // The tool turn did not take the steer, so no model ran in it. + const toolTurn = await session.settled({ + inputId: (await operation.receipt).inputId ?? '', + }) + expect(toolTurn.outcome).toBe('completed') + const steer = await session.settled({ inputId: receipt.inputId }) + expect(steer.outcome).toBe('completed') + expect(steer.operationId).not.toBe(operation.id) + expect(calls).toHaveLength(1) + expect(messageTexts(calls[0]).at(-1)).toBe('are you there?') + + await host.close() + }) + it('types expose.tools and the args from the harness tools', () => { expect(() => defineHarness({ diff --git a/packages/ai-harness/tests/tracing.test.ts b/packages/ai-harness/tests/tracing.test.ts new file mode 100644 index 0000000000..3291c90bb2 --- /dev/null +++ b/packages/ai-harness/tests/tracing.test.ts @@ -0,0 +1,487 @@ +import { afterEach, describe, expect, it, vi } from 'vitest' +import { EventType, defineAgent } from '@tanstack/ai' +import { + memoryLogBackend, + memoryLogStore, + memoryPersistence, +} from '@tanstack/ai-persistence' +import { createHarnessHost, defineHarness } from '../src' +import { createHarnessClient } from '../src/client' +import { noopTracer } from '../src/observability' +import { createHarnessHandler } from '../src/server' +import { trace } from '@opentelemetry/api' +import { after, fakeSpan, gate, mockAdapter, text } from './helpers' +import type { + AttributeValue, + Span, + SpanContext, + SpanOptions, + Tracer, +} from '@opentelemetry/api' +import type { StreamChunk } from '@tanstack/ai' +import type { LogStore } from '@tanstack/ai-persistence' +import type { AnyHarness, HarnessHostOptions, OperationReport } from '../src' +import type { Reply } from './helpers' + +// Coverage runs in CI are slow, so polled checks wait up to 5 s. +const waitFor = (check: () => T | Promise) => + vi.waitFor(check, { timeout: 5_000 }) + +const THREAD = 't1' +const TRACE_ID = '4bf92f3577b34da6a3ce929d0e0e4736' +const CALLER = `00-${TRACE_ID}-00f067aa0ba902b7-01` +/** The span context that `CALLER` names. */ +const CALLER_CONTEXT: SpanContext = { + traceId: TRACE_ID, + spanId: '00f067aa0ba902b7', + traceFlags: 1, + isRemote: true, +} + +/** A span that the fake tracer started. */ +interface RecordedSpan { + name: string + options: SpanOptions + attributes: Record + /** The span of the context that `startSpan` got. */ + parent: RecordedSpan | undefined + status?: { code: number; message?: string } + ended: boolean + span: Span +} + +/** A tracer that records every span, and reads the parent from the context. */ +function fakeTracer() { + const spans: Array = [] + const tracer: Tracer = { + startSpan: (name, options = {}, context) => { + const parentSpan = context && trace.getSpan(context) + const span = fakeSpan( + { + traceId: 'a'.repeat(32), + spanId: (spans.length + 1).toString(16).padStart(16, '0'), + traceFlags: 1, + }, + { + setAttribute: (key, value) => { + recorded.attributes[key] = value + }, + setStatus: (status) => { + recorded.status = status + }, + end: () => { + recorded.ended = true + }, + }, + ) + const recorded: RecordedSpan = { + name, + options, + attributes: { ...options.attributes }, + parent: spans.find((item) => item.span === parentSpan), + ended: false, + span, + } + spans.push(recorded) + return span + }, + startActiveSpan: () => { + throw new Error('The harness does not call startActiveSpan.') + }, + } + const operations = () => + spans.filter((item) => item.name === 'harness.operation') + return { tracer, spans, operations } +} + +/** `content` as one model call that reports 7 prompt and 3 completion tokens. */ +const answerWithUsage = + (content: string): Reply => + () => + text(content).map( + (chunk): StreamChunk => + chunk.type === EventType.RUN_FINISHED + ? { + ...chunk, + usage: { promptTokens: 7, completionTokens: 3, totalTokens: 10 }, + } + : chunk, + ) + +const helper = defineAgent({ + name: 'helper', + description: 'Helps', + run: async () => 'helped', +}) + +/** Starts a helper of its own, so the tree has three levels. */ +const nester = defineAgent({ + name: 'nester', + description: 'Starts a helper', + run: async (ctx) => await ctx.harness.agents.run('helper'), +}) + +/** Runs a helper and a nester in a parallel group. */ +const lead = defineAgent({ + name: 'lead', + description: 'Runs a group', + run: async (ctx) => { + const results = await ctx.harness.agents.group({}, (g) => [ + g.start('helper'), + g.start('nester'), + ]) + return results.join(' ') + }, +}) + +const agentHarness = defineHarness({ + name: 'test/tracing-agents', + adapter: mockAdapter([]).adapter, + agents: [lead, nester, helper], + subagents: { agents: [], limits: { maxDepth: 3 } }, +}) + +const chatHarness = (replies: Array) => + defineHarness({ name: 'test/tracing', adapter: mockAdapter(replies).adapter }) + +/** A host on `log` with the session of `THREAD` open. */ +async function hostOn( + log: LogStore, + harness: AnyHarness, + options: Omit = {}, +) { + const host = createHarnessHost({ + persistence: { stores: { log } }, + ...options, + }) + const session = await host.open(harness, { threadId: THREAD }) + return { host, session } +} + +afterEach(() => { + vi.restoreAllMocks() +}) + +describe('operation spans', () => { + it('writes a span for each run, and a child run on the same host is a child span', async () => { + const { tracer, operations } = fakeTracer() + const reports: Array = [] + const host = createHarnessHost({ + persistence: memoryPersistence(), + observability: { + tracer, + onOperation: (report) => void reports.push(report), + }, + }) + const session = await host.open(agentHarness, { threadId: THREAD }) + + const run = session.agents.lead.start() + await expect(run).resolves.toBe('helped helped') + + // The lead, its helper and nester, and the helper of the nester. + await waitFor(() => expect(operations()).toHaveLength(4)) + const tree = new Map( + session.agentRuns().map((item) => [item.runId, item.parentRunId]), + ) + for (const span of operations()) { + const runId = span.attributes['harness.run_id'] + expect(typeof runId).toBe('string') + expect(span.attributes).toMatchObject({ + 'harness.thread_id': THREAD, + 'harness.operation.kind': 'agent', + 'harness.operation.status': 'completed', + }) + expect(span.ended).toBe(true) + const parentRunId = tree.get(String(runId)) + expect(span.attributes['harness.parent_run_id']).toBe(parentRunId) + expect(span.parent?.attributes['harness.run_id']).toBe(parentRunId) + } + const root = operations().find((span) => span.parent === undefined) + expect(root?.attributes).toMatchObject({ + 'harness.run_id': run.id, + 'harness.agent': 'lead', + }) + + // One started and one settled report for each run. + await waitFor(() => expect(reports).toHaveLength(8)) + for (const runId of tree.keys()) { + const phases = reports + .filter((report) => report.operation.id === runId) + .map((report) => report.phase) + expect(phases).toEqual(['started', 'settled']) + } + await host.close() + }) + + it('links the traceparent of POST run, and does not take it as the parent', async () => { + const { tracer, operations } = fakeTracer() + const host = createHarnessHost({ + persistence: memoryPersistence(), + observability: { tracer }, + }) + const harness = chatHarness([() => text('Hello.')]) + const handler = createHarnessHandler({ host, harness, openAccess: true }) + + const response = await handler( + new Request('http://x/api/harness/run', { + method: 'POST', + headers: { 'content-type': 'application/json', traceparent: CALLER }, + body: JSON.stringify({ + threadId: THREAD, + runId: 'run-a', + messages: [{ id: 'm-1', role: 'user', content: 'Hi.' }], + tools: [], + context: [], + }), + }), + ) + expect(await response.text()).toContain('Hello.') + + await waitFor(() => expect(operations()[0]?.ended).toBe(true)) + const [span] = operations() + expect(span?.options.links).toEqual([{ context: CALLER_CONTEXT }]) + expect(span?.parent).toBeUndefined() + expect(span?.attributes).toMatchObject({ + 'harness.run_id': 'run-a', + 'harness.operation.kind': 'chat', + }) + await host.close() + }) + + it('links a traceparent sent to one host on the span of the host that runs the turn', async () => { + const backend = memoryLogBackend() + const ownerTracing = fakeTracer() + const otherTracing = fakeTracer() + const otherReports: Array = [] + const release = gate() + const owner = await hostOn( + memoryLogStore({ backend }), + chatHarness([after(release.opened, 'One.'), () => text('Two.')]), + { observability: { tracer: ownerTracing.tracer } }, + ) + const harness = chatHarness([() => text('Not here.')]) + const other = await hostOn(memoryLogStore({ backend }), harness, { + observability: { + tracer: otherTracing.tracer, + onOperation: (report) => void otherReports.push(report), + }, + }) + const handler = createHarnessHandler({ + host: other.host, + harness, + openAccess: true, + }) + const sentHeaders: Array = [] + const client = createHarnessClient({ + url: 'http://x/api/harness/', + threadId: THREAD, + fetch: (input, init) => { + const request = new Request(input, init) + sentHeaders.push(request.headers.get('traceparent')) + return handler(request) + }, + }) + + await owner.session.send('Go.').receipt + const receipt = await client.send('And then?', { + mode: 'followUp', + traceparent: CALLER, + }) + expect(receipt.status).toBe('queued') + expect(sentHeaders).toEqual([CALLER]) + release.open() + expect( + (await other.session.settled({ inputId: receipt.inputId })).outcome, + ).toBe('completed') + + await waitFor(() => + expect(ownerTracing.operations().filter((s) => s.ended)).toHaveLength(2), + ) + const linked = ownerTracing + .operations() + .filter((span) => span.options.links !== undefined) + expect(linked).toHaveLength(1) + expect(linked[0]?.options.links).toEqual([{ context: CALLER_CONTEXT }]) + expect(linked[0]?.parent).toBeUndefined() + // The host that only offered the input runs nothing, so it reports nothing. + expect(otherTracing.operations()).toEqual([]) + expect(otherReports).toEqual([]) + await owner.host.close() + await other.host.close() + }) + + it('writes a harness.transfer span', async () => { + const { tracer, spans } = fakeTracer() + const host = createHarnessHost({ + persistence: memoryPersistence(), + observability: { tracer }, + }) + const harness = chatHarness([() => text('Hello.')]) + const session = await host.open(harness, { threadId: THREAD }) + await session.send('Hi.') + + const { transferId, threadId } = await host.transfer({ + from: { harness, threadId: THREAD }, + to: { harness }, + envelope: { note: 'Moved.' }, + }) + + const transfer = spans.find((span) => span.name === 'harness.transfer') + expect(transfer?.attributes).toMatchObject({ + 'harness.thread_id': THREAD, + 'harness.transfer.id': transferId, + 'harness.transfer.to_thread_id': threadId, + }) + expect(transfer?.ended).toBe(true) + await host.close() + }) +}) + +describe('onOperation', () => { + it('reports a turn once when it starts and once when it settles, with its usage', async () => { + const reports: Array = [] + const host = createHarnessHost({ + persistence: memoryPersistence(), + observability: { onOperation: (report) => void reports.push(report) }, + }) + const session = await host.open(chatHarness([answerWithUsage('Hello.')]), { + threadId: THREAD, + }) + + const turn = session.send('Hi.') + await turn + + await waitFor(() => expect(reports).toHaveLength(2)) + expect(reports[0]).toEqual({ + phase: 'started', + operation: { + id: turn.id, + kind: 'chat', + threadId: THREAD, + status: 'running', + }, + }) + expect(reports[1]).toMatchObject({ + phase: 'settled', + operation: { + id: turn.id, + kind: 'chat', + threadId: THREAD, + status: 'completed', + usage: { calls: 1, promptTokens: 7, totalTokens: 10 }, + }, + }) + expect(reports[1]?.operation.durationMs).toBeGreaterThanOrEqual(0) + await host.close() + }) + + it('reports the error of a failed turn, and a listener that throws changes nothing', async () => { + const reports: Array = [] + const host = createHarnessHost({ + persistence: memoryPersistence(), + observability: { + onOperation: (report) => { + reports.push(report) + throw new Error('listener broke') + }, + }, + }) + const session = await host.open( + chatHarness([ + () => { + throw new Error('model down') + }, + ]), + { threadId: THREAD }, + ) + + const turn = session.send('Hi.') + await expect(Promise.resolve(turn)).rejects.toThrow('model down') + + await waitFor(() => expect(reports).toHaveLength(2)) + expect(reports[1]?.operation).toMatchObject({ + status: 'failed', + error: { message: 'model down' }, + }) + await host.close() + }) + + it('a tracer that throws fails no turn and stops no report', async () => { + let calls = 0 + const reports: Array = [] + const host = createHarnessHost({ + persistence: memoryPersistence(), + observability: { + tracer: { + // The first span does not start. The second one does not end. + startSpan: () => { + calls += 1 + if (calls === 1) throw new Error('startSpan broke') + return fakeSpan(CALLER_CONTEXT, { + end: () => { + throw new Error('end broke') + }, + }) + }, + startActiveSpan: () => { + throw new Error('The harness does not call startActiveSpan.') + }, + }, + onOperation: (report) => void reports.push(report), + }, + }) + const session = await host.open( + chatHarness([() => text('One.'), () => text('Two.')]), + { threadId: THREAD }, + ) + + await session.send('Hi.') + await session.send('Again.') + + await waitFor(() => expect(reports).toHaveLength(4)) + expect(reports.map((report) => report.phase)).toEqual([ + 'started', + 'settled', + 'started', + 'settled', + ]) + await host.close() + }) +}) + +describe('without observability', () => { + it('starts no span', async () => { + const startSpan = vi.spyOn(noopTracer, 'startSpan') + const host = createHarnessHost({ persistence: memoryPersistence() }) + const session = await host.open(chatHarness([() => text('Hello.')]), { + threadId: THREAD, + }) + + await session.send('Hi.') + + expect(startSpan).not.toHaveBeenCalled() + await host.close() + }) + + it('uses one shared span when only onOperation is set', async () => { + const startSpan = vi.spyOn(noopTracer, 'startSpan') + const reports: Array = [] + const host = createHarnessHost({ + persistence: memoryPersistence(), + observability: { onOperation: (report) => void reports.push(report) }, + }) + const session = await host.open( + chatHarness([() => text('One.'), () => text('Two.')]), + { threadId: THREAD }, + ) + + await session.send('Hi.') + await session.send('Again.') + + await waitFor(() => expect(reports).toHaveLength(4)) + const returned = startSpan.mock.results.map((result) => result.value) + expect(returned).toHaveLength(2) + expect(returned[0]).toBe(returned[1]) + await host.close() + }) +}) diff --git a/packages/ai-harness/tests/transfer.test.ts b/packages/ai-harness/tests/transfer.test.ts index 1e17fbb606..8710d8c66d 100644 --- a/packages/ai-harness/tests/transfer.test.ts +++ b/packages/ai-harness/tests/transfer.test.ts @@ -317,12 +317,12 @@ describe('host.transfer to a new thread', () => { const id = parts.map(mediaIdOf).find((value) => value !== undefined) if (!id) throw new Error('The transfer lost the media part.') expect(id).not.toBe(record.id) - expect(await target.getMedia(id)).toMatchObject({ + expect(await target.getMedia({ id })).toMatchObject({ threadId, name: 'cat.png', }) expect(await target.loadMedia({ id })).toEqual(bytes) - expect(await target.getMedia(record.id)).toBeNull() + expect(await target.getMedia({ id: record.id })).toBeNull() // Nothing else moved. expect(target.settings()).toEqual({}) expect(target.config().tone?.value).toBe('plain') @@ -349,7 +349,7 @@ describe('host.transfer to a new thread', () => { }) it("adds one summary message by the source model for messages: 'summary', and counts its usage in the source", async () => { - const { host, support, billing, supportCalls } = accessHost({ + const { host, log, support, billing, supportCalls } = accessHost({ support: [ () => text('A1'), () => withUsage(text('The user wants a refund.'), 7, 3), @@ -383,6 +383,17 @@ describe('host.transfer to a new thread', () => { totalTokens: 10, }), ) + // The summary call counts under the source `transfer`. + expect(source.usage().bySource).toEqual({ + transfer: source.usage().bySender.olga, + }) + expect( + (await recordsOf(log, SOURCE)).filter( + (record) => record.type === 'harness.usage', + ), + ).toEqual([ + expect.objectContaining({ source: 'transfer', label: 'summary' }), + ]) expect(target.usage().total.calls).toBe(0) await host.close() }) diff --git a/packages/ai-harness/tests/transport-edges.test.ts b/packages/ai-harness/tests/transport-edges.test.ts index 33475fd337..3af651e2f4 100644 --- a/packages/ai-harness/tests/transport-edges.test.ts +++ b/packages/ai-harness/tests/transport-edges.test.ts @@ -21,6 +21,12 @@ import { untilAborted, } from './helpers' +// Coverage runs in CI are slow, so polled checks wait up to 5 s. +const waitFor = ( + check: () => T | Promise, + options?: { timeout?: number; interval?: number }, +) => vi.waitFor(check, { timeout: 5_000, ...options }) + const pricer = defineAgent({ name: 'pricer', description: 'Prices a vendor', @@ -173,7 +179,7 @@ describe('HTTP handler edges', () => { expect((await control.json()).status).toBe('accepted') const snapshotOf = async () => (await call('snapshot?threadId=user-1-resume')).json() - await vi.waitFor(async () => + await waitFor(async () => expect((await snapshotOf()).status).toBe('requires_action'), ) const [interrupt] = (await snapshotOf()).pendingInterrupts @@ -215,7 +221,7 @@ describe('HTTP handler edges', () => { }), }) await prompt('first') - await vi.waitFor(async () => + await waitFor(async () => expect( (await (await call('snapshot?threadId=user-1-from')).json()).status, ).toBe('idle'), @@ -249,8 +255,15 @@ describe('WebSocket edges', () => { const { socket, sent, closed, receive } = fakeSocket() handleHarnessSocket({ host, harness, socket, request: upgrade }) receive({ type: 'harness.subscribe', threadId: 'user-2-x' }) - await vi.waitFor(() => expect(closed).toEqual([4403])) - expect(sent).toEqual([{ type: 'harness.error', message: 'forbidden' }]) + await waitFor(() => expect(closed).toEqual([4403])) + expect(sent).toEqual([ + { + type: 'harness.error', + message: 'forbidden', + code: 'forbidden', + requestId: expect.any(String), + }, + ]) await host.close() }) @@ -264,20 +277,22 @@ describe('WebSocket edges', () => { request: upgrade, }) first.receive({ type: 'harness.subscribe', threadId: 'user-1-ws2' }) - await vi.waitFor(() => + await waitFor(() => expect(first.sent.some((frame) => frame.type === 'harness.hello')).toBe( true, ), ) first.receive({ type: 'harness.subscribe', threadId: 'user-1-ws2' }) - await vi.waitFor(() => + await waitFor(() => expect(first.sent.at(-1)).toEqual({ type: 'harness.error', message: 'Already subscribed.', + code: 'invalid_frame', + requestId: expect.any(String), }), ) first.receive('[1, 2]') - await vi.waitFor(() => + await waitFor(() => expect(first.sent.at(-1)).toMatchObject({ type: 'harness.error' }), ) first.receive({ @@ -285,7 +300,7 @@ describe('WebSocket edges', () => { requestId: 'p1', input: { op: 'prompt', message: 'hi' }, }) - await vi.waitFor(() => + await waitFor(() => expect( first.sent.some( (frame) => @@ -313,7 +328,7 @@ describe('WebSocket edges', () => { threadId: 'user-1-ws2', from: cursor && 'cursor' in cursor ? cursor.cursor : undefined, }) - await vi.waitFor(() => + await waitFor(() => expect(second.sent.some((frame) => frame.type === 'harness.hello')).toBe( true, ), @@ -341,7 +356,7 @@ describe('WebSocket edges', () => { requestId: 'p1', input: { op: 'prompt', message: 'hi' }, }) - await vi.waitFor(() => + await waitFor(() => expect( sent.find((frame) => frame.type === 'harness.receipt'), ).toMatchObject({ requestId: 'p1', status: 'accepted' }), @@ -351,6 +366,41 @@ describe('WebSocket edges', () => { await host.close() }) + it('closes the socket on a frame over media.maxBodyBytes, and runs nothing', async () => { + const { host } = setup() + const { adapter, calls } = mockAdapter([() => text('hello')]) + const harness = defineHarness({ + name: 'test/transport', + adapter, + media: { maxBodyBytes: 200 }, + }) + const { socket, sent, closed, receive } = fakeSocket() + handleHarnessSocket({ host, harness, socket, request: upgrade }) + + receive({ type: 'harness.subscribe', threadId: 'user-1-big' }) + await vi.waitFor(() => + expect(sent.some((frame) => frame.type === 'harness.hello')).toBe(true), + ) + receive({ + type: 'harness.input', + requestId: 'p1', + input: { op: 'prompt', message: 'x'.repeat(300) }, + }) + + // The frame is not parsed, so its error frame names the socket, not p1. + expect(sent.at(-1)).toEqual({ + type: 'harness.error', + message: 'The frame is bigger than the limit of 200 bytes.', + code: 'invalid_frame', + requestId: expect.any(String), + }) + expect(sent.at(-1)).not.toMatchObject({ requestId: 'p1' }) + expect(closed).toEqual([1009]) + expect(sent.some((frame) => frame.type === 'harness.receipt')).toBe(false) + expect(calls).toHaveLength(0) + await host.close() + }) + it('stops reading when the socket cannot send', async () => { const { host, harness } = setup() const { socket, receive } = fakeSocket({ failSend: true }) @@ -452,7 +502,7 @@ describe('protocol parsing and inputs', () => { }), ).toHaveProperty('status') release.open() - await vi.waitFor(() => expect(session.snapshot().status).toBe('idle')) + await waitFor(() => expect(session.snapshot().status).toBe('idle')) expect( await applyInput({ @@ -628,7 +678,7 @@ describe('an operation that waits for cancel', () => { principal: user, }) const operation = session.send('wait') - await vi.waitFor(() => expect(session.snapshot().status).toBe('running')) + await waitFor(() => expect(session.snapshot().status).toBe('running')) await applyInput({ harness, session, diff --git a/packages/ai-harness/tests/transport-wire.test.ts b/packages/ai-harness/tests/transport-wire.test.ts new file mode 100644 index 0000000000..92636e34ee --- /dev/null +++ b/packages/ai-harness/tests/transport-wire.test.ts @@ -0,0 +1,382 @@ +import { PassThrough } from 'node:stream' +import { afterEach, describe, expect, it, vi } from 'vitest' +import { EventType, encodeWsFrame } from '@tanstack/ai' +import { memoryPersistence } from '@tanstack/ai-persistence' +import { createHarnessHost, defineHarness } from '../src' +import { createHarnessClient } from '../src/client' +import { createHarnessHandler, handleHarnessSocket } from '../src/server' +import { HEARTBEAT_MS } from '../src/sse' +import { runHarnessWorker } from '../src/worker' +import { after, fakeSocket, gate, mockAdapter, readSse, text } from './helpers' +import type { StreamChunk } from '@tanstack/ai' +import type { SessionEvent } from '../src/types' + +// Coverage runs in CI are slow, so polled checks wait up to 5 s. +const waitFor = ( + check: () => T | Promise, + options?: { timeout?: number; interval?: number }, +) => vi.waitFor(check, { timeout: 5_000, ...options }) + +const BASE = 'http://x/api/harness' + +/** + * One model call that answers `content`, with usage. The log keeps the + * TanStack `usage` object of `RUN_FINISHED`. The wire sends the AG-UI + * `usage` array. + */ +const lossyText = (content: string) => + text(content).map((chunk) => + chunk.type === EventType.RUN_FINISHED + ? { + ...chunk, + usage: { promptTokens: 1, completionTokens: 2, totalTokens: 3 }, + } + : chunk, + ) + +/** The AG-UI wire form of a logged chunk (`toWireChunk`). */ +const wireOf = (event: StreamChunk): unknown => + JSON.parse(encodeWsFrame(event, undefined)) + +const isFinished = (event: { type: string; name?: string }) => + event.type === EventType.CUSTOM && event.name === 'harness.operation.finished' + +/** Frames until the first `harness.operation.finished` event. */ +const isTurnEnd = (frame: any) => + frame.type === 'harness.event' && isFinished(frame.event) + +/** The logged entries of `events`, until the turn ends. */ +async function logged(events: AsyncIterable) { + const entries: Array = [] + for await (const entry of events) { + entries.push(entry) + if (isFinished(entry.event)) break + } + return entries +} + +/** + * Each event frame has the wire form of the logged chunk at its cursor. At + * least one logged chunk has keys that the wire drops, so an internal chunk + * on the wire fails this check. + */ +function expectWireFrames( + frames: ReadonlyArray, + entries: ReadonlyArray, +) { + const events = frames.filter((frame) => frame.type === 'harness.event') + expect(events.map((frame) => frame.cursor)).toEqual( + entries.map((entry) => entry.cursor), + ) + events.forEach((frame, index) => { + expect(frame.event).toEqual(wireOf(entries[index]!.event)) + }) + const isLossy = entries.some( + (entry) => + JSON.stringify(wireOf(entry.event)) !== JSON.stringify(entry.event), + ) + expect(isLossy).toBe(true) +} + +function setup() { + const { adapter } = mockAdapter([() => lossyText('hello')]) + const harness = defineHarness({ name: 'test/transport-wire', adapter }) + const host = createHarnessHost({ persistence: memoryPersistence() }) + const handler = createHarnessHandler({ host, harness, openAccess: true }) + return { host, harness, handler } +} + +afterEach(() => { + vi.useRealTimers() +}) + +describe('every transport sends AG-UI wire chunks', () => { + it('GET events starts with hello, then sends the wire form of each logged chunk', async () => { + const { host, harness, handler } = setup() + const session = await host.open(harness, { threadId: 't1' }) + await session.send('hi') + + const frames = await readSse( + await handler(new Request(`${BASE}/events?threadId=t1&v=2`)), + isTurnEnd, + ) + + expect(frames[0]).toEqual({ type: 'harness.hello', v: 2, threadId: 't1' }) + expectWireFrames(frames, await logged(session.events())) + await session.close() + await host.close() + }) + + it('the socket sends the wire form of each logged chunk', async () => { + const { host, harness } = setup() + const { socket, sent, receive } = fakeSocket() + handleHarnessSocket({ + host, + harness, + socket, + request: new Request(`${BASE}/socket`), + openAccess: true, + }) + + receive({ type: 'harness.subscribe', threadId: 't2', v: 2 }) + receive({ + type: 'harness.input', + requestId: 'r1', + input: { op: 'prompt', message: 'hi' }, + }) + await waitFor(() => expect(sent.some(isTurnEnd)).toBe(true), { + timeout: 5000, + }) + + expect(sent[0]).toEqual({ type: 'harness.hello', v: 2, threadId: 't2' }) + const view = await host.observe(harness, { threadId: 't2' }) + expectWireFrames(sent, await logged(view.events())) + socket.close() + await host.close() + }) + + it('the stdio worker sends the wire form of each logged chunk', async () => { + const { adapter } = mockAdapter([() => lossyText('hello')]) + const harness = defineHarness({ name: 'test/transport-wire', adapter }) + const persistence = memoryPersistence() + const input = new PassThrough() + const frames: Array = [] + const running = runHarnessWorker(harness, { + input, + output: { write: (line: string) => frames.push(JSON.parse(line)) }, + persistence, + }) + const send = (frame: unknown) => input.write(`${JSON.stringify(frame)}\n`) + + send({ type: 'harness.subscribe', threadId: 't3', v: 2 }) + send({ + type: 'harness.input', + requestId: 'r1', + input: { op: 'prompt', message: 'hi' }, + }) + await waitFor(() => expect(frames.some(isTurnEnd)).toBe(true), { + timeout: 5000, + }) + input.end() + await running + + // A second host reads the log that the worker wrote. + const reader = createHarnessHost({ persistence }) + const view = await reader.observe(harness, { threadId: 't3' }) + expectWireFrames(frames, await logged(view.events())) + await reader.close() + }) +}) + +describe('request ids, versions, and heartbeats', () => { + it('sends back a safe X-Request-Id, and makes a new one for an unsafe value', async () => { + const { host, handler } = setup() + + const safe = await handler( + new Request(`${BASE}/capabilities`, { + headers: { 'X-Request-Id': 'req-1.a:b_c' }, + }), + ) + const unsafe = await handler( + new Request(`${BASE}/capabilities`, { + headers: { 'X-Request-Id': 'bad id\twith space' }, + }), + ) + const none = await handler(new Request(`${BASE}/sessions`)) + + expect(safe.headers.get('X-Request-Id')).toBe('req-1.a:b_c') + expect(unsafe.headers.get('X-Request-Id')).toMatch(/^[0-9a-f-]{36}$/) + expect(none.headers.get('X-Request-Id')).toMatch(/^[0-9a-f-]{36}$/) + await host.close() + }) + + it('answers an unsupported ?v= with one error frame that names the request', async () => { + const { host, harness, handler } = setup() + await (await host.open(harness, { threadId: 't4' })).close() + + const frames = await readSse( + await handler( + new Request(`${BASE}/events?threadId=t4&v=9`, { + headers: { 'X-Request-Id': 'req-9' }, + }), + ), + () => false, + ) + const hostFrames = await readSse( + await handler(new Request(`${BASE}/host-events?v=9`)), + () => false, + ) + + const refused = { + type: 'harness.error', + message: + 'Protocol version 9 is not supported. This host speaks version 2.', + code: 'unsupported_version', + } + expect(frames).toEqual([{ ...refused, requestId: 'req-9' }]) + expect(hostFrames).toEqual([{ ...refused, requestId: expect.any(String) }]) + await host.close() + }) + + it('the socket refuses an unsupported v and closes, and names the request of a bad frame', async () => { + const { host, harness } = setup() + const bad = fakeSocket() + handleHarnessSocket({ + host, + harness, + socket: bad.socket, + request: new Request(`${BASE}/socket`, { + headers: { 'X-Request-Id': 'up-1' }, + }), + openAccess: true, + }) + bad.receive({ type: 'harness.input', requestId: 'r9', input: {} }) + await waitFor( + () => + expect(bad.sent.at(-1)).toMatchObject({ + type: 'harness.error', + code: 'invalid_frame', + requestId: 'r9', + }), + { timeout: 5000 }, + ) + + bad.receive({ type: 'harness.subscribe', threadId: 't5', v: 1 }) + await waitFor(() => expect(bad.closed).toEqual([4400]), { + timeout: 5000, + }) + expect(bad.sent.at(-1)).toEqual({ + type: 'harness.error', + message: + 'Protocol version 1 is not supported. This host speaks version 2.', + code: 'unsupported_version', + // The subscribe has no requestId, so the id of the upgrade request. + requestId: 'up-1', + }) + await host.close() + }) + + it('host-events starts with hello, and the client skips it', async () => { + const { host, harness, handler } = setup() + const frames = readSse( + await handler(new Request(`${BASE}/host-events`)), + (frame) => frame.type !== 'harness.hello', + ) + const client = createHarnessClient({ + url: BASE, + threadId: 't6', + fetch: (input, init) => handler(new Request(input, init)), + }) + const reading = new AbortController() + const first = (async () => { + for await (const event of client.hostEvents({ signal: reading.signal })) + return event + return undefined + })() + + const session = await host.open(harness, { threadId: 't6' }) + await session.send('hi') + + expect((await frames)[0]).toEqual({ type: 'harness.hello', v: 2 }) + // The first frame that the client gives is a host event. + expect((await first)?.type).not.toMatch(/^harness./) + reading.abort() + await session.close() + await host.close() + }) + + it.each([ + ['the default', HEARTBEAT_MS, undefined], + ['stream.heartbeatMs', 1000, { heartbeatMs: 1000 }], + ])( + 'POST run gets a ping comment while the turn is quiet, after %s', + async (_name, wait, stream) => { + vi.useFakeTimers({ toFake: ['setInterval', 'clearInterval'] }) + const release = gate() + const { adapter } = mockAdapter([after(release.opened, 'late')]) + const harness = defineHarness({ name: 'test/transport-wire', adapter }) + const host = createHarnessHost({ persistence: memoryPersistence() }) + const handler = createHarnessHandler({ + host, + harness, + openAccess: true, + ...(stream ? { stream } : {}), + }) + + const response = await handler( + new Request(`${BASE}/run`, { + method: 'POST', + headers: { 'content-type': 'application/json' }, + body: JSON.stringify({ + threadId: 't7', + runId: 'run-7', + messages: [{ id: 'u1', role: 'user', content: 'hi' }], + tools: [], + context: [], + state: {}, + forwardedProps: {}, + }), + }), + ) + const reader = response.body!.getReader() + const decoder = new TextDecoder() + let body = '' + // Read the first event, then wait in a quiet turn. + while (!body.includes('data: ')) { + body += decoder.decode((await reader.read()).value) + } + vi.advanceTimersByTime(wait) + body += decoder.decode((await reader.read()).value) + + expect(body).toContain(': ping\n\n') + release.open() + await reader.cancel() + await host.close() + }, + ) + + it('the socket sends a ping after stream.heartbeatMs', async () => { + vi.useFakeTimers({ toFake: ['setInterval', 'clearInterval'] }) + const { host, harness } = setup() + const { socket, sent } = fakeSocket() + handleHarnessSocket({ + host, + harness, + socket, + request: new Request(`${BASE}/socket`), + openAccess: true, + stream: { heartbeatMs: 1000 }, + }) + vi.advanceTimersByTime(1000) + expect(sent).toEqual([{ type: 'harness.ping' }]) + socket.close() + await host.close() + }) + + it('the handler and the socket throw at the start for a bad stream value', async () => { + const { host, harness } = setup() + expect(() => + createHarnessHandler({ + host, + harness, + openAccess: true, + stream: { heartbeatMs: 0 }, + }), + ).toThrow( + 'createHarnessHandler needs a stream.heartbeatMs that is a finite number above 0.', + ) + expect(() => + handleHarnessSocket({ + host, + harness, + socket: fakeSocket().socket, + request: new Request(`${BASE}/socket`), + openAccess: true, + stream: { socketBufferBytes: Number.NaN }, + }), + ).toThrow( + 'handleHarnessSocket needs a stream.socketBufferBytes that is a finite number above 0.', + ) + await host.close() + }) +}) diff --git a/packages/ai-harness/tests/turn-control.test.ts b/packages/ai-harness/tests/turn-control.test.ts index 9a14a231a3..66f55eab41 100644 --- a/packages/ai-harness/tests/turn-control.test.ts +++ b/packages/ai-harness/tests/turn-control.test.ts @@ -171,9 +171,11 @@ describe('turn.onModelError', () => { turn: { onModelError: async ({ session: current, error, retries }) => { if (retries > 0 || error.code !== CONTEXT_OVERFLOW_CODE) return - await current.append([ - { type: 'app.compaction', summary: 'summary of the chat' }, - ]) + await current.append({ + records: [ + { type: 'app.compaction', summary: 'summary of the chat' }, + ], + }) return 'retry' }, }, diff --git a/packages/ai-harness/tests/types.test-d.ts b/packages/ai-harness/tests/types.test-d.ts index 3e3bcdbfc9..f4b3122b3d 100644 --- a/packages/ai-harness/tests/types.test-d.ts +++ b/packages/ai-harness/tests/types.test-d.ts @@ -2,6 +2,7 @@ import { expectTypeOf, it } from 'vitest' import { z } from 'zod' import { defineAgent } from '@tanstack/ai' import { + HARNESS_EVENTS, createHarnessHost, defineHarness, definePlugin, @@ -18,6 +19,7 @@ import type { AgentSettled, HarnessAgentContext, HarnessConfig, + HarnessEventValues, HarnessHost, HarnessInput, HarnessSession, @@ -244,6 +246,12 @@ it('types the openAccess principal by the host, and needs an explicit login mode oauthConnector({ id: 'svc', label: 'Service', oauth, login: 'device' }) }) +it('has a value type for every harness event name, and no other', () => { + expectTypeOf().toEqualTypeOf< + (typeof HARNESS_EVENTS)[keyof typeof HARNESS_EVENTS] + >() +}) + type Priced = { vendor: string; cents: number } declare const agents: AgentActions diff --git a/packages/ai-harness/tests/usage.test.ts b/packages/ai-harness/tests/usage.test.ts index 5f9462806d..3eb039d45b 100644 --- a/packages/ai-harness/tests/usage.test.ts +++ b/packages/ai-harness/tests/usage.test.ts @@ -324,6 +324,50 @@ describe('usage of each run', () => { const isRecordWithUsage = (value: unknown) => typeof value === 'object' && value !== null && 'usage' in value +describe('usage totals by source', () => { + const counts = (totalTokens: number) => ({ + promptTokens: totalTokens, + completionTokens: 0, + totalTokens, + cachedTokens: 0, + cacheWriteTokens: 0, + }) + + it('adds a call with a source under it, and keeps a turn call out of bySource', () => { + const totals = emptyUsage() + + addUsage(totals, { model: 'm', usage: counts(5) }) + addUsage(totals, { + model: 'm', + source: 'plugin:tanstack/title', + label: 'title', + usage: counts(2), + }) + addUsage(totals, { + model: 'm', + source: 'plugin:tanstack/title', + usage: counts(3), + }) + addUsage(totals, { + model: 'm', + source: 'plugin:tanstack/goal', + usage: counts(4), + }) + + expect(totals.total).toMatchObject({ calls: 4, totalTokens: 14 }) + expect(totals.bySource).toEqual({ + 'plugin:tanstack/title': expect.objectContaining({ + calls: 2, + totalTokens: 5, + }), + 'plugin:tanstack/goal': expect.objectContaining({ + calls: 1, + totalTokens: 4, + }), + }) + }) +}) + describe('usage totals of an odd key', () => { it('keeps a sender with the id __proto__ as its own entry', () => { const totals = emptyUsage() diff --git a/packages/ai-harness/tests/view-client-tools.test.ts b/packages/ai-harness/tests/view-client-tools.test.ts index ffa0e56524..e073250b90 100644 --- a/packages/ai-harness/tests/view-client-tools.test.ts +++ b/packages/ai-harness/tests/view-client-tools.test.ts @@ -86,6 +86,7 @@ describe('session view client tools', () => { ) expect(controls.at(-1)).toEqual({ op: 'resolve', + inputId: expect.any(String), resume: [ { interruptId: call?.id, @@ -125,7 +126,7 @@ describe('session view client tools', () => { approval?.approve() // A client tool id is not an approval, so a yes does not answer it. - view.approve(call?.id ?? '') + view.approve({ id: call?.id ?? '' }) expect(resolve).not.toHaveBeenCalled() call?.resolve({ opened: true }) call?.resolve({ opened: false }) diff --git a/packages/ai-harness/tests/view-fixtures.ts b/packages/ai-harness/tests/view-fixtures.ts index df0fc47cc1..5819013e36 100644 --- a/packages/ai-harness/tests/view-fixtures.ts +++ b/packages/ai-harness/tests/view-fixtures.ts @@ -8,7 +8,7 @@ let cursor = 0 /** A session event with the next cursor. */ export function at(event: StreamChunk, operationId = 'op-1'): SessionEvent { cursor += 1 - return { cursor: String(cursor), operationId, event } + return { cursor: String(cursor), threadId: 't1', operationId, event } } /** A `CUSTOM` event that the session itself sends. */ diff --git a/packages/ai-harness/tests/view-media.test.ts b/packages/ai-harness/tests/view-media.test.ts index a60825afb7..3264da07fc 100644 --- a/packages/ai-harness/tests/view-media.test.ts +++ b/packages/ai-harness/tests/view-media.test.ts @@ -154,7 +154,7 @@ async function openView(over: Partial = {}) { /** A source whose URLs expire an hour after each call, and the ids it was asked for. */ function expiringUrls() { const asked: Array = [] - const mediaUrl = async (id: string) => { + const mediaUrl = async ({ id }: { id: string }) => { asked.push(id) return { url: `https://media.test/${id}?v=${asked.length}`, diff --git a/packages/ai-harness/tests/view-paths.test.ts b/packages/ai-harness/tests/view-paths.test.ts index 66078f903f..1bd3a8828e 100644 --- a/packages/ai-harness/tests/view-paths.test.ts +++ b/packages/ai-harness/tests/view-paths.test.ts @@ -16,6 +16,12 @@ import type { } from '../src' import type { SessionView, SessionViewSource } from '../src/view' +// Coverage runs in CI are slow, so polled checks wait up to 5 s. +const waitFor = ( + check: () => T | Promise, + options?: { timeout?: number; interval?: number }, +) => vi.waitFor(check, { timeout: 5_000, ...options }) + type StreamOptions = Parameters[0] const accepted: Receipt = { inputId: 'i', status: 'accepted' } @@ -211,7 +217,7 @@ describe('session view start', () => { stream.push(says('Still here.')) - await vi.waitFor(() => expect(assistantText(view)).toBe('Still here.')) + await waitFor(() => expect(assistantText(view)).toBe('Still here.')) expect(errors).toEqual([`${read} down`]) expect(notices(view)).toEqual([{ kind: 'error', text: `${read} down` }]) expect(stream.opened[0]?.from).toBe(from) @@ -244,7 +250,7 @@ describe('session view events', () => { // A source can throw a value that is not an Error. stream.fail('socket closed') - await vi.waitFor(() => expect(view.store.get().connection).toBe('closed')) + await waitFor(() => expect(view.store.get().connection).toBe('closed')) expect(errors).toEqual(['socket closed']) expect(notices(view)).toEqual([{ kind: 'error', text: 'socket closed' }]) }) @@ -254,7 +260,7 @@ describe('session view events', () => { stream.end() - await vi.waitFor(() => expect(view.store.get().connection).toBe('closed')) + await waitFor(() => expect(view.store.get().connection).toBe('closed')) expect(errors).toEqual([]) }) @@ -282,7 +288,7 @@ describe('session view events', () => { }), ) - await vi.waitFor(() => expect(stream.closed()).toBe(true)) + await waitFor(() => expect(stream.closed()).toBe(true)) expect(calls).toEqual([]) expect(view.store.get().messages).toEqual([]) }) @@ -294,7 +300,7 @@ describe('session view events', () => { stream.push(at(childRun.started), at(childRun.finished), runFinished()) - await vi.waitFor(() => expect(view.store.get().queuedTurns).toBe(1)) + await waitFor(() => expect(view.store.get().queuedTurns).toBe(1)) await drain() expect(snapshot).toHaveBeenCalledTimes(2) }) @@ -315,10 +321,10 @@ describe('session view events', () => { custom(HARNESS_EVENTS.operationFinished, { operationId: 'op-1' }), says('Done.'), ) - await vi.waitFor(() => expect(assistantText(view)).toBe('Done.')) + await waitFor(() => expect(assistantText(view)).toBe('Done.')) hold.open() - await vi.waitFor(() => expect(view.store.get().queuedTurns).toBe(2)) + await waitFor(() => expect(view.store.get().queuedTurns).toBe(2)) await drain() expect(snapshot).toHaveBeenCalledTimes(3) }) @@ -332,7 +338,7 @@ describe('session view events', () => { stream.push(runFinished()) - await vi.waitFor(() => expect(errors).toEqual(['snapshot down'])) + await waitFor(() => expect(errors).toEqual(['snapshot down'])) expect(notices(view)).toEqual([{ kind: 'error', text: 'snapshot down' }]) }) @@ -349,7 +355,7 @@ describe('session view events', () => { custom(HARNESS_EVENTS.configChanged, { key: 'mode', value: 'plan' }), ) - await vi.waitFor(() => + await waitFor(() => expect(view.store.get().commands).toEqual([planCommand]), ) }) @@ -365,7 +371,7 @@ describe('session view events', () => { custom(HARNESS_EVENTS.configChanged, { key: 'mode', value: 'plan' }), ) - await vi.waitFor(() => expect(errors).toEqual(['describe down'])) + await waitFor(() => expect(errors).toEqual(['describe down'])) }) it('ignores reads that end after dispose', async () => { @@ -391,7 +397,7 @@ describe('session view events', () => { runFinished(), custom(HARNESS_EVENTS.configChanged, { key: 'mode', value: 'plan' }), ) - await vi.waitFor(() => { + await waitFor(() => { expect(snapshot).toHaveBeenCalledTimes(2) expect(describeSession).toHaveBeenCalledTimes(2) }) @@ -481,7 +487,7 @@ describe('session view actions', () => { await act(view) - await vi.waitFor(() => + await waitFor(() => expect(errors).toEqual(['Harness request failed (500): down']), ) expect(notices(view)).toEqual([ @@ -562,10 +568,10 @@ describe('session view approvals', () => { ) const { view, stream, inputs } = await openView({ snapshot }) - view.approve('late') + view.approve({ id: 'late' }) stream.push(runFinished()) - await vi.waitFor(() => expect(view.store.get().approvals).toHaveLength(2)) - view.approve('a') + await waitFor(() => expect(view.store.get().approvals).toHaveLength(2)) + view.approve({ id: 'a' }) // `late` still waits for its own answer. expect(inputs).toEqual([]) @@ -581,7 +587,7 @@ describe('session view approvals', () => { view.approveAll() expect(view.store.get().approvals).toEqual([]) - await vi.waitFor(() => expect(view.store.get().approvals).toHaveLength(2)) + await waitFor(() => expect(view.store.get().approvals).toHaveLength(2)) expect(errors).toEqual(['resolve down']) expect(snapshot).toHaveBeenCalledTimes(2) }) @@ -635,7 +641,7 @@ describe('session view on()', () => { }), ) - await vi.waitFor(() => expect(pings).toEqual([{ count: 3 }])) + await waitFor(() => expect(pings).toEqual([{ count: 3 }])) expect(numbered).toEqual([]) }) @@ -676,7 +682,7 @@ describe('session view on()', () => { }), ) - await vi.waitFor(() => expect(agents).toHaveLength(5)) + await waitFor(() => expect(agents).toHaveLength(5)) expect(calls).toEqual([{ id: 'c1', name: 'read' }]) expect(agents).toEqual([ { id: 's1', name: 'helper', status: 'running' }, @@ -701,7 +707,7 @@ describe('session view on()', () => { await view.ready stream.push(runFinished()) - await vi.waitFor(() => expect(view.store.get().queuedTurns).toBe(1)) + await waitFor(() => expect(view.store.get().queuedTurns).toBe(1)) expect(asked).toEqual(['q1', 'q2']) // Strict: a question without a schema has no `schema` key. @@ -732,7 +738,7 @@ describe('session view on()', () => { says('Waiting.'), ) - await vi.waitFor(() => expect(assistantText(view)).toBe('Waiting.')) + await waitFor(() => expect(assistantText(view)).toBe('Waiting.')) expect(signIns).toEqual([ { connector: 'notion', @@ -760,7 +766,7 @@ describe('session view on()', () => { custom(HARNESS_EVENTS.operationFinished, { operationId: 'op-t' }), ) - await vi.waitFor(() => expect(ended).toEqual(['op-t'])) + await waitFor(() => expect(ended).toEqual(['op-t'])) }) it('reports a lead run error and a rejected input, not a child run error', async () => { @@ -772,7 +778,7 @@ describe('session view on()', () => { at({ type: EventType.RUN_ERROR, message: 'boom' }), ) - await vi.waitFor(() => expect(errors).toHaveLength(2)) + await waitFor(() => expect(errors).toHaveLength(2)) expect(errors).toEqual(['Not accepted: unknown reason', 'Error: boom']) }) }) @@ -842,11 +848,11 @@ describe('session view inputs', () => { await view.followUp('then the tests') await view.continue() await view.cancel({ operationId: 'op-2' }) - await view.cancelInput('in-1') + await view.cancelInput({ inputId: 'in-1' }) await view.setDelivery({ inputId: 'in-2', delivery: 'steer' }) await view.configure({ settings: { model: 'strong' } }) await view.reset({ note: 'Start over.' }) - await view.revert('m1') + await view.revert({ messageId: 'm1' }) await view.unrevert() await drain() @@ -854,11 +860,11 @@ describe('session view inputs', () => { ['followUp', 'then the tests', { mode: 'followUp' }], ['continue'], ['cancel', { operationId: 'op-2' }], - ['cancelInput', 'in-1'], + ['cancelInput', { inputId: 'in-1' }], ['setDelivery', { inputId: 'in-2', delivery: 'steer' }], ['configure', { settings: { model: 'strong' } }], ['reset', { note: 'Start over.' }], - ['revert', 'm1'], + ['revert', { messageId: 'm1' }], ['unrevert'], ]) expect(view.store.get().messages).toEqual([ @@ -896,14 +902,16 @@ describe('session view inputs', () => { const { view } = await openView({ putMedia }) const bare = await openView() - expect(await view.upload({ body: new Uint8Array([1]), ...info })).toBe(file) + expect(await view.putMedia({ body: new Uint8Array([1]), ...info })).toBe( + file, + ) expect(putMedia).toHaveBeenCalledWith({ body: new Uint8Array([1]), name: 'a.png', mimeType: 'image/png', }) await expect( - bare.view.upload({ body: new Uint8Array([1]), ...info }), + bare.view.putMedia({ body: new Uint8Array([1]), ...info }), ).rejects.toThrow('This view source cannot store media.') }) }) @@ -959,7 +967,7 @@ describe('session view interrupts', () => { snapshot.mockReturnValueOnce(stopped) const { view, inputs } = await openView({ snapshot }) - view.reject('ap-1', { reason: 'wrong file' }) + view.reject({ id: 'ap-1', payload: { reason: 'wrong file' } }) view.store.get().interrupts[0]?.cancel() expect(inputs).toEqual([ @@ -988,7 +996,7 @@ describe('session view interrupts', () => { stream.push(...turn('op-a', 'interrupted'), ...turn('op-b', 'completed')) - await vi.waitFor(() => expect(ended).toEqual(['op-b'])) + await waitFor(() => expect(ended).toEqual(['op-b'])) }) }) @@ -997,7 +1005,12 @@ const entry = ( cursor: number, event: StreamChunk, operationId = 'op-1', -): SessionEvent => ({ cursor: String(cursor), operationId, event }) +): SessionEvent => ({ + cursor: String(cursor), + threadId: 't1', + operationId, + event, +}) describe('session view reload', () => { it('shows a running turn from its start, without what the transcript has', async () => { @@ -1055,7 +1068,7 @@ describe('session view reload', () => { }), ) - await vi.waitFor(() => expect(assistantText(view)).toBe('Looking.Done.')) + await waitFor(() => expect(assistantText(view)).toBe('Looking.Done.')) expect(stream.opened[0]?.from).toBe('0') expect(view.store.get().messages).toMatchObject([ { role: 'user', text: 'Clean up' }, @@ -1135,7 +1148,7 @@ describe('session view reload', () => { entry(12, says('m3', 'C.')), ) - await vi.waitFor(() => expect(assistantText(view)).toBe('A.B.C.')) + await waitFor(() => expect(assistantText(view)).toBe('A.B.C.')) expect(view.store.get().messages.at(-1)).toMatchObject({ id: 'turn-op-1', parts: [ @@ -1173,7 +1186,7 @@ describe('session view reload', () => { }), ) - await vi.waitFor(() => expect(assistantText(view)).toBe('Answer.')) + await waitFor(() => expect(assistantText(view)).toBe('Answer.')) expect(view.store.get().messages).toMatchObject([ { parts: [{ type: 'text', text: 'Answer.' }] }, ]) @@ -1196,7 +1209,7 @@ describe('session view reload', () => { entry(2, says('m1', 'Still working.')), ) - await vi.waitFor(() => expect(assistantText(view)).toBe('Still working.')) + await waitFor(() => expect(assistantText(view)).toBe('Still working.')) }) it('shows how a background child ended, from its saved card', async () => { @@ -1221,7 +1234,7 @@ describe('session view reload', () => { started('c1', 's1'), started('c2', 's2'), ], - agentTranscript: async (id) => [ + agentTranscript: async ({ subagentRunId: id }) => [ { role: 'assistant', content: id === 's1' ? 'All done.' : '', @@ -1236,8 +1249,10 @@ describe('session view reload', () => { }, ], }) + await view.loadAgent({ id: 's1' }) + await view.loadAgent({ id: 's2' }) - await vi.waitFor(() => + await waitFor(() => expect(view.store.get().messages[0]).toMatchObject({ parts: [ { type: 'tool-call', id: 'c1' }, @@ -1283,6 +1298,7 @@ describe('session view reload', () => { ], snapshot: () => sessionSnapshot({ cursor: '3' }), }) + await view.loadAgent({ id: 's1' }) const child = { subagentRunId: 's1' } stream.push( @@ -1295,7 +1311,7 @@ describe('session view reload', () => { entry(5, { ...child, type: EventType.SUBAGENT_FINISHED, result: 2 }), ) - await vi.waitFor(() => + await waitFor(() => expect(view.store.get().messages[0]).toMatchObject({ parts: [ { type: 'tool-call', id: 'c1' }, @@ -1311,7 +1327,7 @@ describe('session view reload', () => { ) }) - it('fills the agent part of a saved agent call from its own messages', async () => { + it('fills the agent part of a saved agent call from its own messages, only on loadAgent', async () => { const asked: Array = [] const { view } = await openView({ transcript: async () => [ @@ -1333,7 +1349,7 @@ describe('session view reload', () => { content: JSON.stringify({ subagentRunId: 's1', result: 'Note 1' }), }, ], - agentTranscript: async (id) => { + agentTranscript: async ({ subagentRunId: id }) => { asked.push(id) return [ { role: 'user', content: 'Notes' }, @@ -1341,24 +1357,31 @@ describe('session view reload', () => { ] }, }) + await drain() + // Until it loads, the child shows its card from the saved result. + expect(asked).toEqual([]) + expect(view.store.get().messages[0]).toMatchObject({ + parts: [{ type: 'tool-call' }, { type: 'agent', id: 's1', parts: [] }], + }) - await vi.waitFor(() => - expect(view.store.get().messages).toMatchObject([ - { - role: 'assistant', - parts: [ - { type: 'tool-call', id: 'c1', status: 'done' }, - { - type: 'agent', - id: 's1', - name: 'writer', - status: 'done', - parts: [{ type: 'text', text: 'Note 1' }], - }, - ], - }, - ]), - ) + await view.loadAgent({ id: 's1' }) + await view.loadAgent({ id: 's1' }) + + expect(view.store.get().messages).toMatchObject([ + { + role: 'assistant', + parts: [ + { type: 'tool-call', id: 'c1', status: 'done' }, + { + type: 'agent', + id: 's1', + name: 'writer', + status: 'done', + parts: [{ type: 'text', text: 'Note 1' }], + }, + ], + }, + ]) expect(asked).toEqual(['s1']) }) @@ -1388,7 +1411,69 @@ describe('session view reload', () => { }) stream.push(custom(HARNESS_EVENTS.usage, { model: 'fast' })) - await vi.waitFor(() => expect(view.store.get().usage?.total.calls).toBe(2)) + await waitFor(() => expect(view.store.get().usage?.total.calls).toBe(2)) + }) + + it('adds the next usage call from its event, with its source, and no new snapshot read', async () => { + const counts = (calls: number) => ({ + ...emptyUsage().total, + calls, + totalTokens: calls * 10, + }) + const snapshot = vi.fn(() => + sessionSnapshot({ usage: { ...emptyUsage(), total: counts(1) } }), + ) + const { view, stream } = await openView({ snapshot }) + const call = { ...counts(1), calls: undefined } + const usage = (calls: number) => + custom(HARNESS_EVENTS.usage, { + model: 'fast', + sender: 'ada', + source: 'plugin:tanstack/title', + usage: call, + total: counts(calls), + }) + + // The first call is in the snapshot already (a replay). The second is new. + stream.push(usage(1), usage(2)) + await drain() + + expect(snapshot).toHaveBeenCalledTimes(1) + expect(view.store.get().usage).toMatchObject({ + total: { calls: 2, totalTokens: 20 }, + byModel: { fast: { calls: 1, totalTokens: 10 } }, + bySender: { ada: { calls: 1 } }, + bySource: { 'plugin:tanstack/title': { calls: 1 } }, + }) + }) + + it('walks only the messages that changed on a commit', async () => { + const { view, stream } = await openView({ + transcript: async () => [ + { id: 'm0', role: 'assistant', content: 'Old.' }, + ], + }) + const old = view.store.get().messages[0] + let reads = 0 + if (old?.role === 'assistant') { + const { parts } = old + Object.defineProperty(old, 'parts', { + get: () => { + reads += 1 + return parts + }, + }) + } + + stream.push(says('Hello'), says(' there')) + + await vi.waitFor(() => + expect(view.store.get().messages[1]).toMatchObject({ + parts: [{ type: 'text', text: 'Hello there' }], + }), + ) + expect(view.store.get().messages[0]).toBe(old) + expect(reads).toBe(0) }) }) @@ -1406,7 +1491,7 @@ describe('session view dispose', () => { action: string act: (view: SessionView) => unknown }> = [ - { action: 'notice', act: (view) => view.notice('hi') }, + { action: 'notice', act: (view) => view.notice({ text: 'hi' }) }, { action: 'command', act: (view) => view.command({ name: 'ping' }) }, { action: 'setConfig', diff --git a/packages/ai-harness/tests/view.test.ts b/packages/ai-harness/tests/view.test.ts index 07ba0f02c1..826b20bd97 100644 --- a/packages/ai-harness/tests/view.test.ts +++ b/packages/ai-harness/tests/view.test.ts @@ -262,7 +262,7 @@ describe('createSessionView', () => { const second = await open([], persistence) const view = createSessionView(second.session) - view.notice('Welcome back.') + view.notice({ text: 'Welcome back.' }) await view.ready expect(view.store.get().messages).toEqual([ @@ -303,7 +303,7 @@ describe('createSessionView', () => { expect(view.store.get()).toEqual({ ...before, connection: 'closed' }) expect(snapshot).not.toHaveBeenCalled() await expect(view.send('hi')).rejects.toThrow('disposed') - expect(() => view.approve('x')).toThrow('disposed') + expect(() => view.approve({ id: 'x' })).toThrow('disposed') await host.close() }) @@ -489,7 +489,7 @@ describe('the agents of a session view', () => { expect(view.store.get().agents).toMatchObject([ { id: run.id, status: 'waiting' }, ]) - view.approve(wait.waitId) + view.approve({ id: wait.waitId }) await expect(run).resolves.toBe('Done.') expect(removed).toEqual(['old.log']) await waitFor(() => expect(view.store.get().waits).toEqual([])) @@ -523,21 +523,25 @@ describe('the agents of a session view', () => { const view = createSessionView(session) await view.ready + const message = view.store.get().messages[1] + const card = message?.role === 'assistant' ? message.parts[1] : undefined + // Until it loads, the child shows only its card. + expect(card).toMatchObject({ type: 'agent', status: 'done', parts: [] }) - await vi.waitFor(() => - expect(view.store.get().messages[1]).toMatchObject({ - role: 'assistant', - parts: [ - { type: 'tool-call', id: 'call-1', status: 'done' }, - { - type: 'agent', - name: 'writer', - status: 'done', - parts: [{ type: 'text', text: 'Note 1' }], - }, - ], - }), - ) + await view.loadAgent({ id: card?.type === 'agent' ? card.id : '' }) + + expect(view.store.get().messages[1]).toMatchObject({ + role: 'assistant', + parts: [ + { type: 'tool-call', id: 'call-1', status: 'done' }, + { + type: 'agent', + name: 'writer', + status: 'done', + parts: [{ type: 'text', text: 'Note 1' }], + }, + ], + }) view.dispose() await host.close() }) diff --git a/packages/ai-harness/tests/wire.test.ts b/packages/ai-harness/tests/wire.test.ts new file mode 100644 index 0000000000..6d7cc762e1 --- /dev/null +++ b/packages/ai-harness/tests/wire.test.ts @@ -0,0 +1,73 @@ +import { describe, expect, it } from 'vitest' +import { EventType } from '@tanstack/ai' +import { wireFrame } from '../src/wire' +import type { StreamChunk } from '@tanstack/ai' +import type { SessionEvent } from '../src/types' + +const entryOf = (event: Record): SessionEvent => ({ + cursor: '7', + threadId: 't1', + operationId: 'op-1', + event: event as unknown as StreamChunk, +}) + +describe('wireFrame', () => { + it('sends the event as an AG-UI wire chunk, and keeps subagentRunId', () => { + const entry = entryOf({ + type: EventType.TEXT_MESSAGE_CONTENT, + messageId: 'm1', + delta: 'hi', + timestamp: 1, + // Not AG-UI keys of this event: the wire moves them to the metadata. + threadId: 'thread-1', + runId: 'run-1', + subagentRunId: 'child-1', + }) + expect(JSON.parse(wireFrame(entry))).toEqual({ + type: 'harness.event', + cursor: '7', + threadId: 't1', + operationId: 'op-1', + event: { + type: EventType.TEXT_MESSAGE_CONTENT, + messageId: 'm1', + delta: 'hi', + timestamp: 1, + subagentRunId: 'child-1', + metadata: { tanstack: { threadId: 'thread-1', runId: 'run-1' } }, + }, + }) + }) + + it('keeps parentRunId on RUN_STARTED, where AG-UI has it', () => { + const entry = entryOf({ + type: EventType.RUN_STARTED, + threadId: 'thread-1', + runId: 'run-2', + parentRunId: 'run-1', + timestamp: 1, + }) + expect(JSON.parse(wireFrame(entry)).event).toEqual({ + type: EventType.RUN_STARTED, + threadId: 'thread-1', + runId: 'run-2', + parentRunId: 'run-1', + timestamp: 1, + }) + }) + + it('encodes one entry object once', () => { + const entry = entryOf({ + type: EventType.TEXT_MESSAGE_CONTENT, + messageId: 'm1', + delta: 'first', + timestamp: 1, + }) + const first = wireFrame(entry) + // A second call reads the cache: a change to the object does not show. + ;(entry.event as unknown as { delta: string }).delta = 'changed' + expect(wireFrame(entry)).toBe(first) + // Another object with the same content is encoded again. + expect(wireFrame({ ...entry })).toContain('changed') + }) +}) diff --git a/packages/ai-harness/tests/workspace-tools.test.ts b/packages/ai-harness/tests/workspace-tools.test.ts index 8c361781bf..74ec2f69fd 100644 --- a/packages/ai-harness/tests/workspace-tools.test.ts +++ b/packages/ai-harness/tests/workspace-tools.test.ts @@ -3,7 +3,12 @@ import { tmpdir } from 'node:os' import * as nodePath from 'node:path' import { afterAll, beforeAll, describe, expect, it, vi } from 'vitest' import { memoryPersistence } from '@tanstack/ai-persistence' -import { createHarnessHost, defineHarness, definePlugin } from '../src' +import { + createHarnessHost, + defineHarness, + definePlugin, + harnessAgent, +} from '../src' import { permissions } from '../src/first-party/permissions' import { hostBackend } from '../src/first-party/coding/backend' import { @@ -1236,7 +1241,7 @@ describe('a workspace for each thread', () => { ) } - it('closes the workspace that bind gave when the session closes', async () => { + it('closes the workspace that bind gave as moved when the session closes, so the thread can reattach it', async () => { const { workspace } = threadWorkspaces() const closed: Array = [] const closing: WorkspaceBinder = { @@ -1258,7 +1263,153 @@ describe('a workspace for each thread', () => { await host.close() - expect(closed).toEqual(['closed']) + // A session close does not end the thread: a sandbox stays. + expect(closed).toEqual(['moved']) + }) + + it('destroys the sandbox of a thread that host.sessions.delete purges, also when no session has the thread open', async () => { + const { workspace } = threadWorkspaces() + const closed: Array = [] + const reattached: Array = [] + const sandbox = async (input: Parameters[0]) => ({ + ...(await workspace.bind(input)), + binding: () => ({ + kind: 'sandbox' as const, + provider: 'test', + sandboxId: `sb-${input.run.threadId}`, + }), + close: async ({ reason }: { reason: 'closed' | 'moved' }) => { + closed.push(reason) + }, + }) + const owned: WorkspaceBinder = { + bind: sandbox, + reattach: async ({ run, binding }) => { + reattached.push(binding.kind === 'sandbox' ? binding.sandboxId : '') + return sandbox({ run }) + }, + } + const { adapter } = mockAdapter([ + () => toolCall('read_file', { path: 'mine.txt' }, 'c0'), + () => text('done'), + ]) + const host = createHarnessHost({ + persistence: memoryPersistence(), + lifecycle: { idleCloseMs: 0 }, + }) + const session = await openThread(host, owned, adapter, 'a') + await session.send('read') + await session.close() + await waitFor(() => expect(closed).toEqual(['moved'])) + + await host.sessions.delete({ threadId: 'a' }) + + // The purge reattached the sandbox from the log, then destroyed it. + expect(reattached).toEqual(['sb-a']) + expect(closed).toEqual(['moved', 'closed']) + await host.close() + }) + + it('destroys the sandbox on a host that never opened the harness, when its harnesses option has it, and warns when not', async () => { + const { workspace } = threadWorkspaces() + const closed: Array = [] + const sandbox = async (input: Parameters[0]) => ({ + ...(await workspace.bind(input)), + binding: () => ({ + kind: 'sandbox' as const, + provider: 'test', + sandboxId: `sb-${input.run.threadId}`, + }), + close: async ({ reason }: { reason: 'closed' | 'moved' }) => { + closed.push(`${input.run.threadId}:${reason}`) + }, + }) + const owned: WorkspaceBinder = { + bind: sandbox, + reattach: async ({ run }) => sandbox({ run }), + } + const { adapter } = mockAdapter([ + () => toolCall('read_file', { path: 'mine.txt' }, 'c0'), + () => text('done'), + () => toolCall('read_file', { path: 'mine.txt' }, 'c1'), + () => text('done'), + ]) + const harness = defineHarness({ + name: 'test/workspace-delete', + adapter, + plugins: () => [workspaceTools({ workspace: owned, web: false })], + }) + const persistence = memoryPersistence() + const first = createHarnessHost({ persistence }) + for (const threadId of ['a', 'b']) { + const session = await first.open(harness, { threadId }) + await session.send('read') + } + await first.close() + closed.length = 0 + const warn = vi.spyOn(console, 'warn').mockImplementation(() => {}) + + // A new host, after a deploy: it never opened the harness. + const knows = createHarnessHost({ persistence, harnesses: [harness] }) + await knows.sessions.delete({ threadId: 'a' }) + const blind = createHarnessHost({ persistence }) + await blind.sessions.delete({ threadId: 'b' }) + + expect(closed).toEqual(['a:closed']) + expect(warn).toHaveBeenCalledWith( + expect.stringContaining('does not know the harness'), + ) + warn.mockRestore() + await knows.close() + await blind.close() + }) + + it('destroys the sandbox of a harness child thread when host.sessions.delete purges its root', async () => { + const { workspace } = threadWorkspaces() + const closed: Array = [] + const sandbox = async (input: Parameters[0]) => ({ + ...(await workspace.bind(input)), + binding: () => ({ + kind: 'sandbox' as const, + provider: 'test', + sandboxId: `sb-${input.run.threadId}`, + }), + close: async ({ reason }: { reason: 'closed' | 'moved' }) => { + closed.push(reason) + }, + }) + const child = defineHarness({ + name: 'test/workspace-child', + adapter: mockAdapter([ + () => toolCall('read_file', { path: 'mine.txt' }, 'c0'), + () => text('Child done.'), + ]).adapter, + plugins: () => [ + workspaceTools({ + workspace: { bind: sandbox, reattach: ({ run }) => sandbox({ run }) }, + web: false, + }), + ], + }) + const agent = harnessAgent(child) + const lead = defineHarness({ + name: 'test/workspace-lead', + adapter: mockAdapter([ + () => toolCall(agent.name, { prompt: 'Read it.' }), + () => text('Done.'), + ]).adapter, + subagents: { agents: [agent] }, + }) + const host = createHarnessHost({ persistence: memoryPersistence() }) + const session = await host.open(lead, { threadId: 't' }) + await session.send('Go.') + await waitFor(() => expect(closed).toEqual(['moved'])) + + await host.sessions.delete({ threadId: 't' }) + + // The child lives in the log of its root: the purge reads it there. + expect(closed).toEqual(['moved', 'closed']) + await host.close() }) it('keeps a thread in its root: no path through .. or a link reaches another root', async () => { diff --git a/packages/ai-mcp/src/client.ts b/packages/ai-mcp/src/client.ts index 1ad773326c..0cb25127b0 100644 --- a/packages/ai-mcp/src/client.ts +++ b/packages/ai-mcp/src/client.ts @@ -65,6 +65,8 @@ type ToolPolicy = Pick< 'toolFilter' | 'needsApproval' | 'toolName' | 'requestOptions' > +type ClientEvents = Pick + export interface MCPClient< TServer extends ServerDescriptor = AutomaticDescriptor, > { @@ -172,6 +174,9 @@ class MCPClientImpl< instructions: string | undefined readonly #client: Client #closed = false + // Set when connect succeeds. A close before that is a failed connect. + #connected = false + readonly #events: ClientEvents #toolDefinitions?: Map // Open only after a spec 2026 connect when the server can report tool changes. #toolListSubscription?: McpSubscription @@ -193,11 +198,13 @@ class MCPClientImpl< clientOptions?: ClientOptions, policy: ToolPolicy = {}, onElicit?: OnElicit, + events: ClientEvents = {}, ) { this.prefix = prefix this.#transport = transport this.#clientOptions = clientOptions this.#policy = policy + this.#events = events // Try spec 2026-07-28 first. If the server does not support it, the // client uses the 2025 initialize handshake. Caller options still apply. // `mode` stays `auto` even when `clientOptions` sets another mode. @@ -256,8 +263,15 @@ class MCPClientImpl< 'notifications/tools/list_changed', () => { this.#toolDefinitions = undefined + this.#events.onToolListChanged?.() }, ) + // The SDK calls `onclose` for every close, also for our own `close()`. + // Only a close that this client did not start is reported. + this.#client.onclose = () => { + const isDropped = this.#connected && !this.#closed + if (isDropped) this.#events.onClose?.() + } // The connect handshake gets the request options too, so `timeout` // also limits a server that does not answer while it connects. await this.#client.connect( @@ -269,6 +283,7 @@ class MCPClientImpl< // A failed listen only means that tool changes are not pushed. // The connection still works, so the error does not fail connect. await this.#listenForToolListChanges().catch(() => undefined) + this.#connected = true } catch (err) { await this.#toolListSubscription?.close().catch(() => undefined) await this.#client.close().catch(() => undefined) @@ -607,7 +622,8 @@ export type OnElicit = ( /** * Connects a client to a transport. The harness passes `onElicit` to answer - * the `elicitation/create` requests of spec 2025 servers. + * the `elicitation/create` requests of spec 2025 servers. `options.onClose` + * and `options.onToolListChanged` get the events of the connection. */ export async function connectTransport< TServer extends ServerDescriptor = AutomaticDescriptor, @@ -628,6 +644,10 @@ export async function connectTransport< requestOptions: options.requestOptions, }, onElicit, + { + onClose: options.onClose, + onToolListChanged: options.onToolListChanged, + }, ) await impl.connect(transport) return impl diff --git a/packages/ai-mcp/src/connector.ts b/packages/ai-mcp/src/connector.ts index 52570614dc..de920c569c 100644 --- a/packages/ai-mcp/src/connector.ts +++ b/packages/ai-mcp/src/connector.ts @@ -227,6 +227,10 @@ interface Connection { connected: boolean client?: MCPClient tools?: ReadonlyArray + /** The server sent `tools/list_changed`: the next turn lists again. */ + stale?: boolean + /** Wraps a tool list of `client` for the session. */ + wrap?: (tools: ReadonlyArray) => ReadonlyArray } /** @@ -240,6 +244,8 @@ interface Connection { * - After sign-in, the next turn has the server's tools, named * `_` by default. `toolName` can change the names. * Tools that can change data ask for approval. + * - When the connection closes, the next turn connects again. When the + * server sends `tools/list_changed`, the next turn lists its tools again. * * @example * ```ts @@ -313,8 +319,10 @@ export function mcpConnector(options: McpConnectorOptions) { : undefined) ?? 'http://127.0.0.1/callback' // The tool calls of this client that run now, newest last. const calling: Array = [] + const current = connection + let opened: MCPClient | undefined try { - connection.client = await connectTransport( + opened = await connectTransport( { transport: { type: 'http', @@ -333,9 +341,20 @@ export function mcpConnector(options: McpConnectorOptions) { options.needsApproval ?? ((tool) => tool.annotations?.readOnlyHint !== true), clientOptions: harnessClientOptions, + // A dropped connection: the next turn connects again. + onClose: () => { + if (current.client !== opened) return + current.client = undefined + current.tools = undefined + }, + onToolListChanged: () => { + if (current.client === opened) current.stale = true + }, }, elicitFor(calling), ) + connection.client = opened + connection.stale = false } catch (error) { // The server wants a new sign-in: ask the user to run /connect. // Other failures, for example a server that is down, stay @@ -348,7 +367,7 @@ export function mcpConnector(options: McpConnectorOptions) { ctx.session.authRequired({ connector: id }) throw error.cause } - const client = connection.client + const client = opened // A failed tool list closes the client: the next turn connects // again, and must not leave this one open. const listed = await client @@ -358,14 +377,22 @@ export function mcpConnector(options: McpConnectorOptions) { await client.close().catch(() => {}) throw error }) - connection.tools = askForInput( - trackCaller( + const wrap = (found: ReadonlyArray) => + askForInput(trackCaller(found, ctx.session, calling), ctx.session) + connection.wrap = wrap + connection.tools = wrap(listed as ReadonlyArray) + } else if (connection.stale && connection.client) { + // After `tools/list_changed`: list the tools again. On a + // failure, the old list stays, and the next turn tries again. + connection.stale = false + const client = connection.client + const listed = await client.tools().catch(() => undefined) + if (listed === undefined) connection.stale = true + else if (connection.client === client && connection.wrap) { + connection.tools = connection.wrap( listed as ReadonlyArray, - ctx.session, - calling, - ), - ctx.session, - ) + ) + } } return connection.tools }, diff --git a/packages/ai-mcp/src/harness-plugin.ts b/packages/ai-mcp/src/harness-plugin.ts index 5362248284..d055618230 100644 --- a/packages/ai-mcp/src/harness-plugin.ts +++ b/packages/ai-mcp/src/harness-plugin.ts @@ -11,8 +11,10 @@ import type { RunContext, TurnInfo, } from '@tanstack/ai-harness' -import type { OnElicit } from './client' +import type { MCPClient, OnElicit } from './client' +import type { McpConnectorOptions } from './connector' import type { AskInput } from './tools' +import type { MCPClientOptions } from './types' import type { HttpTransportConfig, StdioTransportConfig } from './transport' /** One server of {@link mcp}, keyed by its name. */ @@ -41,18 +43,31 @@ export type McpServerConfig = ( * changes. */ codeMode?: boolean + /** + * Ask the user for approval before a tool of this server runs. `true` asks + * for every tool. A function gets the MCP tool definition and returns + * `true` to ask. Default: no approval, and for `oauth: true` every tool + * that the server does not mark `readOnlyHint`. + */ + needsApproval?: boolean | McpConnectorOptions['needsApproval'] + /** + * The name the model sees for each tool of this server. It gets the MCP + * tool definition. Default: `_`. + */ + toolName?: MCPClientOptions['toolName'] } /** * The status of one server. UIs read the list from - * `snapshot().plugins['tanstack/mcp'].servers`. + * `snapshot().plugins['tanstack/mcp'].servers`. A server is `reconnecting` + * after its connection closed, until a new connection is open. */ export type McpServerStatus = { name: string /** How many tools the server gives the model. */ toolCount: number } & ( - | { status: 'connecting' | 'connected' } + | { status: 'connecting' | 'connected' | 'reconnecting' } | { status: 'failed'; error: string } ) @@ -70,9 +85,31 @@ export type McpServerStatus = { * sessions closes, or at `host.close()`. The other servers connect for * each session. * - An `http` server with `oauth: true` signs in through `/connect `. + * - When the connection of a server closes, the server is `reconnecting`. + * The plugin connects again after 0.5 s, 1 s, 2 s, and so on (at most + * 30 s), until it is connected or its session or host closes. The wait + * starts again at 0.5 s only after a connection stayed open for 30 s, so a + * server that crashes right after it connects waits longer each time. The + * next turn gets the tools of the new connection. A server that fails its + * first connect is `failed`, and connects again the same way. + * - A spec 2025 server that asks for input (`elicitation/create`) does not + * name the tool call that asks. The user of the call is asked only when + * all the calls of that server that run now belong to one run. Else the + * request is declined, so a question never goes to another user. + * - When a server sends `tools/list_changed`, the next turn lists its tools + * again. + * - When two servers give a tool with the same name, the model gets the + * tool of the server that comes first in `servers`. The plugin warns once. + * With `onDuplicateTool: 'error'`, the later server fails instead. * * @param options.servers - The servers, keyed by name. The name is the tool * prefix. + * @param options.onDuplicateTool - What happens when a server gives a tool + * name that a server earlier in `servers` gives too. `'warn'` (default): + * the model does not get that one tool, and the plugin warns once. + * `'error'`: the later server gets the status `failed` with an error that + * names both servers and the tool, and the model gets none of its tools. + * The other servers keep working. * * @example * ```ts @@ -91,6 +128,7 @@ export type McpServerStatus = { */ export function mcp(options: { servers: Readonly> + onDuplicateTool?: 'warn' | 'error' }) { const servers = Object.entries(options.servers) return definePlugin({ @@ -100,17 +138,20 @@ export function mcp(options: { setup: (host) => { // The tool calls of a shared server that run now, newest last. const calling: Array = [] - // ponytail: a spec 2025 request of a shared server goes to the newest - // running call. Two sessions that call one stdio server at the same - // time can get the question of the other one. + // A spec 2025 request of a shared server goes to the run that calls, + // only when one run calls. Else it is declined (see `elicitFor`). const askCaller = elicitFor(calling) const shared = new Map() for (const [name, server] of servers) { if (server.type !== 'stdio') continue shared.set(name, openServer(host.resources, name, server, askCaller)) } + // The tool name clashes that were warned about, for all sessions. + const warned = new Set() + const onDuplicate = options.onDuplicateTool ?? 'warn' return { - session: (ctx) => sessionPart(ctx, servers, shared, calling), + session: (ctx) => + sessionPart(ctx, servers, shared, calling, warned, onDuplicate), } }, }) @@ -122,6 +163,8 @@ async function sessionPart( servers: ReadonlyArray<[string, McpServerConfig]>, shared: ReadonlyMap, calling: Array, + warned: Set, + onDuplicate: 'warn' | 'error', ) { const statuses = new Map() for (const [name] of servers) { @@ -138,29 +181,43 @@ async function sessionPart( await state.update(() => ({ servers: [...statuses.values()] })) /** - * The tools of one server for this session, in the background. A - * stdio server is the shared one of the host. Never rejects. + * Connect one server in the background. Returns the tools of the server + * for each turn. A stdio server is the shared one of the host. Never + * rejects. */ - const connect = async (name: string, server: McpServerConfig) => { + const connect = (name: string, server: McpServerConfig) => { const sharedServer = shared.get(name) // The tool calls of this server for this session that run now. const own: Array = [] - const opened = await (sharedServer ?? - openServer(ctx.resources, name, server, elicitFor(own))) - const failed = opened instanceof Error - const status: McpServerStatus = failed - ? { name, status: 'failed', error: opened.message, toolCount: 0 } - : { name, status: 'connected', toolCount: opened.tools.length } + const opened = + sharedServer ?? openServer(ctx.resources, name, server, elicitFor(own)) // After the session closed, the state store can be gone. `/mcp` // reads `statuses`, which is set first. - await setStatus(status).catch(() => undefined) - if (failed) return [] - const tools = trackCaller( - opened.tools, - ctx.session, - sharedServer ? calling : own, - ) - return askForInput(tools, ctx.session) + const opening = opened.then(async (link) => { + if (link instanceof Error) { + await setStatus({ + name, + status: 'failed', + error: link.message, + toolCount: 0, + }).catch(() => undefined) + return undefined + } + const update = () => setStatus(link.status()).catch(() => undefined) + link.watch(() => void update(), ctx.resources.signal) + await update() + return link + }) + return async (_turn: TurnInfo) => { + const link = await opening + if (!link) return [] + const tools = trackCaller( + await link.tools(), + ctx.session, + sharedServer ? calling : own, + ) + return askForInput(tools, ctx.session) + } } const commands: Record = { @@ -175,16 +232,15 @@ async function sessionPart( servers.map(async (entry) => { const [name, server] = entry const isOAuth = server.type === 'http' && server.oauth === true - if (!isOAuth) { - const connecting = connect(name, server) - return (_turn: TurnInfo) => connecting - } + if (!isOAuth) return { name, list: connect(name, server) } // Loaded here: the connector needs node:crypto and node:http. const { mcpConnector } = await import('./connector') const connector = await mcpConnector({ id: name, label: name, url: server.url, + toolName: server.toolName, + needsApproval: approvalOf(server), requestOptions: requestOptionsOf(server), ...(server.fetch ? { fetch: server.fetch } : {}), }).setup(ctx) @@ -192,7 +248,7 @@ async function sessionPart( prompts.push(...connector.prompts) // The sign-in belongs to the sender of the turn, so the status is // the one of the last turn. - return async (turn: TurnInfo) => { + const list = async (turn: TurnInfo) => { try { const tools = await connector.discoverTools(turn) const signedIn = @@ -219,6 +275,7 @@ async function sessionPart( return [] } } + return { name, list } }), ) @@ -226,18 +283,31 @@ async function sessionPart( commands, prompts, discoverTools: async (turn: TurnInfo) => { - const found = await Promise.all(lists.map((list) => list(turn))) - return found.flat() + const found = await Promise.all( + lists.map(async ({ name, list }) => ({ + name, + tools: await list(turn), + })), + ) + return onDuplicate === 'error' + ? failOnDuplicate(found, setStatus) + : keepFirstTools(found, warned) }, } } -/** A server that the sessions of a host share: its tools, or its error. */ +/** A server that the sessions of a host share: its connection, or its error. */ type SharedServer = ReturnType /** - * Connect one server in `resources`. Resolves to the error of a server that - * failed, so it never rejects. + * Connect one server in `resources`. It never rejects: an error that stops + * the setup itself resolves to that error. + * + * When the first connect fails, the status is `failed` with the error, and + * the server connects again with backoff. When the connection closes later, + * the status is `reconnecting`, with the same backoff. Both go on until a + * connection is open or `resources` is disposed. After + * `tools/list_changed`, the next `tools()` lists the tools again. */ function openServer( resources: PluginResources, @@ -245,36 +315,244 @@ function openServer( server: McpServerConfig, onElicit: OnElicit, ) { + const { signal } = resources + const watchers = new Set<() => void>() + let status: McpServerStatus = { name, status: 'connecting', toolCount: 0 } + let client: MCPClient | undefined + let tools: ReadonlyArray = [] + // Set by tools/list_changed. The next `tools()` lists the tools again. + let stale = false + let listing: Promise | undefined + // The backoff goes on across drops, so a server that crashes right after + // it connects waits longer each time. It resets only after a connection + // stayed open for `stableMs`. + let attempt = 0 + let openedAt = 0 + + const setStatus = (next: McpServerStatus) => { + status = next + for (const watcher of watchers) watcher() + } + + /** One connection with its tool list. Throws when the server fails. */ + const open = async () => { + // Loaded here: the stdio transport needs node:child_process. + const transport = + server.type === 'stdio' + ? (await import('./stdio')).stdioTransport(server) + : server + const next = await connectTransport( + { + transport, + prefix: name, + toolName: server.toolName, + needsApproval: approvalOf(server), + requestOptions: requestOptionsOf(server), + clientOptions: harnessClientOptions, + onClose: () => dropped(), + onToolListChanged: () => { + stale = true + }, + }, + onElicit, + ) + let listed: ReadonlyArray + try { + listed = forCodeMode(server, await next.tools()) + } catch (error) { + // Connected, but no tool list: the server failed. + await next.close().catch(() => undefined) + throw error + } + // Disposed while this connect ran: nothing uses the connection. + if (signal.aborted) { + await next.close().catch(() => undefined) + throw signal.reason + } + client = next + tools = listed + stale = false + openedAt = Date.now() + setStatus({ name, status: 'connected', toolCount: listed.length }) + } + + /** The connection closed, and this plugin did not close it. */ + const dropped = () => { + if (status.status !== 'connected' || signal.aborted) return + client = undefined + tools = [] + if (Date.now() - openedAt >= stableMs) attempt = 0 + setStatus({ name, status: 'reconnecting', toolCount: 0 }) + void reconnect() + } + + const reconnect = async () => { + for (;;) { + await delay(backoffMs(attempt++), signal) + if (signal.aborted) return + const isOpen = await open().then( + () => true, + () => false, + ) + if (isOpen) return + } + } + + /** List the tools of `current` again. On a failure, the old list stays. */ + const relist = async (current: MCPClient) => { + stale = false + const listed = await current.tools().then( + (found) => forCodeMode(server, found), + () => undefined, + ) + // The connection closed while it listed. + if (client !== current) return + if (listed === undefined) { + stale = true + return + } + tools = listed + setStatus({ name, status: 'connected', toolCount: listed.length }) + } + + const link = { + status: () => status, + /** The tools of the open connection. Never rejects. */ + tools: async () => { + if (stale && client) { + listing ??= relist(client).finally(() => { + listing = undefined + }) + await listing + } + return tools + }, + /** Calls `watcher` after each status change, until `until` aborts. */ + watch: (watcher: () => void, until: AbortSignal) => { + if (until.aborted) return + watchers.add(watcher) + until.addEventListener('abort', () => watchers.delete(watcher), { + once: true, + }) + }, + } + return resources .acquire( async () => { - // Loaded here: the stdio transport needs node:child_process. - const transport = - server.type === 'stdio' - ? (await import('./stdio')).stdioTransport(server) - : server - const client = await connectTransport( - { - transport, - prefix: name, - requestOptions: requestOptionsOf(server), - clientOptions: harnessClientOptions, - }, - onElicit, - ) try { - return { client, tools: forCodeMode(server, await client.tools()) } + await open() } catch (error) { - // Connected, but no tool list: the server failed. - await client.close().catch(() => undefined) - throw error + if (signal.aborted) throw error + // A server that is not up yet (a slow `npx`, a short outage) gets + // its tools at a later turn. + setStatus({ + name, + status: 'failed', + error: errorText(error), + toolCount: 0, + }) + void reconnect() } + return link }, - ({ client }) => client.close(), + () => client?.close(), ) .catch((error: unknown) => new Error(errorText(error))) } +/** How long a connection stays open before the backoff starts again at 0.5 s. */ +const stableMs = 30_000 + +/** + * The wait before reconnect `attempt` (0 is the first): 0.5 s, 1 s, 2 s, + * and so on, at most 30 s. A random part of up to half the wait keeps the + * sessions of many hosts from connecting at the same moment. + */ +function backoffMs(attempt: number) { + const full = Math.min(500 * 2 ** attempt, 30_000) + return full / 2 + Math.random() * (full / 2) +} + +/** Resolves after `ms`, or at once when `signal` aborts. */ +function delay(ms: number, signal: AbortSignal) { + return new Promise((resolve) => { + // An aborted signal never fires `abort` again. + if (signal.aborted) return resolve() + const done = () => { + clearTimeout(timer) + signal.removeEventListener('abort', done) + resolve() + } + const timer = setTimeout(done, ms) + signal.addEventListener('abort', done, { once: true }) + }) +} + +/** The `needsApproval` of a server for the client and the connector. */ +function approvalOf(server: McpServerConfig) { + const { needsApproval } = server + if (typeof needsApproval !== 'boolean') return needsApproval + return () => needsApproval +} + +/** + * The tools of all servers, in server order. A tool whose name an earlier + * server already gives is dropped. Each pair of servers warns once. + */ +function keepFirstTools( + found: ReadonlyArray<{ name: string; tools: ReadonlyArray }>, + warned: Set, +) { + const owners = new Map() + const kept: Array = [] + for (const { name, tools } of found) { + for (const tool of tools) { + const owner = owners.get(tool.name) + if (owner === undefined) { + owners.set(tool.name, name) + kept.push(tool) + continue + } + const pair = JSON.stringify([owner, name]) + if (warned.has(pair)) continue + warned.add(pair) + console.warn( + `[@tanstack/ai-mcp] The MCP servers ${JSON.stringify(owner)} and ${JSON.stringify(name)} both give the tool ${JSON.stringify(tool.name)}. The model gets the tool of ${JSON.stringify(owner)}. Set toolName on one of the servers to rename its tools.`, + ) + } + } + return kept +} + +/** + * The tools of all servers, in server order. A server that gives a tool name + * of an earlier server gets the status `failed`, and none of its tools. + */ +async function failOnDuplicate( + found: ReadonlyArray<{ name: string; tools: ReadonlyArray }>, + setStatus: (status: McpServerStatus) => Promise, +) { + const owners = new Map() + const kept: Array = [] + for (const { name, tools } of found) { + const clash = tools.find((tool) => owners.has(tool.name)) + if (clash) { + const owner = owners.get(clash.name) + await setStatus({ + name, + status: 'failed', + error: `The MCP servers ${JSON.stringify(owner)} and ${JSON.stringify(name)} both give the tool ${JSON.stringify(clash.name)}. Set toolName on one of the servers to rename its tools.`, + toolCount: 0, + }) + continue + } + for (const tool of tools) owners.set(tool.name, name) + kept.push(...tools) + } + return kept +} + /** A tool call that runs now: its session, and the run that calls. */ export interface Caller { session: PluginSessionApi @@ -283,7 +561,7 @@ export interface Caller { /** * Puts each call of `tools` in `calling` while it runs, with its session - * and its run, so a spec 2025 elicitation goes to the run that calls. + * and its run, so a spec 2025 elicitation can find the run that calls. */ export function trackCaller( tools: ReadonlyArray, @@ -350,14 +628,24 @@ export function askForInput( } /** - * Answers the `elicitation/create` requests of a spec 2025 server: it asks - * for the newest call of `calling`, as a live wait of its run. With no - * call, it declines: only a tool call can wait for the user. + * Answers the `elicitation/create` requests of a spec 2025 server. The + * request does not name the tool call that it belongs to. So it asks only + * when the calls of `calling` all belong to one run of one session, as a + * live wait of that run. With no call, or with calls of more than one run + * (two users or tenants can share a server), it declines: a question must + * never go to the user of another run. */ export function elicitFor(calling: ReadonlyArray): OnElicit { return async (params) => { - const caller = calling.at(-1) + const [caller] = calling if (!caller) return { action: 'decline' } + const isOneRun = calling.every( + (other) => + other.session === caller.session && + other.run !== undefined && + other.run.runId === caller.run?.runId, + ) + if (!isOneRun) return { action: 'decline' } return elicitResult(await ask(caller.session, params, caller.run)) } } @@ -463,6 +751,8 @@ function statusLine(server: McpServerStatus) { switch (server.status) { case 'connecting': return `${server.name}: connecting` + case 'reconnecting': + return `${server.name}: reconnecting` case 'connected': { const unit = server.toolCount === 1 ? 'tool' : 'tools' return `${server.name}: connected (${server.toolCount} ${unit})` diff --git a/packages/ai-mcp/src/harness.ts b/packages/ai-mcp/src/harness.ts index 4b922908c1..4f9f8b9419 100644 --- a/packages/ai-mcp/src/harness.ts +++ b/packages/ai-mcp/src/harness.ts @@ -265,6 +265,18 @@ export async function createHarnessMcpServer

( // question, or a wait of a background agent run). `answer` or `resolve` // continues it. const waiting = new Map>() + /** + * Keep `operation` for its wait `waitId`. The entry goes when the + * operation settles, also when the wait ended another way: answered by + * another client, expired, or cancelled. + */ + const keepWaiting = (waitId: string, operation: Operation) => { + waiting.set(waitId, operation) + const forget = () => { + if (waiting.get(waitId) === operation) waiting.delete(waitId) + } + void operation.then(forget, forget) + } // The user of each request, by the token in its context. A token that // this server did not make is in no map, so a context cannot name a user. const callers = new WeakMap() @@ -379,7 +391,7 @@ export async function createHarnessMcpServer

( if (!outcome.done) { const snapshot = session.snapshot() for (const wait of snapshot.pendingWaits ?? []) { - if (wait.runId === operation.id) waiting.set(wait.waitId, operation) + if (wait.runId === operation.id) keepWaiting(wait.waitId, operation) } return { status: 'waiting', ...pending(snapshot, operation.id) } } @@ -576,6 +588,9 @@ export async function createHarnessMcpServer

( }).server( onSession(async (session, args, context) => { const questionId = textArg(args, 'questionId') + // Read before the answer: the work can settle at once after it, and + // then its entry is gone. + const operation = waiting.get(questionId) const receipt = await session.answer({ questionId, value: isRecord(args) ? args.value : undefined, @@ -585,7 +600,6 @@ export async function createHarnessMcpServer

( `The harness refused the answer: ${receipt.reason ?? 'no reason'}.`, ) } - const operation = waiting.get(questionId) waiting.delete(questionId) // A question that no tool call waits on: nothing to continue. if (operation === undefined) { @@ -766,7 +780,7 @@ export async function createHarnessMcpServer

( try { // A read: no session opens. const view = await observe(threadId, ctx.context) - const record = await view.getMedia(id) + const record = await view.getMedia({ id }) if (record === null) { throw new Error(`Media ${id} was not found in thread ${threadId}.`) } diff --git a/packages/ai-mcp/src/types.ts b/packages/ai-mcp/src/types.ts index 380cd742eb..f3c48667ac 100644 --- a/packages/ai-mcp/src/types.ts +++ b/packages/ai-mcp/src/types.ts @@ -195,6 +195,24 @@ export interface MCPClientOptions { * ``` */ needsApproval?: (tool: McpToolDef) => boolean + /** + * Runs when the connection closes and this client did not close it. For + * example, the process of a stdio server stops. `close()` does not call + * it. After it runs, the client cannot send requests: connect a new one. + * + * ```ts + * const mcp = await createMCPClient({ + * transport: { type: 'stdio', command: 'npx', args: ['my-server'] }, + * onClose: () => console.warn('The MCP server stopped.'), + * }) + * ``` + */ + onClose?: () => void + /** + * Runs when the server sends `notifications/tools/list_changed`. Call + * `tools()` again to get the new list. + */ + onToolListChanged?: () => void } export interface ToolsOptions { diff --git a/packages/ai-mcp/tests/client.test.ts b/packages/ai-mcp/tests/client.test.ts index 8bb42ece2f..37e3184897 100644 --- a/packages/ai-mcp/tests/client.test.ts +++ b/packages/ai-mcp/tests/client.test.ts @@ -486,6 +486,44 @@ describe('createMCPClient', () => { }) }) + it('calls onToolListChanged on a tools/list_changed notification', async () => { + const { server, clientTransport } = await makeServerWithChangingTools() + const onToolListChanged = vi.fn() + await using client = await createMCPClient({ + transport: clientTransport, + onToolListChanged, + }) + expect(onToolListChanged).not.toHaveBeenCalled() + await server.sendToolListChanged() + await vi.waitFor(() => expect(onToolListChanged).toHaveBeenCalledTimes(1)) + // The hook only reports. The next tools() call lists again. + expect((await client.tools()).map((tool) => tool.name)).toEqual(['tool_a']) + }) + + it('calls onClose when the server closes the connection', async () => { + const { server, clientTransport } = await makeServerWithChangingTools() + const onClose = vi.fn() + const client = await createMCPClient({ + transport: clientTransport, + onClose, + }) + await server.close() + await vi.waitFor(() => expect(onClose).toHaveBeenCalledTimes(1)) + await client.close() + expect(onClose).toHaveBeenCalledTimes(1) + }) + + it('does not call onClose for its own close()', async () => { + const { clientTransport } = await makeServerWithChangingTools() + const onClose = vi.fn() + const client = await createMCPClient({ + transport: clientTransport, + onClose, + }) + await client.close() + expect(onClose).not.toHaveBeenCalled() + }) + it('opens subscriptions/listen on a spec 2026 server that reports tool list changes', async () => { const { clientTransport, sawListen } = await makeModernChangingServer() await using client = await createMCPClientFromTransport(clientTransport) diff --git a/packages/ai-mcp/tests/connector-reconnect.test.ts b/packages/ai-mcp/tests/connector-reconnect.test.ts new file mode 100644 index 0000000000..8f0f2ee197 --- /dev/null +++ b/packages/ai-mcp/tests/connector-reconnect.test.ts @@ -0,0 +1,97 @@ +import { memoryCredentialStore } from '@tanstack/ai/byok/server' +import { afterEach, describe, expect, it, vi } from 'vitest' +import { memoryPersistence } from '@tanstack/ai-persistence' +import { createHarnessHost, defineHarness } from '@tanstack/ai-harness' +import { mcpConnector } from '../src/connector' +import { recorder } from './connector-helpers' +import type { AnyTool } from '@tanstack/ai' +import type { MCPClientOptions } from '../src/types' + +// Each connect of the connector gets a fake client. The test keeps the +// options of each connect, so it can close a connection or change its +// tool list like a server does. +const connects = vi.hoisted(() => ({ + options: [] as Array, + tools: [] as Array>, +})) +vi.mock('../src/client', async (importOriginal) => { + const actual = await importOriginal() + return { + ...actual, + connectTransport: async (options: unknown) => { + const index = connects.options.length + connects.options.push(options) + connects.tools.push(['echo']) + return { + tools: async () => + (connects.tools[index] ?? []).map( + (name): AnyTool => ({ + name: `demo_${name}`, + description: name, + execute: async () => name, + }), + ), + close: async () => {}, + } + }, + } +}) + +const cleanups: Array<() => unknown> = [] +afterEach(async () => { + for (const cleanup of cleanups.splice(0).reverse()) await cleanup() + connects.options.length = 0 + connects.tools.length = 0 +}) + +/** The callbacks that connect `index` gave the client. */ +function eventsOf(index: number) { + const options = connects.options[index] as Pick< + MCPClientOptions, + 'onClose' | 'onToolListChanged' + > + return options +} + +describe('mcpConnector connection events', () => { + it('lists the tools again after tools/list_changed, and connects again after a close', async () => { + const credentials = memoryCredentialStore() + await credentials.set({ + scope: { threadId: 't', userId: 'user-1' }, + id: 'demo', + credential: { type: 'oauth', accessToken: 'access-1' }, + }) + const host = createHarnessHost({ + persistence: memoryPersistence(), + credentials, + }) + cleanups.push(() => host.close()) + const model = recorder() + const session = await host.open( + defineHarness({ + name: 'test/connector-events', + adapter: model.adapter, + plugins: () => [ + mcpConnector({ id: 'demo', label: 'Demo', url: 'http://mcp.test' }), + ], + }), + { threadId: 't', principal: { id: 'user-1' } }, + ) + + await session.send('one') + expect(model.toolNames(0)).toEqual(['demo_echo']) + + // The server adds a tool and says so. + connects.tools[0] = ['echo', 'time'] + eventsOf(0).onToolListChanged?.() + await session.send('two') + expect(model.toolNames(1)).toEqual(['demo_echo', 'demo_time']) + expect(connects.options).toHaveLength(1) + + // The connection drops: the next turn connects again. + eventsOf(0).onClose?.() + await session.send('three') + expect(connects.options).toHaveLength(2) + expect(model.toolNames(2)).toEqual(['demo_echo']) + }) +}) diff --git a/packages/ai-mcp/tests/harness-elicit.test.ts b/packages/ai-mcp/tests/harness-elicit.test.ts index 140ed4439b..23f1f3a7b2 100644 --- a/packages/ai-mcp/tests/harness-elicit.test.ts +++ b/packages/ai-mcp/tests/harness-elicit.test.ts @@ -56,7 +56,7 @@ describe('an MCP elicitation is a live wait of the run that calls', () => { }) }) - it('asks a spec 2025 elicitation as a wait of the newest running call', async () => { + it('asks a spec 2025 elicitation as a wait of the run of the running call', async () => { const { session, ask } = askingSession({ city: 'Paris' }) const calling: Parameters[2] = [] const onElicit = elicitFor(calling) @@ -90,6 +90,36 @@ describe('an MCP elicitation is a live wait of the run that calls', () => { expect(ask).not.toHaveBeenCalled() }) + it('declines a spec 2025 elicitation while calls of two runs share the server, and asks nobody', async () => { + const first = askingSession({ city: 'Paris' }) + const second = askingSession({ city: 'Rome' }) + const calling: Parameters[2] = [ + { session: first.session, run: runOf('op-user-a') }, + { session: second.session, run: runOf('op-user-b') }, + ] + + await expect(elicitFor(calling)(form as never)).resolves.toEqual({ + action: 'decline', + }) + expect(first.ask).not.toHaveBeenCalled() + expect(second.ask).not.toHaveBeenCalled() + }) + + it('asks a spec 2025 elicitation when every running call is of one run', async () => { + const { session, ask } = askingSession({ city: 'Paris' }) + const run = runOf('op-chat-1') + const calling: Parameters[2] = [ + { session, run }, + { session, run }, + ] + + await expect(elicitFor(calling)(form as never)).resolves.toEqual({ + action: 'accept', + content: { city: 'Paris' }, + }) + expect(ask).toHaveBeenCalledTimes(1) + }) + it('declines a spec 2025 elicitation that comes when no call runs', async () => { const onElicit = elicitFor([]) await expect(onElicit(form as never)).resolves.toEqual({ diff --git a/packages/ai-mcp/tests/harness-plugin.test.ts b/packages/ai-mcp/tests/harness-plugin.test.ts index 445193a64c..785a2ce003 100644 --- a/packages/ai-mcp/tests/harness-plugin.test.ts +++ b/packages/ai-mcp/tests/harness-plugin.test.ts @@ -25,10 +25,10 @@ const childEnv = 'TANSTACK_AI_MCP_PLUGIN_CHILD' const cleanups: Array<() => unknown> = [] /** A server with one `echo` tool, in this process. */ -function echoServer() { +function echoServer(description = 'Echo text') { const echo = toolDefinition({ name: 'echo', - description: 'Echo text', + description, inputSchema: z.object({ text: z.string() }), }).server(async ({ text }) => text) return createMCPServer({ name: 'echo', version: '1.0.0', tools: [echo] }) @@ -37,7 +37,10 @@ function echoServer() { /** An `http` server config whose requests go to `handle`, with no network. */ function httpServer( handle: (request: Request) => Promise, - options: Pick = {}, + options: Pick< + McpServerConfig, + 'timeoutMs' | 'codeMode' | 'needsApproval' | 'toolName' + > = {}, ) { const config: McpServerConfig = { type: 'http', @@ -209,6 +212,120 @@ if (process.env[childEnv] === '1') { ]) }) + it.each([ + ['true', true], + ['a function', (tool: { name: string }) => tool.name === 'echo'], + ])( + 'asks for approval before the tools of a server with needsApproval %s', + async (_name, needsApproval) => { + const gated = echoServer() + const open = echoServer() + const { session, model } = await openWith({ + gated: httpServer((request) => gated.fetch(request), { + needsApproval, + }), + open: httpServer((request) => open.fetch(request)), + }) + await session.send('hi') + + const marks = (model.calls[0]?.tools ?? []).map((tool) => [ + tool.name, + tool.needsApproval === true, + ]) + expect(marks).toEqual([ + ['gated_echo', true], + ['open_echo', false], + ]) + }, + ) + + it('names the tools of a server with its toolName', async () => { + const files = echoServer() + const { session, model } = await openWith({ + files: httpServer((request) => files.fetch(request), { + toolName: (tool) => `fs_${tool.name}`, + }), + }) + await session.send('hi') + + expect(model.toolNames(0)).toEqual(['fs_echo']) + }) + + it('keeps the tool of the first server when two servers give one tool name, and warns once', async () => { + const warn = vi.spyOn(console, 'warn').mockImplementation(() => {}) + cleanups.push(() => warn.mockRestore()) + const first = echoServer('First echo') + const second = echoServer('Second echo') + const { session, model } = await openWith({ + first: httpServer((request) => first.fetch(request), { + toolName: (tool) => tool.name, + }), + second: httpServer((request) => second.fetch(request), { + toolName: (tool) => tool.name, + }), + }) + await session.send('hi') + await session.send('again') + + for (const call of [0, 1]) { + const tools = (model.calls[call]?.tools ?? []).map((tool) => [ + tool.name, + tool.description, + ]) + expect(tools).toEqual([['echo', 'First echo']]) + } + expect(warn.mock.calls).toEqual([ + [ + '[@tanstack/ai-mcp] The MCP servers "first" and "second" both give the tool "echo". The model gets the tool of "first". Set toolName on one of the servers to rename its tools.', + ], + ]) + }) + + it("fails the later server when two servers give one tool name, with onDuplicateTool 'error'", async () => { + const first = echoServer('First echo') + const second = echoServer('Second echo') + const third = echoServer('Third echo') + const model = recorder() + const host = createHarnessHost({ persistence: memoryPersistence() }) + cleanups.push(() => host.close()) + const harness = defineHarness({ + name: 'test/mcp-plugin', + adapter: model.adapter, + plugins: () => [ + mcp({ + onDuplicateTool: 'error', + servers: { + first: httpServer((request) => first.fetch(request), { + toolName: (tool) => tool.name, + }), + second: httpServer((request) => second.fetch(request)), + third: httpServer((request) => third.fetch(request), { + toolName: (tool) => tool.name, + }), + }, + }), + ], + }) + const session = await host.open(harness, { threadId: 't' }) + await session.send('hi') + + const tools = (model.calls[0]?.tools ?? []).map((tool) => [ + tool.name, + tool.description, + ]) + expect(tools).toEqual([ + ['echo', 'First echo'], + ['second_echo', 'Second echo'], + ]) + await expect(session.command({ name: 'mcp' })).resolves.toBe( + [ + 'first: connected (1 tool)', + 'second: connected (1 tool)', + 'third: failed: The MCP servers "first" and "third" both give the tool "echo". Set toolName on one of the servers to rename its tools.', + ].join('\n'), + ) + }) + it.each([ [ 'accepts', diff --git a/packages/ai-mcp/tests/harness-reconnect.test.ts b/packages/ai-mcp/tests/harness-reconnect.test.ts new file mode 100644 index 0000000000..0ca80d4c96 --- /dev/null +++ b/packages/ai-mcp/tests/harness-reconnect.test.ts @@ -0,0 +1,268 @@ +import { afterEach, describe, expect, it, vi } from 'vitest' +import { z } from 'zod' +import { InMemoryTransport } from '@modelcontextprotocol/client' +import { McpServer } from '@modelcontextprotocol/server' +import { memoryPersistence } from '@tanstack/ai-persistence' +import { createHarnessHost, defineHarness } from '@tanstack/ai-harness' +import { mcp } from '../src/harness' +import { recorder } from './connector-helpers' +import type { HarnessSession } from '@tanstack/ai-harness' + +// The stdio transport starts a child process. Here each start is an SDK +// server in this process, so a test can stop it and start the next one. +const stdio = vi.hoisted(() => ({ + start: (): unknown => { + throw new Error('stdio.start is not set') + }, +})) +vi.mock('../src/stdio', () => ({ stdioTransport: () => stdio.start() })) + +/** One start of the stdio server. */ +interface Child { + server: McpServer + /** The next start fails to connect, like a child that exits at once. */ + refused: boolean +} + +const children: Array = [] +// How many of the next starts fail. +let refuse = 0 + +stdio.start = () => { + const [clientSide, serverSide] = InMemoryTransport.createLinkedPair() + const server = new McpServer({ name: 'local', version: '1.0.0' }) + const number = children.length + 1 + server.registerTool( + 'pid', + { + description: 'Name this start of the server', + inputSchema: z.object({ text: z.string() }), + }, + async () => ({ + content: [{ type: 'text' as const, text: `child:${number}` }], + }), + ) + const refused = refuse > 0 + refuse -= 1 + children.push({ server, refused }) + if (refused) void serverSide.close() + else void server.connect(serverSide) + return clientSide +} + +const cleanups: Array<() => unknown> = [] + +afterEach(async () => { + vi.useRealTimers() + vi.restoreAllMocks() + for (const cleanup of cleanups.splice(0).reverse()) await cleanup() + children.length = 0 + refuse = 0 +}) + +/** + * Opens `count` sessions of one harness with one stdio server `local`. The + * model calls `local_pid` on the first model call of each turn. + */ +async function openSessions(count: number) { + const host = createHarnessHost({ persistence: memoryPersistence() }) + cleanups.push(() => host.close()) + const model = recorder('local_pid', (call) => call % 2 === 1) + const harness = defineHarness({ + name: 'test/mcp-reconnect', + adapter: model.adapter, + plugins: () => [ + mcp({ + servers: { local: { type: 'stdio', command: 'local-server' } }, + }), + ], + }) + const sessions = await Promise.all( + Array.from({ length: count }, (_, index) => + host.open(harness, { threadId: `t-${index}` }), + ), + ) + return { host, sessions, model } +} + +function statusOf(session: HarnessSession) { + return session.command({ name: 'mcp' }) +} + +/** The text of the last model call: it has the result of the tool call. */ +function lastSeen(model: ReturnType) { + return JSON.stringify(model.calls.at(-1)?.messages) +} + +/** Stop the current start of the server, with fake timers for the backoff. */ +async function stopServer() { + vi.useFakeTimers({ toFake: ['setTimeout', 'clearTimeout'] }) + await children.at(-1)?.server.close() + await vi.advanceTimersByTimeAsync(0) +} + +describe('mcp reconnect', () => { + it('connects a shared stdio server again for every session, and gives its tools at the next turn', async () => { + const { sessions, model } = await openSessions(2) + const [first, second] = sessions + if (!first || !second) throw new Error('two sessions expected') + await first.send('which server?') + expect(lastSeen(model)).toContain('child:1') + + await stopServer() + await expect(statusOf(first)).resolves.toBe('local: reconnecting') + await expect(statusOf(second)).resolves.toBe('local: reconnecting') + + // The first wait is 0.5 s at most. + await vi.advanceTimersByTimeAsync(500) + vi.useRealTimers() + await vi.waitFor(async () => + expect(await statusOf(second)).toBe('local: connected (1 tool)'), + ) + await expect(statusOf(first)).resolves.toBe('local: connected (1 tool)') + // One new start for both sessions. + expect(children).toHaveLength(2) + + await second.send('which server now?') + expect(lastSeen(model)).toContain('child:2') + await first.send('and here?') + expect(lastSeen(model)).toContain('child:2') + }) + + it('waits 0.5 s, 1 s, then 2 s (with jitter) between the tries', async () => { + const { sessions } = await openSessions(1) + const [session] = sessions + if (!session) throw new Error('a session expected') + await session.send('hi') + + refuse = 2 + // No jitter: each wait is half of 0.5 s, 1 s, 2 s, and so on. The + // harness uses Math.random for its ids, so the mock ends before the + // next session call. + const random = vi.spyOn(Math, 'random').mockReturnValue(0) + await stopServer() + // Attempt 1 after 250 ms (refused), 2 after 500 ms more (refused), 3 + // after 1000 ms more. + for (const [wait, starts] of [ + [249, 1], + [1, 2], + [499, 2], + [1, 3], + [999, 3], + [1, 4], + ] as const) { + await vi.advanceTimersByTimeAsync(wait) + expect(children).toHaveLength(starts) + } + random.mockRestore() + vi.useRealTimers() + await vi.waitFor(async () => + expect(await statusOf(session)).toBe('local: connected (1 tool)'), + ) + }) + + it('keeps the backoff when a server drops again right after it connects', async () => { + const { sessions } = await openSessions(1) + const [session] = sessions + if (!session) throw new Error('a session expected') + await session.send('hi') + + const random = vi.spyOn(Math, 'random').mockReturnValue(0) + await stopServer() + await vi.advanceTimersByTimeAsync(250) + expect(children).toHaveLength(2) + await vi.waitFor(async () => + expect(await statusOf(session)).toBe('local: connected (1 tool)'), + ) + + // The new start crashes at once: the next wait is 0.5 s, not 0.25 s. + await stopServer() + await vi.advanceTimersByTimeAsync(499) + expect(children).toHaveLength(2) + await vi.advanceTimersByTimeAsync(1) + expect(children).toHaveLength(3) + random.mockRestore() + }) + + it('leaves no timer when the host closes while a connect runs', async () => { + const { host, sessions } = await openSessions(1) + await sessions[0]?.send('hi') + + // The next start hangs until the test refuses it. + let refuseNow = () => {} + const hanging = new Promise((_, reject) => { + refuseNow = () => reject(new Error('refused')) + }) + const start = stdio.start + cleanups.push(() => { + stdio.start = start + }) + stdio.start = () => ({ + start: () => hanging, + send: async () => {}, + close: async () => {}, + }) + await stopServer() + await vi.advanceTimersByTimeAsync(500) + + const closing = host.close() + refuseNow() + await closing + await vi.advanceTimersByTimeAsync(0) + expect(vi.getTimerCount()).toBe(0) + }) + + it('connects again after a failed first connect, and gives its tools at a later turn', async () => { + refuse = 1 + const { sessions, model } = await openSessions(1) + const [session] = sessions + if (!session) throw new Error('a session expected') + expect(await statusOf(session)).toMatch(/^local: failed: /) + + // The first wait is 0.5 s at most. + await vi.waitFor( + async () => + expect(await statusOf(session)).toBe('local: connected (1 tool)'), + { timeout: 5_000 }, + ) + expect(children).toHaveLength(2) + await session.send('which server?') + expect(lastSeen(model)).toContain('child:2') + }) + + it('stops trying when the host closes', async () => { + const { host, sessions } = await openSessions(1) + await sessions[0]?.send('hi') + + refuse = Number.POSITIVE_INFINITY + await stopServer() + await vi.advanceTimersByTimeAsync(500) + expect(children).toHaveLength(2) + + await host.close() + await vi.advanceTimersByTimeAsync(120_000) + expect(children).toHaveLength(2) + }) +}) + +describe('mcp tools/list_changed', () => { + it('lists the tools again at the next turn, and updates the tool count', async () => { + const { sessions, model } = await openSessions(1) + const [session] = sessions + if (!session) throw new Error('a session expected') + await session.send('hi') + expect(model.toolNames(0)).toEqual(['local_pid']) + + const server = children.at(-1)?.server + server?.registerTool( + 'time', + { description: 'The time', inputSchema: z.object({}) }, + async () => ({ content: [{ type: 'text' as const, text: 'noon' }] }), + ) + await server?.sendToolListChanged() + await session.send('and now?') + + expect(model.toolNames(2)).toEqual(['local_pid', 'local_time']) + await expect(statusOf(session)).resolves.toBe('local: connected (2 tools)') + }) +}) diff --git a/packages/ai-mcp/tests/harness.test.ts b/packages/ai-mcp/tests/harness.test.ts index 81bcd7862d..bb3ed89a0b 100644 --- a/packages/ai-mcp/tests/harness.test.ts +++ b/packages/ai-mcp/tests/harness.test.ts @@ -1318,6 +1318,36 @@ describe('createHarnessMcpServer', () => { expect(answered.content).toEqual([{ type: 'text', text: 'Confirmed.' }]) }) + it('forgets a waiting question when another client answers it', async () => { + // The server keeps the waiting work in a Map that it does not export. + // Find it: the Map that gets the question id with an operation. + const sets = vi.spyOn(Map.prototype, 'set') + const { server, session } = await harnessServer([]) + const { client } = await connect(server) + const waiting = await client.callTool({ + name: 'command_confirm', + arguments: {}, + }) + const [question] = pendingShape.parse(waiting.structuredContent).questions + if (question === undefined) throw new Error('The command asked nothing') + const index = sets.mock.calls.findIndex( + ([key, value]) => + key === question.id && + typeof value === 'object' && + value !== null && + 'stream' in value, + ) + const map: unknown = sets.mock.contexts[index] + sets.mockRestore() + if (!(map instanceof Map)) throw new Error('No waiting map found') + expect(map.has(question.id)).toBe(true) + + // Another client answers, so the MCP `answer` tool never runs. + const live = await session() + await live.answer({ questionId: question.id, value: true }) + await waitFor(() => expect(map.has(question.id)).toBe(false)) + }) + it('returns the result of its own turn when another operation asks a question', async () => { const held = heldTurn(textTurn('Done.')) const { server, session } = await harnessServer([held]) diff --git a/packages/ai-persistence/src/memory.ts b/packages/ai-persistence/src/memory.ts index a8038e6264..c169ac4c84 100644 --- a/packages/ai-persistence/src/memory.ts +++ b/packages/ai-persistence/src/memory.ts @@ -867,6 +867,9 @@ class MemoryLogStore implements LogStore { return { epoch: owner.epoch } } const epoch = (owner?.epoch ?? UNOWNED_EPOCH) + 1 + // Another owner that did not release: its ownership expired. + const expired = + owner !== undefined && !owner.released && owner.ownerId !== ownerId this.backend.owners.set(logId, { harness, ownerId, @@ -874,7 +877,7 @@ class MemoryLogStore implements LogStore { expiresAt: now + ttlMs, released: false, }) - return { epoch } + return expired ? { epoch, expired: true } : { epoch } } async renew({ diff --git a/packages/ai-persistence/src/testkit/conformance.ts b/packages/ai-persistence/src/testkit/conformance.ts index 580c87ccac..9a0b8454f0 100644 --- a/packages/ai-persistence/src/testkit/conformance.ts +++ b/packages/ai-persistence/src/testkit/conformance.ts @@ -1956,7 +1956,7 @@ export function runPersistenceConformance( const { taken } = await handOver(store) - expect(taken).toEqual({ epoch: 2 }) + expect(taken).toEqual({ epoch: 2, expired: true }) }) it('fences a write from the previous owner', async (ctx) => { diff --git a/packages/ai-persistence/src/types.ts b/packages/ai-persistence/src/types.ts index 8fcffae160..89d98d5438 100644 --- a/packages/ai-persistence/src/types.ts +++ b/packages/ai-persistence/src/types.ts @@ -851,6 +851,8 @@ export interface LogStore { * - Else increments the epoch, stores `harness`, clears the wake time, and * returns the new epoch. The epoch never goes down, also after `release` * and `deleteLog`. + * - When the log had another owner whose ownership expired with no + * `release`, the result has `expired: true`: that owner stopped. * * Rejects with a `RangeError` when `ttlMs` is not more than 0. */ @@ -860,7 +862,7 @@ export interface LogStore { ownerId: string ttlMs: number affinity?: string - }) => Promise<{ epoch: number } | null> + }) => Promise<{ epoch: number; expired?: boolean } | null> /** * Extend the ownership to `ttlMs` milliseconds from now. Returns `true` only * when the log is not released, `epoch` is the current epoch, and `ownerId` diff --git a/packages/ai-persistence/tests/memory.test.ts b/packages/ai-persistence/tests/memory.test.ts index 0d94e6ce5e..747259de72 100644 --- a/packages/ai-persistence/tests/memory.test.ts +++ b/packages/ai-persistence/tests/memory.test.ts @@ -307,7 +307,7 @@ describe('memoryLogBackend({ now })', () => { clock.now = 1_500 expect( await store.acquire({ logId, harness, ownerId: 'b', ttlMs: 500 }), - ).toEqual({ epoch: 2 }) + ).toEqual({ epoch: 2, expired: true }) }) it('lists a released log for its wake only when its clock reaches the wake time', async () => { diff --git a/packages/ai-skills/src/harness.ts b/packages/ai-skills/src/harness.ts index 200a0f113c..68cf5b7871 100644 --- a/packages/ai-skills/src/harness.ts +++ b/packages/ai-skills/src/harness.ts @@ -7,6 +7,7 @@ import { estimateTokens, withSkills } from './middleware' import { skillDirectory } from './node' import type { FSWatcher } from 'node:fs' import type { AnyChatMiddleware } from '@tanstack/ai' +import type { PluginResources, PluginSetupContext } from '@tanstack/ai-harness' import type { SkillsOptions } from './middleware' import type { ModelFamily, SkillMetadata, SkillSource } from './types' @@ -30,23 +31,10 @@ interface FoundSkill { dir: string } -/** - * The skills of `dirs`, by name, in folder order. The first folder wins a - * name. A folder that is missing or cannot be read gives no skills, and a - * broken `SKILL.md` is left out. - */ -async function scan(dirs: ReadonlyArray) { - const found = new Map() - for (const dir of dirs) { - // A folder that throws, at once or later, gives no skills. - const listed = await Promise.resolve() - .then(() => skillDirectory(dir, { strict: false }).list()) - .catch(() => []) - for (const skill of listed) { - if (!found.has(skill.name)) found.set(skill.name, { skill, dir }) - } - } - return found +/** The `code` of a Node.js error, for example `ENOENT`. */ +function codeOf(error: unknown) { + const hasCode = error instanceof Error && 'code' in error + return hasCode ? error.code : undefined } /** The catalog with only the first sentence of each description. */ @@ -68,15 +56,328 @@ function skillPrompt(name: string, input: unknown) { : `Use the ${name} skill: ${task}` } +/** + * The part of {@link skills} that every session of a harness on a host + * shares: the watchers of the folders. A change only tells the sessions to + * sync. Each sync reads the skills from disk again. + */ +async function sharedFolders( + dirs: ReadonlyArray, + resources: PluginResources, +) { + /** The watcher of each folder, and the path it watches. */ + const watchers = await resources.acquire( + () => new Map(), + (open) => { + for (const { watcher } of open.values()) watcher.close() + open.clear() + }, + ) + /** The folders that cannot be watched. They were warned about once. */ + const failed = new Set() + /** The sessions to tell when a folder changed. */ + const listeners = new Set<() => void>() + + // A save often fires several watch events: wait until they stop. + let timer: ReturnType | undefined + const soon = () => { + clearTimeout(timer) + timer = setTimeout(() => { + for (const listener of listeners) listener() + }, 200) + timer.unref?.() + } + resources.signal.addEventListener('abort', () => clearTimeout(timer)) + + const drop = (dir: string) => { + watchers.get(dir)?.watcher.close() + watchers.delete(dir) + } + + /** Watch `path` for the folder `dir`. Throws when `watch` throws. */ + const start = (dir: string, path: string, recursive: boolean) => { + const watcher = watch( + path, + { recursive, persistent: false }, + (event, name) => { + // The watched folder itself went away (a delete or a move). Its + // watcher is dead, so the next sync watches it again. + if (event === 'rename' && name === basename(path)) drop(dir) + soon() + }, + ) + watcher.on('error', () => { + drop(dir) + soon() + }) + watchers.set(dir, { watcher, path }) + } + + /** Watch the folder `dir` at `path`: the folder itself, or a parent. */ + const watchAt = (dir: string, path: string) => { + const isOwn = path === dir + try { + return start(dir, path, isOwn) + } catch (error) { + // ponytail: Linux has no recursive watch before Node 19.1. Then only + // the folder is watched: an added or removed skill shows at once, an + // edited SKILL.md at the next turn. + if (codeOf(error) !== 'ERR_FEATURE_UNAVAILABLE_ON_PLATFORM') throw error + console.warn( + `skills: cannot watch the subfolders of ${dir} on this platform. An edited SKILL.md shows at the next turn.`, + ) + return start(dir, path, false) + } + } + + /** + * Watch each skill folder. While a folder does not exist, watch its + * nearest folder that does, so the skills show as soon as you make the + * folders, one level at a time. + */ + const watchFolders = () => { + if (resources.signal.aborted) return + for (const dir of dirs) { + if (failed.has(dir)) continue + // Some watchers stop with no event when their folder goes away, so + // check that the folder is still there. + let path = dir + while (!existsSync(path) && dirname(path) !== path) path = dirname(path) + if (watchers.get(dir)?.path === path) continue + drop(dir) + try { + watchAt(dir, path) + } catch (error) { + // The folder went away since the check. The next sync tries again. + if (codeOf(error) === 'ENOENT') continue + failed.add(dir) + const reason = error instanceof Error ? error.message : String(error) + console.warn( + `skills: cannot watch ${dir} (${reason}). Its skills change at the next turn.`, + ) + } + } + } + + /** + * The skills of `dirs`, by name, in folder order. The first folder wins a + * name. A folder that is missing or cannot be read gives no skills, and a + * broken `SKILL.md` is left out. + */ + const read = async () => { + watchFolders() + const found = new Map() + for (const dir of dirs) { + // A folder that throws, at once or later, gives no skills. + const listed = await Promise.resolve() + .then(() => skillDirectory(dir, { strict: false }).list()) + .catch(() => []) + for (const skill of listed) { + if (!found.has(skill.name)) found.set(skill.name, { skill, dir }) + } + } + return found + } + + /** The source of the folder that has the skill `name` now. */ + const folderOf = async (name: string) => { + const entry = (await read()).get(name) + if (!entry) { + throw new Error(`No skill named "${name}" in ${dirs.join(', ')}.`) + } + return skillDirectory(entry.dir, { strict: false }) + } + + const source: SkillSource = { + list: async () => [...(await read()).values()].map((entry) => entry.skill), + load: async (name) => (await folderOf(name)).load(name), + listResources: async (name) => + (await (await folderOf(name)).listResources?.(name)) ?? [], + listScripts: async (name) => + (await (await folderOf(name)).listScripts?.(name)) ?? [], + readResource: async (name, path) => { + const folder = await folderOf(name) + if (!folder.readResource) { + throw new Error(`Skill "${name}" has no resource files.`) + } + return folder.readResource(name, path) + }, + } + + watchFolders() + + return { + source, + read, + /** Call `listener` after a folder changed. Returns a function that stops it. */ + onChange: (listener: () => void) => { + listeners.add(listener) + return () => { + listeners.delete(listener) + } + }, + } +} + +/** The part of {@link skills} that each session sets up. */ +async function sessionPart( + ctx: PluginSetupContext, + folders: Awaited>, + options: SkillsPluginOptions, +) { + const { dirs, reserved = [], ...skillsOptions } = options + // `skills` is this plugin's own command, so a skill named `skills` gets + // the prefix too. + const taken = new Set(['skills', ...reserved]) + // The skill commands of this session: command name to skill name and the + // description the command has. + const mine = new Map() + // Skills called with their command stay in the model's list. + const called = new Set() + let found = new Map() + const closed = () => ctx.resources.signal.aborted + + const skillCommand = (name: string, description: string) => + defineCommand({ + description: `Skill: ${firstSentence(description)}`, + input: { + type: 'string', + description: 'What the skill should do. Optional.', + }, + // The command's own session, so the turn runs as the user who ran the + // command. + run: (input: unknown, command) => { + called.add(name) + void command.session.send(skillPrompt(name, input)) + }, + }) + + /** Make the commands match the skills in the folders now. */ + const update = async () => { + if (closed()) return + found = await folders.read() + if (closed()) return + const next = new Map() + for (const { skill } of found.values()) { + const isFree = + !taken.has(skill.name) && + (mine.has(skill.name) || !ctx.commands.has(skill.name)) + next.set(isFree ? skill.name : `skill:${skill.name}`, { + name: skill.name, + description: skill.description, + }) + } + for (const [command, entry] of mine) { + if (next.get(command)?.name === entry.name) continue + ctx.commands.delete(command) + mine.delete(command) + } + for (const [command, entry] of next) { + const isSame = mine.get(command)?.description === entry.description + if (isSame) continue + try { + ctx.commands.set(command, skillCommand(entry.name, entry.description)) + mine.set(command, entry) + } catch { + // Another plugin took `skill:` too. The skill stays in the + // model's list, with no command. + } + } + } + // Updates run one after the other. + let updating = Promise.resolve() + const sync = () => { + updating = updating.then(update).catch(() => {}) + return updating + } + + await ctx.resources.acquire( + () => folders.onChange(() => void sync()), + (stop) => stop(), + ) + // Commands of the plugins after this one exist only once all are set up, + // so the first sync waits for that. + void ctx.commands.ready.then(sync) + + const hasOwnCatalog = + skillsOptions.renderCatalog !== undefined || + skillsOptions.instructionTemplate !== undefined + // The catalog that `withSkills` renders, so the count matches it. + const render = + skillsOptions.renderCatalog ?? + (hasOwnCatalog ? renderCatalog : shortCatalog) + + /** + * Over the token cap: the called skills first, then the folder order, + * until the list is full. The XML list of Anthropic models is the longest + * form, so the count uses it. + */ + const keepEarlier = (listed: Array, limit: number) => { + const ordered = [ + ...listed.filter((skill) => called.has(skill.name)), + ...listed.filter((skill) => !called.has(skill.name)), + ] + const kept: Array = [] + for (const skill of ordered) { + const tooLong = + estimateTokens(render([...kept, skill], 'anthropic')) > limit + if (tooLong) break + kept.push(skill) + } + return kept + } + + // The sync runs first, so the commands match the folders at the start of + // each turn. + const middleware: Array = [ + { name: 'tanstack/skills:sync', setup: () => sync() }, + withSkills(folders.source, { + ...skillsOptions, + ...(hasOwnCatalog ? {} : { renderCatalog: shortCatalog }), + onLimitExceeded: skillsOptions.onLimitExceeded ?? keepEarlier, + }), + ] + + return { + middleware, + commands: { + skills: defineCommand({ + description: 'List the skills, their folders, and their commands', + run: async () => { + await sync() + if (found.size === 0) { + return `No skills yet. Add a folder with a SKILL.md to one of: ${dirs.join(', ')}` + } + const commandOf = new Map( + [...mine].map(([command, entry]) => [entry.name, command]), + ) + const lines = [...found.values()].map(({ skill, dir }) => { + const command = commandOf.get(skill.name) + const name = command ? `/${command}` : skill.name + const note = + command && command !== skill.name + ? ` (/${skill.name} is taken)` + : '' + return `${name} ${dir}${note}` + }) + return lines.join('\n') + }, + }), + }, + } +} + /** * A harness plugin that gives the model the skills of `dirs`, and makes each * skill a command: `/ [task]` starts a turn that tells the model to use * the skill. A skill whose name another command has is `/skill:`. * `/skills` lists every skill, its folder, and its command. * - * The plugin watches the folders, and it lists them again at the start of - * each turn. A skill that you add, change, or remove while the session runs - * changes the commands at once, with no restart. + * The plugin watches each folder once for each host, and every session of + * the harness shares the watchers. It reads the folders again at the start + * of each turn. A skill that you add, change, or remove while a session runs + * changes the commands at once, with no restart. When a folder cannot be + * watched, the plugin warns once, and its skills change at the next turn. * * The model's skill list has one line for each skill: its name and the first * sentence of its description. When the list is longer than @@ -102,258 +403,15 @@ function skillPrompt(name: string, input: unknown) { * ``` */ export function skills(options: SkillsPluginOptions) { - const { dirs, reserved = [], ...skillsOptions } = options - // `skills` is this plugin's own command, so a skill named `skills` gets - // the prefix too. - const taken = new Set(['skills', ...reserved]) - return definePlugin({ name: 'tanstack/skills', - setup: async (ctx) => { - // The skill commands of this plugin: command name to skill name and - // the description the command has. - const mine = new Map() - // Skills called with their command stay in the model's list. - const called = new Set() - let found = new Map() - const watchers = await ctx.resources.acquire( - () => new Map(), - (open) => { - for (const watcher of open.values()) watcher.close() - open.clear() - }, - ) - const closed = () => ctx.resources.signal.aborted - - const skillCommand = (name: string, description: string) => - defineCommand({ - description: `Skill: ${firstSentence(description)}`, - input: { - type: 'string', - description: 'What the skill should do. Optional.', - }, - // The command's own session, so the turn runs as the user who - // ran the command. - run: (input: unknown, command) => { - called.add(name) - void command.session.send(skillPrompt(name, input)) - }, - }) - - /** The folder each parent watcher watches, by key. */ - const parents = new Map() - const drop = (key: string) => { - watchers.get(key)?.close() - watchers.delete(key) - parents.delete(key) - } - - /** Watch `path` under `key`. `false` when `path` does not exist. */ - const tryWatch = ( - key: string, - path: string, - recursive: boolean, - ): boolean => { - try { - const watcher = watch( - path, - { recursive, persistent: false }, - (event, name) => { - // The folder itself went away (a delete or a move). Its - // watcher is dead, so the next sync watches it again. - if (event === 'rename' && name === basename(path)) drop(key) - soon() - }, - ) - watcher.on('error', () => { - drop(key) - soon() - }) - watchers.set(key, watcher) - return true - } catch (error) { - // ponytail: Linux has no recursive watch before Node 19.1. Then - // only the folder is watched: an added or removed skill shows at - // once, an edited SKILL.md at the next turn. Watch each skill - // folder too if edits must show at once there. - const isUnavailable = - error instanceof Error && - 'code' in error && - error.code === 'ERR_FEATURE_UNAVAILABLE_ON_PLATFORM' - return recursive && isUnavailable ? tryWatch(key, path, false) : false - } - } - - /** - * Watch each skill folder. While a folder does not exist, watch its - * nearest folder that does, so the skills show as soon as you make the - * folders, one level at a time. - */ - const watchFolders = () => { - for (const dir of dirs) { - // Some watchers stop with no event when their folder goes away. - if (watchers.has(dir) && !existsSync(dir)) drop(dir) - if (watchers.has(dir)) continue - const parentKey = `${dir} (parent)` - if (tryWatch(dir, dir, true)) { - drop(parentKey) - continue - } - let parent = dirname(dir) - while (!existsSync(parent) && dirname(parent) !== parent) { - parent = dirname(parent) - } - if (parents.get(parentKey) === parent) continue - drop(parentKey) - if (tryWatch(parentKey, parent, false)) parents.set(parentKey, parent) - } - } - - /** Make the commands match the skills in the folders now. */ - const update = async () => { - if (closed()) return - found = await scan(dirs) - if (closed()) return - watchFolders() - const next = new Map() - for (const { skill } of found.values()) { - const isFree = - !taken.has(skill.name) && - (mine.has(skill.name) || !ctx.commands.has(skill.name)) - next.set(isFree ? skill.name : `skill:${skill.name}`, { - name: skill.name, - description: skill.description, - }) - } - for (const [command, entry] of mine) { - if (next.get(command)?.name === entry.name) continue - ctx.commands.delete(command) - mine.delete(command) - } - for (const [command, entry] of next) { - const isSame = mine.get(command)?.description === entry.description - if (isSame) continue - try { - ctx.commands.set( - command, - skillCommand(entry.name, entry.description), - ) - mine.set(command, entry) - } catch { - // Another plugin took `skill:` too. The skill stays in the - // model's list, with no command. - } - } - } - // Updates run one after the other. - let updating = Promise.resolve() - const sync = () => { - updating = updating.then(update).catch(() => {}) - return updating - } - - // A save often fires several watch events: wait until they stop. - let timer: ReturnType | undefined - function soon() { - clearTimeout(timer) - timer = setTimeout(() => void sync(), 200) - timer.unref?.() - } - ctx.resources.signal.addEventListener('abort', () => clearTimeout(timer)) - - // Commands of the plugins after this one exist only once all are set - // up, so the first sync waits for that. - void ctx.commands.ready.then(sync) - - const source: SkillSource = { - list: async () => - [...(await scan(dirs)).values()].map((entry) => entry.skill), - load: async (name) => (await folderOf(name)).load(name), - listResources: async (name) => - (await (await folderOf(name)).listResources?.(name)) ?? [], - listScripts: async (name) => - (await (await folderOf(name)).listScripts?.(name)) ?? [], - readResource: async (name, path) => { - const folder = await folderOf(name) - if (!folder.readResource) { - throw new Error(`Skill "${name}" has no resource files.`) - } - return folder.readResource(name, path) - }, - } - /** The source of the folder that has the skill `name` now. */ - async function folderOf(name: string) { - const entry = (await scan(dirs)).get(name) - if (!entry) - throw new Error(`No skill named "${name}" in ${dirs.join(', ')}.`) - return skillDirectory(entry.dir, { strict: false }) - } - - const hasOwnCatalog = - skillsOptions.renderCatalog !== undefined || - skillsOptions.instructionTemplate !== undefined - // The catalog that `withSkills` renders, so the count matches it. - const render = - skillsOptions.renderCatalog ?? - (hasOwnCatalog ? renderCatalog : shortCatalog) - - /** - * Over the token cap: the called skills first, then the folder order, - * until the list is full. The XML list of Anthropic models is the - * longest form, so the count uses it. - */ - const keepEarlier = (listed: Array, limit: number) => { - const ordered = [ - ...listed.filter((skill) => called.has(skill.name)), - ...listed.filter((skill) => !called.has(skill.name)), - ] - const kept: Array = [] - for (const skill of ordered) { - const tooLong = - estimateTokens(render([...kept, skill], 'anthropic')) > limit - if (tooLong) break - kept.push(skill) - } - return kept - } - - // The sync runs first, so the commands match the folders at the start - // of each turn. - const middleware: Array = [ - { name: 'tanstack/skills:sync', setup: () => sync() }, - withSkills(source, { - ...skillsOptions, - ...(hasOwnCatalog ? {} : { renderCatalog: shortCatalog }), - onLimitExceeded: skillsOptions.onLimitExceeded ?? keepEarlier, - }), - ] - + // The folder watchers are shared by the sessions of the host. The + // commands and the model's skill list are per session. + lifetime: 'host', + setup: async (host) => { + const folders = await sharedFolders(options.dirs, host.resources) return { - middleware, - commands: { - skills: defineCommand({ - description: 'List the skills, their folders, and their commands', - run: async () => { - await sync() - if (found.size === 0) { - return `No skills yet. Add a folder with a SKILL.md to one of: ${dirs.join(', ')}` - } - const commandOf = new Map( - [...mine].map(([command, entry]) => [entry.name, command]), - ) - const lines = [...found.values()].map(({ skill, dir }) => { - const command = commandOf.get(skill.name) - const name = command ? `/${command}` : skill.name - const note = - command && command !== skill.name - ? ` (/${skill.name} is taken)` - : '' - return `${name} ${dir}${note}` - }) - return lines.join('\n') - }, - }), - }, + session: (ctx) => sessionPart(ctx, folders, options), } }, }) diff --git a/packages/ai-skills/tests/harness.test.ts b/packages/ai-skills/tests/harness.test.ts index 1a52b5dc82..4a6bff538f 100644 --- a/packages/ai-skills/tests/harness.test.ts +++ b/packages/ai-skills/tests/harness.test.ts @@ -12,9 +12,14 @@ import { definePlugin, } from '@tanstack/ai-harness' import { skills } from '../src/harness' +import type { FSWatcher } from 'node:fs' import type { TextOptions } from '@tanstack/ai' import type { FakeResponse } from '@tanstack/ai/testing' -import type { HarnessHost, HarnessPlugin } from '@tanstack/ai-harness' +import type { + HarnessHost, + HarnessPlugin, + HostLifecycleOptions, +} from '@tanstack/ai-harness' import type { SkillsPluginOptions } from '../src/harness' import type { SkillSource } from '../src/types' @@ -51,10 +56,14 @@ vi.mock('../src/node', async (importOriginal) => { const key = String(args[0]) if (failing.open.has(key)) throw new Error(`cannot open ${key}`) const source = original.skillDirectory(...args) - if (!failing.list.has(key)) return source return { ...source, - list: () => Promise.reject(new Error(`cannot read ${key}`)), + list: () => { + if (failing.list.has(key)) { + return Promise.reject(new Error(`cannot read ${key}`)) + } + return source.list() + }, } }, } @@ -73,6 +82,7 @@ afterEach(async () => { failing.list.clear() failing.open.clear() vi.mocked(watch).mockClear() + vi.restoreAllMocks() }) async function tempDir() { @@ -105,14 +115,27 @@ function watchersOf(path: string) { ) } +/** Is `watcher` closed? Read it after {@link nextLoop}: the close event comes then. */ +function closedFlag(watcher: FSWatcher | undefined) { + const flag = { closed: false } + watcher?.on('close', () => (flag.closed = true)) + return flag +} + +/** Wait one turn of the event loop. */ +const nextLoop = () => new Promise((resolve) => setImmediate(resolve)) + +/** Stop the warnings in the test output, and keep them. */ +const warnings = () => vi.spyOn(console, 'warn').mockImplementation(() => {}) + /** An answer of the fake model, or a function that makes it. */ type Answer = FakeResponse | (() => Promise) /** - * A session with the skills plugin and `others` after it. The model gives + * A harness with the skills plugin and `others` after it. The model gives * `answers` first and then answers "ok". It keeps each request. */ -async function open( +function harnessOf( options: SkillsPluginOptions, { others = [], @@ -129,23 +152,41 @@ async function open( return typeof answer === 'function' ? answer() : answer }), ) - const host = createHarnessHost() + const harness = defineHarness({ + name: 'test/skills', + adapter: fake, + plugins: () => [skills(options), ...others], + }) + return { harness, requests } +} + +/** A host that the test closes at its end. */ +function hostOf(lifecycle: HostLifecycleOptions = {}) { + const host = createHarnessHost({ lifecycle }) hosts.push(host) - const session = await host.open( - defineHarness({ - name: 'test/skills', - adapter: fake, - plugins: () => [skills(options), ...others], - }), - { - threadId: 't', - }, - ) + return host +} + +/** Open the session `threadId` of `harness` on `host`. */ +async function openOn( + host: HarnessHost, + harness: ReturnType['harness'], + threadId = 't', +) { + const session = await host.open(harness, { threadId }) const names = () => session.commands().map((command) => command.name) - // `/skills` syncs the commands with the folders first, so a test needs - // no watch event. + // `/skills` syncs the commands with the folders first. const list = async () => String(await session.command({ name: 'skills' })) - return { session, names, list, requests } + return { session, names, list } +} + +/** A session of a new host, with a harness from {@link harnessOf}. */ +async function open( + options: SkillsPluginOptions, + extra: Parameters[1] = {}, +) { + const { harness, requests } = harnessOf(options, extra) + return { ...(await openOn(hostOf(), harness)), requests } } /** The system prompt text of a model request. */ @@ -279,8 +320,13 @@ describe('skills()', () => { expect(await list()).toContain('No skills yet.') expect(names()).not.toContain('alpha') await addSkill(dir, 'alpha') - await list() - expect(names()).toContain('alpha') + await vi.waitFor( + async () => { + await list() + expect(names()).toContain('alpha') + }, + { timeout: 5000 }, + ) }) it('keeps the earlier folders, and the called skills, when the list is too long', async () => { @@ -392,7 +438,12 @@ describe('skills()', () => { // The first folder wins a name, also for a skill made after the start. await addSkill(first, 'alpha') await addGuide(first, 'alpha', 'first guide') - expect(await read('alpha')).toBe('first guide') + await vi.waitFor( + async () => expect(await read('alpha')).toBe('first guide'), + { + timeout: 5000, + }, + ) await expect(read('missing')).rejects.toThrow('No skill named "missing"') }) @@ -446,6 +497,7 @@ describe('skills()', () => { }) it('watches a folder without recursive where the platform has none', async () => { + const warn = warnings() const dir = await tempDir() await addSkill(dir, 'alpha') vi.mocked(watch).mockImplementationOnce(() => { @@ -464,6 +516,86 @@ describe('skills()', () => { [dir, { recursive: false, persistent: false }], ]) expect(names()).toContain('alpha') + expect(warn.mock.calls).toEqual([ + [ + `skills: cannot watch the subfolders of ${dir} on this platform. An edited SKILL.md shows at the next turn.`, + ], + ]) + // Without a recursive watcher, an edited SKILL.md shows at the next read. + await addSkill(dir, 'alpha', 'an edit') + expect(await list()).toContain(`/alpha ${dir}`) + expect( + seen.sources.at(-1) && (await seen.sources.at(-1)?.list())?.[0], + ).toMatchObject({ + description: 'Helps with an edit. Second sentence here.', + }) + }) + + it('watches each folder once for all the sessions of a host', async () => { + const dir = await tempDir() + await addSkill(dir, 'alpha') + const { harness } = harnessOf({ dirs: [dir] }) + const host = hostOf() + const first = await openOn(host, harness, 't-0') + const second = await openOn(host, harness, 't-1') + + expect(await first.list()).toBe(`/alpha ${dir}`) + expect(await second.list()).toBe(`/alpha ${dir}`) + expect(watchersOf(dir)).toHaveLength(1) + await addSkill(dir, 'gamma') + await vi.waitFor( + () => { + expect(first.names()).toContain('gamma') + expect(second.names()).toContain('gamma') + }, + { timeout: 5000 }, + ) + }) + + it('warns once, and reads the folder at each turn, when it cannot be watched', async () => { + const warn = warnings() + const dir = await tempDir() + await addSkill(dir, 'alpha') + vi.mocked(watch).mockImplementationOnce(() => { + throw Object.assign(new Error('too many open files'), { code: 'EMFILE' }) + }) + const { list } = await open({ dirs: [dir] }) + + expect(await list()).toBe(`/alpha ${dir}`) + await addSkill(dir, 'gamma') + expect(await list()).toContain(`/gamma ${dir}`) + await list() + expect(vi.mocked(watch)).toHaveBeenCalledTimes(1) + expect(warn.mock.calls).toEqual([ + [ + `skills: cannot watch ${dir} (too many open files). Its skills change at the next turn.`, + ], + ]) + }) + + it('closes the folder watcher after the last session, and at host.close()', async () => { + const dir = await tempDir() + await addSkill(dir, 'alpha') + const { harness } = harnessOf({ dirs: [dir] }) + const host = hostOf({ idleCloseMs: 0 }) + const first = await openOn(host, harness, 't-0') + const second = await openOn(host, harness, 't-1') + const shared = closedFlag(watchersOf(dir)[0]?.value) + + await first.session.close() + await nextLoop() + expect(shared.closed).toBe(false) + await second.session.close() + await nextLoop() + expect(shared.closed).toBe(true) + + // The next session watches the folder again, until the host closes. + await openOn(host, harness, 't-2') + expect(watchersOf(dir)).toHaveLength(2) + const again = closedFlag(watchersOf(dir)[1]?.value) + await Promise.all(hosts.splice(0).map((open) => open.close())) + await nextLoop() + expect(again.closed).toBe(true) }) it('skips a folder whose skills cannot be read', async () => { diff --git a/packages/ai/src/activities/chat/index.ts b/packages/ai/src/activities/chat/index.ts index b9c427f13b..437147dade 100644 --- a/packages/ai/src/activities/chat/index.ts +++ b/packages/ai/src/activities/chat/index.ts @@ -3990,6 +3990,14 @@ class TextEngine< toolCallId: toolCall.id, }) }, + onGuardToolCall: async (toolCall, tool, args) => + this.middlewareRunner.runOnGuardToolCall(this.middlewareCtx, { + toolCall, + tool, + args, + toolName: toolCall.function.name, + toolCallId: toolCall.id, + }), onAfterToolCall: async (info) => { this.logger.tools(`phase=after name=${info.toolName}`, { name: info.toolName, diff --git a/packages/ai/src/activities/chat/middleware/compose.ts b/packages/ai/src/activities/chat/middleware/compose.ts index 6940b5fbf2..f2f90d1ae1 100644 --- a/packages/ai/src/activities/chat/middleware/compose.ts +++ b/packages/ai/src/activities/chat/middleware/compose.ts @@ -11,6 +11,7 @@ import type { ChatMiddlewareContext, ErrorInfo, FinishInfo, + GuardToolCallDecision, InterruptBoundaryPhase, InterruptResolutionCollection, InterruptToolResume, @@ -539,6 +540,43 @@ export class MiddlewareRunner< return undefined } + /** + * Run onGuardToolCall through middleware in order. Every guard runs until + * one returns a decision, so an earlier `onBeforeToolCall` decision never + * skips a guard. Returns that decision, or undefined to allow the call. + */ + async runOnGuardToolCall( + ctx: ChatMiddlewareContext, + hookCtx: ToolCallHookContext, + ): Promise { + for (const mw of this.middlewares) { + if (mw.onGuardToolCall) { + const skip = shouldSkipInstrumentation(mw) + const start = Date.now() + const decision = await mw.onGuardToolCall(ctx, hookCtx) + const decided = decision !== undefined && decision !== null + if (!skip) { + this.logger.middleware( + `hook=onGuardToolCall middleware=${mw.name ?? 'unnamed'}`, + { middleware: mw.name ?? 'unnamed', hook: 'onGuardToolCall' }, + ) + aiEventClient.emit('middleware:hook:executed', { + ...instrumentCtx(ctx), + middlewareName: mw.name || 'unnamed', + hookName: 'onGuardToolCall', + iteration: ctx.iteration, + duration: Date.now() - start, + hasTransform: decided, + }) + } + if (decided) { + return decision + } + } + } + return undefined + } + /** * Run onAfterToolCall on all middleware in order. * A `replaceResult` decision changes `info.result` for the next middleware. diff --git a/packages/ai/src/activities/chat/middleware/index.ts b/packages/ai/src/activities/chat/middleware/index.ts index f7ffe81cd5..1e06958612 100644 --- a/packages/ai/src/activities/chat/middleware/index.ts +++ b/packages/ai/src/activities/chat/middleware/index.ts @@ -8,6 +8,7 @@ export type { StructuredOutputMiddlewareConfig, ToolCallHookContext, BeforeToolCallDecision, + GuardToolCallDecision, AfterToolCallDecision, AfterToolCallInfo, IterationInfo, diff --git a/packages/ai/src/activities/chat/middleware/types.ts b/packages/ai/src/activities/chat/middleware/types.ts index 101b920a6f..4669881560 100644 --- a/packages/ai/src/activities/chat/middleware/types.ts +++ b/packages/ai/src/activities/chat/middleware/types.ts @@ -432,7 +432,8 @@ export interface StructuredOutputMiddlewareConfig extends Omit< // =========================== /** - * Context provided to tool call hooks (onBeforeToolCall / onAfterToolCall). + * Context provided to tool call hooks (onBeforeToolCall / onGuardToolCall / + * onAfterToolCall). */ export interface ToolCallHookContext { /** The tool call being executed */ @@ -441,7 +442,8 @@ export interface ToolCallHookContext { tool: Tool | undefined /** * Parsed arguments for the tool call, before the `inputSchema` check: no - * schema defaults or transforms yet, and possibly invalid. + * schema defaults or transforms yet, and possibly invalid. In + * `onGuardToolCall`, these are the final checked arguments. */ args: unknown /** Name of the tool */ @@ -465,6 +467,20 @@ export type BeforeToolCallDecision = | { type: 'skip'; result: unknown } | { type: 'abort'; reason?: string } +/** + * Decision returned from onGuardToolCall. A guard cannot change the args. + * - undefined/void: allow the call. The next guard runs. + * - { type: 'skip', result }: deny the call. The tool does not run, and the + * model gets `result`, as for an `onBeforeToolCall` skip. + * - { type: 'abort', reason }: abort the entire chat run + */ +export type GuardToolCallDecision = + | void + | undefined + | null + | { type: 'skip'; result: unknown } + | { type: 'abort'; reason?: string } + /** * Outcome information provided to onAfterToolCall. */ @@ -858,6 +874,24 @@ export interface ChatMiddleware< hookCtx: ToolCallHookContext, ) => BeforeToolCallDecision | Promise + /** + * The guard tier. Called for every tool call, just before a server tool + * runs and before a client tool is sent to the client. Not called while + * the call waits for approval. + * + * It runs after all `onBeforeToolCall` hooks and after the `inputSchema` + * check, so `hookCtx.args` holds the final checked args that the tool gets. + * It also runs for calls that `onBeforeToolCall` let through unchanged. + * + * It is not first-win. The guard of each middleware runs in order, until + * one returns a decision. Use it for checks that must not be skipped, such + * as permissions. + */ + onGuardToolCall?: ( + ctx: ChatMiddlewareContext, + hookCtx: ToolCallHookContext, + ) => GuardToolCallDecision | Promise + /** * Called after a tool execution completes (success or failure). * Return `{ type: 'replaceResult', result }` to change the result that the diff --git a/packages/ai/src/activities/chat/tools/tool-calls.ts b/packages/ai/src/activities/chat/tools/tool-calls.ts index bd697e8edf..097f824d3b 100644 --- a/packages/ai/src/activities/chat/tools/tool-calls.ts +++ b/packages/ai/src/activities/chat/tools/tool-calls.ts @@ -36,6 +36,7 @@ import type { AfterToolCallDecision, AfterToolCallInfo, BeforeToolCallDecision, + GuardToolCallDecision, } from '../middleware/types' import type { McpResourceReadResult } from '../mcp/types' import type { @@ -210,6 +211,12 @@ export interface ToolExecutionMiddlewareHooks { tool: Tool | undefined, args: unknown, ) => Promise + /** Runs on the final checked args, just before the call runs. */ + onGuardToolCall?: ( + toolCall: ToolCall, + tool: Tool | undefined, + args: unknown, + ) => Promise onAfterToolCall?: (info: AfterToolCallInfo) => Promise } @@ -703,6 +710,31 @@ async function applyBeforeToolCallDecision( } const decision = await middlewareHooks.onBeforeToolCall(toolCall, tool, input) + return applyToolCallDecision( + decision, + toolCall, + tool, + input, + toolName, + middlewareHooks, + results, + ) +} + +/** + * Apply an onBeforeToolCall or onGuardToolCall decision. + * Returns the input to run with, or proceed false when the call was skipped + * (result already pushed). Throws MiddlewareAbortError on 'abort'. + */ +async function applyToolCallDecision( + decision: BeforeToolCallDecision | GuardToolCallDecision, + toolCall: ToolCall, + tool: Tool, + input: unknown, + toolName: string, + middlewareHooks: ToolExecutionMiddlewareHooks, + results: Array, +): Promise<{ proceed: true; input: unknown } | { proceed: false }> { if (!decision) { return { proceed: true, input } } @@ -734,7 +766,9 @@ async function applyBeforeToolCallDecision( return { proceed: false } } - return { proceed: true, input: decision.args } + return decision.type === 'transformArgs' + ? { proceed: true, input: decision.args } + : { proceed: true, input } } /** @@ -1285,6 +1319,27 @@ export async function* executeToolCalls( continue } + // The guard tier sees the final checked input. It always runs, also + // when an onBeforeToolCall decision came first, and cannot change it. + if (middlewareHooks?.onGuardToolCall) { + let guard: Awaited> + try { + guard = await applyToolCallDecision( + await middlewareHooks.onGuardToolCall(toolCall, tool, input), + toolCall, + tool, + input, + toolName, + middlewareHooks, + results, + ) + } catch (error) { + failures.push(error) + break + } + if (!guard.proceed) continue + } + // Create a ToolExecutionContext for this tool call with event emission const pendingEvents: Array = [] const inputResponse = resumeState?.inputResponses?.get(toolCall.id) diff --git a/packages/ai/src/index.ts b/packages/ai/src/index.ts index 6e51a8aaf2..969a71e75f 100644 --- a/packages/ai/src/index.ts +++ b/packages/ai/src/index.ts @@ -281,6 +281,7 @@ export type { StructuredOutputMiddlewareConfig, ToolCallHookContext, BeforeToolCallDecision, + GuardToolCallDecision, AfterToolCallDecision, AfterToolCallInfo, IterationInfo, diff --git a/packages/ai/src/stream-to-websocket.ts b/packages/ai/src/stream-to-websocket.ts index 202487fef2..290bac4afa 100644 --- a/packages/ai/src/stream-to-websocket.ts +++ b/packages/ai/src/stream-to-websocket.ts @@ -15,6 +15,12 @@ import type { ModelMessage, StreamChunk, UIMessage } from './types' export interface WebSocketLike { send: (data: string) => void close: (code?: number, reason?: string) => void + /** + * Bytes that `send` queued and the socket did not write yet. A server can + * read it to stop sending to a slow client. A socket without it gets no + * backpressure. + */ + readonly bufferedAmount?: number addEventListener: { (type: 'message', handler: (ev: { data: unknown }) => void): void (type: 'close' | 'error', handler: () => void): void diff --git a/packages/ai/tests/middleware.test.ts b/packages/ai/tests/middleware.test.ts index 1929a591f9..c727defc8c 100644 --- a/packages/ai/tests/middleware.test.ts +++ b/packages/ai/tests/middleware.test.ts @@ -1504,6 +1504,198 @@ describe('chat() middleware', () => { }) }) + // ========================================================================== + // onGuardToolCall + // ========================================================================== + describe('onGuardToolCall', () => { + function oneCallAdapter(args: string) { + return createMockAdapter({ + iterations: [ + [ + ev.runStarted(), + ev.toolStart('tc-1', 'myTool'), + ev.toolArgs('tc-1', args), + ev.toolEnd('tc-1', 'myTool', { input: JSON.parse(args) }), + ev.runFinished('tool_calls'), + ], + [ev.runStarted(), ev.textContent('ok'), ev.runFinished('stop')], + ], + }) + } + + it('runs after a transformArgs decision and sees the final checked args', async () => { + const executeFn = vi.fn(() => ({ result: 'done' })) + const tool = { + name: 'myTool', + description: 'Test tool', + inputSchema: z.object({ x: z.number(), unit: z.string().default('c') }), + execute: executeFn, + } + const { adapter } = oneCallAdapter('{"x":1}') + const guard = vi.fn() + + const stream = chat({ + adapter, + messages: [{ role: 'user', content: 'Hi' }], + tools: [tool], + middleware: [ + { + name: 'transformer', + onBeforeToolCall: () => ({ + type: 'transformArgs' as const, + args: { x: 42 }, + }), + }, + { name: 'guard', onGuardToolCall: guard }, + ], + }) + await collectChunks(stream as AsyncIterable) + + expect(guard).toHaveBeenCalledOnce() + expect(guard.mock.calls[0]![1]).toMatchObject({ + toolName: 'myTool', + toolCallId: 'tc-1', + args: { x: 42, unit: 'c' }, + }) + expect(executeFn).toHaveBeenCalledWith( + { x: 42, unit: 'c' }, + expect.anything(), + ) + }) + + it('a deny stops execute and the model gets it like an onBeforeToolCall skip', async () => { + const denial = { error: 'blocked by policy' } + + async function toolMessage(middleware: ChatMiddleware) { + const executeFn = vi.fn(() => ({ result: 'should not run' })) + const onAfterToolCall = vi.fn() + const { adapter, calls } = oneCallAdapter('{}') + const stream = chat({ + adapter, + messages: [{ role: 'user', content: 'Hi' }], + tools: [serverTool('myTool', executeFn)], + middleware: [{ ...middleware, onAfterToolCall }], + }) + await collectChunks(stream as AsyncIterable) + expect(executeFn).not.toHaveBeenCalled() + expect(onAfterToolCall).toHaveBeenCalledOnce() + expect(onAfterToolCall.mock.calls[0]![1]).toMatchObject({ + ok: true, + result: denial, + }) + const messages = calls[1]!.messages as Array<{ role: string }> + return messages.find((m) => m.role === 'tool') + } + + const fromGuard = await toolMessage({ + name: 'guard', + onGuardToolCall: () => ({ type: 'skip' as const, result: denial }), + }) + const fromBefore = await toolMessage({ + name: 'before', + onBeforeToolCall: () => ({ type: 'skip' as const, result: denial }), + }) + + expect(fromGuard).toBeDefined() + expect(fromGuard).toEqual(fromBefore) + }) + + it('runs the guard of every middleware in order until one denies', async () => { + const executeFn = vi.fn(() => ({ result: 'done' })) + const order: Array = [] + const { adapter } = oneCallAdapter('{}') + + const stream = chat({ + adapter, + messages: [{ role: 'user', content: 'Hi' }], + tools: [serverTool('myTool', executeFn)], + middleware: [ + { name: 'a', onGuardToolCall: () => void order.push('a') }, + { name: 'b', onGuardToolCall: () => void order.push('b') }, + { + name: 'c', + onGuardToolCall: () => { + order.push('c') + return { type: 'skip' as const, result: { error: 'no' } } + }, + }, + { name: 'd', onGuardToolCall: () => void order.push('d') }, + ], + }) + await collectChunks(stream as AsyncIterable) + + expect(order).toEqual(['a', 'b', 'c']) + expect(executeFn).not.toHaveBeenCalled() + }) + + it('guards each call of a parallel batch', async () => { + const executeA = vi.fn(() => ({ a: 'ran' })) + const executeB = vi.fn(() => ({ b: 'ran' })) + const { adapter } = createMockAdapter({ + iterations: [ + [ + ev.runStarted(), + ev.toolStart('tc-1', 'toolA'), + ev.toolArgs('tc-1', '{}'), + ev.toolEnd('tc-1', 'toolA', { input: {} }), + ev.toolStart('tc-2', 'toolB'), + ev.toolArgs('tc-2', '{}'), + ev.toolEnd('tc-2', 'toolB', { input: {} }), + ev.runFinished('tool_calls'), + ], + [ev.runStarted(), ev.textContent('done'), ev.runFinished('stop')], + ], + }) + + const stream = chat({ + adapter, + messages: [{ role: 'user', content: 'Hi' }], + tools: [serverTool('toolA', executeA), serverTool('toolB', executeB)], + toolExecution: 'parallel', + middleware: [ + { + name: 'guard', + onGuardToolCall: (_ctx, hookCtx) => + hookCtx.toolName === 'toolA' + ? { type: 'skip' as const, result: { error: 'no' } } + : undefined, + }, + ], + }) + await collectChunks(stream as AsyncIterable) + + expect(executeA).not.toHaveBeenCalled() + expect(executeB).toHaveBeenCalledOnce() + }) + + it('an abort from the guard ends the run', async () => { + const executeFn = vi.fn(() => ({ result: 'should not run' })) + const onAbort = vi.fn() + const { adapter } = oneCallAdapter('{}') + + const stream = chat({ + adapter, + messages: [{ role: 'user', content: 'Hi' }], + tools: [serverTool('myTool', executeFn)], + middleware: [ + { + name: 'guard', + onGuardToolCall: () => ({ + type: 'abort' as const, + reason: 'policy violation', + }), + onAbort, + }, + ], + }) + await collectChunks(stream as AsyncIterable) + + expect(executeFn).not.toHaveBeenCalled() + expect(onAbort).toHaveBeenCalledOnce() + expect(onAbort.mock.calls[0]![1].reason).toBe('policy violation') + }) + }) + // ========================================================================== // onConfig tools transform // ========================================================================== diff --git a/packages/ai/tests/stream-to-websocket.test.ts b/packages/ai/tests/stream-to-websocket.test.ts index 8c4e6d94b4..e8b691b913 100644 --- a/packages/ai/tests/stream-to-websocket.test.ts +++ b/packages/ai/tests/stream-to-websocket.test.ts @@ -1,4 +1,4 @@ -import { describe, expect, it, vi } from 'vitest' +import { describe, expect, expectTypeOf, it, vi } from 'vitest' import { buildTurnRequest, decodeWsFrame, @@ -101,6 +101,14 @@ class FakeSocket implements WebSocketLike { } } +describe('WebSocketLike', () => { + it('has an optional bufferedAmount (FakeSocket has none and still fits)', () => { + expectTypeOf().toEqualTypeOf< + number | undefined + >() + }) +}) + describe('buildTurnRequest', () => { it('keys the synthetic request by runId and preserves headers', () => { const handshake = new Request('https://x/api/chat', { diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index bec478d622..34292dadc1 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -2645,6 +2645,9 @@ importers: specifier: ^0.11.1 version: 0.11.2 devDependencies: + '@opentelemetry/api': + specifier: ^1.9.0 + version: 1.9.1 '@tanstack/ai': specifier: workspace:* version: link:../ai diff --git a/testing/e2e/src/routeTree.gen.ts b/testing/e2e/src/routeTree.gen.ts index 1ba9cdf165..af1fb0b5f4 100644 --- a/testing/e2e/src/routeTree.gen.ts +++ b/testing/e2e/src/routeTree.gen.ts @@ -192,6 +192,7 @@ import { Route as ApiAudioStreamRouteImport } from './routes/api.audio.stream' import { Route as ApiHarnessProtocolSplatRouteImport } from './routes/api.harness-protocol.$' import { Route as ApiHarnessSecretAnswerSplatRouteImport } from './routes/api.harness-secret-answer.$' import { Route as ApiHarnessHandoffSplatRouteImport } from './routes/api.harness-handoff.$' +import { Route as ApiHarnessObservabilitySplatRouteImport } from './routes/api.harness-observability.$' import { Route as ApiHarnessWaitsSplatRouteImport } from './routes/api.harness-waits.$' import { Route as ApiImageStreamRouteImport } from './routes/api.image.stream' import { Route as ApiTranscriptionStreamRouteImport } from './routes/api.transcription.stream' @@ -1157,6 +1158,12 @@ const ApiHarnessHandoffSplatRoute = ApiHarnessHandoffSplatRouteImport.update({ path: '/api/harness-handoff/$', getParentRoute: () => rootRouteImport, } as any) +const ApiHarnessObservabilitySplatRoute = + ApiHarnessObservabilitySplatRouteImport.update({ + id: '/api/harness-observability/$', + path: '/api/harness-observability/$', + getParentRoute: () => rootRouteImport, + } as any) const ApiHarnessWaitsSplatRoute = ApiHarnessWaitsSplatRouteImport.update({ id: '/api/harness-waits/$', path: '/api/harness-waits/$', @@ -1367,6 +1374,7 @@ export interface FileRoutesByFullPath { '/api/harness-protocol/$': typeof ApiHarnessProtocolSplatRoute '/api/harness-secret-answer/$': typeof ApiHarnessSecretAnswerSplatRoute '/api/harness-handoff/$': typeof ApiHarnessHandoffSplatRoute + '/api/harness-observability/$': typeof ApiHarnessObservabilitySplatRoute '/api/harness-waits/$': typeof ApiHarnessWaitsSplatRoute '/api/image/stream': typeof ApiImageStreamRoute '/api/transcription/stream': typeof ApiTranscriptionStreamRoute @@ -1557,6 +1565,7 @@ export interface FileRoutesByTo { '/api/harness-protocol/$': typeof ApiHarnessProtocolSplatRoute '/api/harness-secret-answer/$': typeof ApiHarnessSecretAnswerSplatRoute '/api/harness-handoff/$': typeof ApiHarnessHandoffSplatRoute + '/api/harness-observability/$': typeof ApiHarnessObservabilitySplatRoute '/api/harness-waits/$': typeof ApiHarnessWaitsSplatRoute '/api/image/stream': typeof ApiImageStreamRoute '/api/transcription/stream': typeof ApiTranscriptionStreamRoute @@ -1748,6 +1757,7 @@ export interface FileRoutesById { '/api/harness-protocol/$': typeof ApiHarnessProtocolSplatRoute '/api/harness-secret-answer/$': typeof ApiHarnessSecretAnswerSplatRoute '/api/harness-handoff/$': typeof ApiHarnessHandoffSplatRoute + '/api/harness-observability/$': typeof ApiHarnessObservabilitySplatRoute '/api/harness-waits/$': typeof ApiHarnessWaitsSplatRoute '/api/image/stream': typeof ApiImageStreamRoute '/api/transcription/stream': typeof ApiTranscriptionStreamRoute @@ -1940,6 +1950,7 @@ export interface FileRouteTypes { | '/api/harness-protocol/$' | '/api/harness-secret-answer/$' | '/api/harness-handoff/$' + | '/api/harness-observability/$' | '/api/harness-waits/$' | '/api/image/stream' | '/api/transcription/stream' @@ -2130,6 +2141,7 @@ export interface FileRouteTypes { | '/api/harness-protocol/$' | '/api/harness-secret-answer/$' | '/api/harness-handoff/$' + | '/api/harness-observability/$' | '/api/harness-waits/$' | '/api/image/stream' | '/api/transcription/stream' @@ -2320,6 +2332,7 @@ export interface FileRouteTypes { | '/api/harness-protocol/$' | '/api/harness-secret-answer/$' | '/api/harness-handoff/$' + | '/api/harness-observability/$' | '/api/harness-waits/$' | '/api/image/stream' | '/api/transcription/stream' @@ -2510,6 +2523,7 @@ export interface RootRouteChildren { ApiHarnessProtocolSplatRoute: typeof ApiHarnessProtocolSplatRoute ApiHarnessSecretAnswerSplatRoute: typeof ApiHarnessSecretAnswerSplatRoute ApiHarnessHandoffSplatRoute: typeof ApiHarnessHandoffSplatRoute + ApiHarnessObservabilitySplatRoute: typeof ApiHarnessObservabilitySplatRoute ApiHarnessWaitsSplatRoute: typeof ApiHarnessWaitsSplatRoute } @@ -3796,6 +3810,13 @@ declare module '@tanstack/react-router' { preLoaderRoute: typeof ApiHarnessHandoffSplatRouteImport parentRoute: typeof rootRouteImport } + '/api/harness-observability/$': { + id: '/api/harness-observability/$' + path: '/api/harness-observability/$' + fullPath: '/api/harness-observability/$' + preLoaderRoute: typeof ApiHarnessObservabilitySplatRouteImport + parentRoute: typeof rootRouteImport + } '/api/harness-waits/$': { id: '/api/harness-waits/$' path: '/api/harness-waits/$' @@ -4081,6 +4102,7 @@ const rootRouteChildren: RootRouteChildren = { ApiHarnessProtocolSplatRoute: ApiHarnessProtocolSplatRoute, ApiHarnessSecretAnswerSplatRoute: ApiHarnessSecretAnswerSplatRoute, ApiHarnessHandoffSplatRoute: ApiHarnessHandoffSplatRoute, + ApiHarnessObservabilitySplatRoute: ApiHarnessObservabilitySplatRoute, ApiHarnessWaitsSplatRoute: ApiHarnessWaitsSplatRoute, } export const routeTree = rootRouteImport diff --git a/testing/e2e/src/routes/api.coding-tools.ts b/testing/e2e/src/routes/api.coding-tools.ts index 4e5b59db06..c81c9b24bc 100644 --- a/testing/e2e/src/routes/api.coding-tools.ts +++ b/testing/e2e/src/routes/api.coding-tools.ts @@ -103,7 +103,7 @@ async function* revertTurns( await writeFile(join(root, 'other.txt'), 'mine\n') yield* session.send('[coding-revert-2] run test').stream() const [prompt] = await session.transcript() - if (prompt?.id) await session.revert(prompt.id) + if (prompt?.id) await session.revert({ messageId: prompt.id }) } finally { await host.close() await rm(dataDir, { recursive: true, force: true }) diff --git a/testing/e2e/src/routes/api.harness-forwarding.ts b/testing/e2e/src/routes/api.harness-forwarding.ts index 917ce2252e..35d1804499 100644 --- a/testing/e2e/src/routes/api.harness-forwarding.ts +++ b/testing/e2e/src/routes/api.harness-forwarding.ts @@ -1,6 +1,7 @@ import { createFileRoute } from '@tanstack/react-router' import { createAnthropicChat } from '@tanstack/ai-anthropic' import { createHarnessHost, defineHarness } from '@tanstack/ai-harness' +import { createHarnessHandler } from '@tanstack/ai-harness/server' import { memoryLogBackend, memoryLogStore } from '@tanstack/ai-persistence' /** @@ -10,6 +11,13 @@ import { memoryLogBackend, memoryLogStore } from '@tanstack/ai-persistence' * same thread. The follow-up waits in the log for the owner (receipt * `queued`), and the owner runs it after its turn. * + * `?scenario=` picks another case on the same two hosts: + * + * - `status`: the owner runs a turn. The route returns the statuses that + * `host.events()` of the other host shows for the thread. + * - `cancel`: the model of the owner waits until the turn stops. The other + * host gets `POST control` with a `cancel` input for the turn. + * * A custom `fetch` answers with canned Anthropic streams and holds the first * answer until the follow-up is in, so no aimock fixture is needed. */ @@ -62,12 +70,26 @@ const textsOf = (messages: Array<{ role: string; content: unknown }>) => /** * A model whose answers are `texts`, in order. The first answer waits for - * `until`. `calls` counts the requests that reach it. + * `until`. `calls` counts the requests that reach it. With `hang`, each + * request waits until the turn stops, and then fails. */ -function scriptedModel(texts: Array, until: Promise) { +function scriptedModel( + texts: Array, + until: Promise, + { hang = false } = {}, +) { const counter = { calls: 0 } - const scriptedFetch: typeof fetch = async () => { + const scriptedFetch: typeof fetch = async (_input, init) => { counter.calls += 1 + if (hang) { + return new Promise((_resolve, reject) => { + const signal = init?.signal + if (!signal) return + const stop = () => reject(signal.reason ?? new Error('Stopped.')) + if (signal.aborted) stop() + signal.addEventListener('abort', stop, { once: true }) + }) + } if (counter.calls === 1) await until const text = texts[counter.calls - 1] ?? 'No more answers.' return new Response(reply(`msg_forwarding_${counter.calls}`, text), { @@ -80,59 +102,182 @@ function scriptedModel(texts: Array, until: Promise) { return { adapter, counter } } +/** Wait until `check()` is true, for at most 10 s. */ +async function waitUntil(check: () => boolean) { + const deadline = Date.now() + 10_000 + while (!check()) { + if (Date.now() > deadline) return + await new Promise((resolve) => setTimeout(resolve, 20)) + } +} + +const NAME = 'e2e/harness-forwarding' +const THREAD = 'e2e-harness-forwarding' + +/** Two hosts on one log backend, each with the session of the thread open. */ +async function twoHosts( + ownerModel: ReturnType, + otherModel: ReturnType, +) { + const backend = memoryLogBackend() + const ownerHost = createHarnessHost({ + persistence: { stores: { log: memoryLogStore({ backend }) } }, + }) + const otherHost = createHarnessHost({ + persistence: { stores: { log: memoryLogStore({ backend }) } }, + }) + const close = async () => { + await ownerHost.close() + await otherHost.close() + } + try { + const owner = await ownerHost.open( + defineHarness({ name: NAME, adapter: ownerModel.adapter }), + { threadId: THREAD }, + ) + const otherHarness = defineHarness({ + name: NAME, + adapter: otherModel.adapter, + }) + const other = await otherHost.open(otherHarness, { threadId: THREAD }) + return { ownerHost, otherHost, owner, other, otherHarness, close } + } catch (error) { + await close() + throw error + } +} + +/** The follow-up from the other host waits in the log for the owner. */ +async function followUp() { + let release!: () => void + const released = new Promise((resolve) => { + release = resolve + }) + const ownerModel = scriptedModel(['One.', 'Two.'], released) + const otherModel = scriptedModel(['Not on this host.'], released) + const { owner, other, close } = await twoHosts(ownerModel, otherModel) + try { + const turn = owner.send('[harness-forwarding] go') + await turn.receipt + const receipt = await other.send('[harness-forwarding] and then', { + mode: 'followUp', + }).receipt + release() + const first = await turn + const settled = (await other.settled({ inputId: receipt.inputId })).outcome + return { + receipt: { status: receipt.status }, + first: first.text, + followUp: settled, + ownerCalls: ownerModel.counter.calls, + otherCalls: otherModel.counter.calls, + transcript: textsOf(await owner.transcript()), + } + } finally { + await close() + } +} + +/** The statuses of the thread on the host that does not own it. */ +async function status() { + const ownerModel = scriptedModel(['One.'], Promise.resolve()) + const otherModel = scriptedModel(['Not on this host.'], Promise.resolve()) + const { otherHost, owner, close } = await twoHosts(ownerModel, otherModel) + const reading = new AbortController() + try { + const statuses: Array = [] + const read = (async () => { + for await (const event of otherHost.events({ signal: reading.signal })) { + if (event.type === 'status' && event.threadId === THREAD) + statuses.push(event.status) + } + })() + // The first status shows that the feed reads, so start the turn now. + await waitUntil(() => statuses.length > 0) + const turn = await owner.send('[harness-forwarding] status') + await waitUntil(() => statuses.length >= 3) + reading.abort() + await read + return { + statuses, + text: turn.text, + ownerCalls: ownerModel.counter.calls, + otherCalls: otherModel.counter.calls, + } + } finally { + reading.abort() + await close() + } +} + +/** A cancel that the other host gets stops the turn on the owner. */ +async function cancel() { + const ownerModel = scriptedModel([], Promise.resolve(), { hang: true }) + const otherModel = scriptedModel(['Not on this host.'], Promise.resolve()) + const { otherHost, owner, otherHarness, close } = await twoHosts( + ownerModel, + otherModel, + ) + try { + const turn = owner.send('[harness-forwarding] cancel') + const ended = Promise.resolve(turn).then( + (result) => ({ text: result.text }), + (error: unknown) => ({ + error: error instanceof Error ? error.message : String(error), + }), + ) + const { operationId } = await turn.receipt + await waitUntil(() => ownerModel.counter.calls > 0) + const handler = createHarnessHandler({ + host: otherHost, + harness: otherHarness, + openAccess: true, + }) + const response = await handler( + new Request('http://local/api/harness/control', { + method: 'POST', + headers: { 'content-type': 'application/json' }, + body: JSON.stringify({ + threadId: THREAD, + input: { op: 'cancel', operationId }, + }), + }), + ) + const receipt: { status: string } = await response.json() + // A turn that does not stop ends the wait after 10 s. + const outcome = await Promise.race([ + ended, + new Promise((resolve) => + setTimeout(() => resolve({ error: 'The turn did not stop.' }), 10_000), + ), + ]) + return { + receipt: { status: receipt.status }, + outcome, + turnStatus: turn.status(), + ownerCalls: ownerModel.counter.calls, + otherCalls: otherModel.counter.calls, + } + } finally { + await close() + } +} + +const scenarios = { followUp, status, cancel } + export const Route = createFileRoute('/api/harness-forwarding')({ server: { handlers: { - POST: async () => { - let release!: () => void - const released = new Promise((resolve) => { - release = resolve - }) - const ownerModel = scriptedModel(['One.', 'Two.'], released) - const otherModel = scriptedModel(['Not on this host.'], released) - const backend = memoryLogBackend() - const name = 'e2e/harness-forwarding' - const ownerHost = createHarnessHost({ - persistence: { stores: { log: memoryLogStore({ backend }) } }, - }) - const otherHost = createHarnessHost({ - persistence: { stores: { log: memoryLogStore({ backend }) } }, - }) + POST: async ({ request }) => { + const picked = new URL(request.url).searchParams.get('scenario') + const scenario = + picked === 'status' || picked === 'cancel' ? picked : 'followUp' try { - const threadId = 'e2e-harness-forwarding' - const owner = await ownerHost.open( - defineHarness({ name, adapter: ownerModel.adapter }), - { threadId }, - ) - const other = await otherHost.open( - defineHarness({ name, adapter: otherModel.adapter }), - { threadId }, - ) - - const turn = owner.send('[harness-forwarding] go') - await turn.receipt - const receipt = await other.send('[harness-forwarding] and then', { - mode: 'followUp', - }).receipt - release() - const first = await turn - const followUp = (await other.settled({ inputId: receipt.inputId })) - .outcome - return Response.json({ - receipt: { status: receipt.status }, - first: first.text, - followUp, - ownerCalls: ownerModel.counter.calls, - otherCalls: otherModel.counter.calls, - transcript: textsOf(await owner.transcript()), - }) + return Response.json(await scenarios[scenario]()) } catch (error) { return Response.json({ error: error instanceof Error ? error.message : String(error), }) - } finally { - await ownerHost.close() - await otherHost.close() } }, }, diff --git a/testing/e2e/src/routes/api.harness-observability.$.ts b/testing/e2e/src/routes/api.harness-observability.$.ts new file mode 100644 index 0000000000..f2bda818e3 --- /dev/null +++ b/testing/e2e/src/routes/api.harness-observability.$.ts @@ -0,0 +1,475 @@ +import { createFileRoute } from '@tanstack/react-router' +import { EventType, defineAgent } from '@tanstack/ai' +import { createAnthropicChat } from '@tanstack/ai-anthropic' +import { createHarnessHost, defineHarness, logIdOf } from '@tanstack/ai-harness' +import { createHarnessHandler } from '@tanstack/ai-harness/server' +import { + memoryLogBackend, + memoryLogStore, + memoryPersistence, +} from '@tanstack/ai-persistence' +import { trace } from '@opentelemetry/api' +import type { + AttributeValue, + Span, + SpanContext, + SpanOptions, + Tracer, +} from '@opentelemetry/api' +import type { + HarnessHost, + HarnessHostOptions, + OperationReport, + SessionEvent, +} from '@tanstack/ai-harness' + +/** + * What a client sees of a run: the session stream on the wire, the gap + * after a truncation, the run tree, the lifecycle events, and the traces. + * + * - `POST .../start` with `{ scenario }` makes the hosts of one run, and + * returns the run id and the thread id. + * - `wire`: one host on a log runs one turn. + * - `gap`: a first host runs one turn with a snapshot after each record + * and `logRecords: 'since-snapshot'`, so the log keeps only the records + * after the snapshot. Then it stops, and a second host serves the + * thread. + * - `tree`: the `lead` agent runs `helper` and `nester` in a parallel + * group. `nester` runs a `helper` of its own. + * - `lifecycle`: one host on a log starts a turn whose model waits for + * `release`. + * - `tracing`: two hosts on one log backend, each with a fake tracer and + * an `onOperation` listener. The owner starts a turn whose model waits + * for `release`. The handler is the handler of the other host. + * - `GET .../events?runId=&threadId=` and `POST .../control?runId=` go to + * the harness handler of the run. + * - `POST .../release?runId=` lets the waiting model answer. + * - `GET .../report?runId=` returns what the test checks. + * - `POST .../stop?runId=` closes the hosts. + * + * A custom `fetch` answers with canned Anthropic streams, so no aimock + * fixture is needed. + */ + +const DUMMY_KEY = 'sk-ant-e2e-test-dummy-key' +const MODEL = 'claude-sonnet-5-5' +const NAME = 'e2e/harness-observability' + +const reply = (id: string, text: string) => + [ + { + type: 'message_start', + message: { + id, + type: 'message', + role: 'assistant', + content: [], + model: MODEL, + stop_reason: null, + stop_sequence: null, + usage: { input_tokens: 5, output_tokens: 0 }, + }, + }, + { + type: 'content_block_start', + index: 0, + content_block: { type: 'text', text: '' }, + }, + { + type: 'content_block_delta', + index: 0, + delta: { type: 'text_delta', text }, + }, + { type: 'content_block_stop', index: 0 }, + { + type: 'message_delta', + delta: { stop_reason: 'end_turn', stop_sequence: null }, + usage: { output_tokens: 1 }, + }, + { type: 'message_stop' }, + ] + .map((event) => `event: ${event.type}\ndata: ${JSON.stringify(event)}\n\n`) + .join('') + +/** + * A model whose answers are `texts`, in order. The first answer waits for + * `until`. `calls` counts the requests that reach it. + */ +function scriptedModel( + texts: Array, + until: Promise = Promise.resolve(), +) { + const counter = { calls: 0 } + const scriptedFetch: typeof fetch = async () => { + counter.calls += 1 + if (counter.calls === 1) await until + const text = texts[counter.calls - 1] ?? 'No more answers.' + return new Response(reply(`msg_observability_${counter.calls}`, text), { + headers: { 'Content-Type': 'text/event-stream' }, + }) + } + const adapter = createAnthropicChat(MODEL, DUMMY_KEY, { + fetch: scriptedFetch, + }) + return { adapter, counter } +} + +/** A promise that `open()` resolves. */ +function gate() { + let open = () => {} + const opened = new Promise((resolve) => { + open = resolve + }) + return { open, opened } +} + +/** Wait until `check()` is true, for at most 10 s. */ +async function waitUntil(check: () => boolean) { + const deadline = Date.now() + 10_000 + while (!check()) { + if (Date.now() > deadline) throw new Error('Waited 10 s for nothing.') + await new Promise((resolve) => setTimeout(resolve, 20)) + } +} + +const isFinished = ({ event }: SessionEvent) => + event.type === EventType.CUSTOM && event.name === 'harness.operation.finished' + +/** The logged events of `events`, until the first operation ends. */ +async function loggedTurn(events: AsyncIterable) { + const entries: Array = [] + for await (const entry of events) { + entries.push(entry) + if (isFinished(entry)) break + } + return entries +} + +/** A span that the fake tracer started. */ +interface RecordedSpan { + name: string + options: SpanOptions + attributes: Record + /** The span of the context that `startSpan` got. */ + parent: RecordedSpan | undefined + ended: boolean + span: Span +} + +/** + * A tracer that records every span, and reads the parent from the context, + * as the fake tracer of `tracing.test.ts` does. `reports` keeps what + * `onOperation` gets. + */ +function fakeObservability() { + const spans: Array = [] + const reports: Array = [] + const tracer: Tracer = { + startSpan: (name, options = {}, context) => { + const parentSpan = context && trace.getSpan(context) + const spanContext: SpanContext = { + traceId: 'a'.repeat(32), + spanId: (spans.length + 1).toString(16).padStart(16, '0'), + traceFlags: 1, + } + const span: Span = { + spanContext: () => spanContext, + setAttribute: (key, value) => { + recorded.attributes[key] = value + return span + }, + setAttributes: () => span, + addEvent: () => span, + addLink: () => span, + addLinks: () => span, + setStatus: () => span, + updateName: () => span, + end: () => { + recorded.ended = true + }, + isRecording: () => true, + recordException: () => {}, + } + const recorded: RecordedSpan = { + name, + options, + attributes: { ...options.attributes }, + parent: spans.find((item) => item.span === parentSpan), + ended: false, + span, + } + spans.push(recorded) + return span + }, + startActiveSpan: () => { + throw new Error('The harness does not call startActiveSpan.') + }, + } + return { + observability: { + tracer, + onOperation: (report: OperationReport) => void reports.push(report), + }, + /** The spans and the reports, as JSON. */ + read: () => ({ + spans: spans.map((item) => ({ + name: item.name, + attributes: item.attributes, + links: item.options.links ?? null, + hasParent: item.parent !== undefined, + ended: item.ended, + })), + reports, + }), + } +} + +/** What a scenario makes. */ +interface Setup { + handler: (request: Request) => Promise + report: () => Promise + hosts: Array +} + +async function wire(threadId: string): Promise { + const { adapter } = scriptedModel(['Hello on the wire.']) + const harness = defineHarness({ name: NAME, adapter }) + const host = createHarnessHost({ + persistence: { stores: { log: memoryLogStore() } }, + }) + const session = await host.open(harness, { threadId }) + await session.send('[harness-observability] wire') + return { + handler: createHarnessHandler({ host, harness, openAccess: true }), + report: async () => ({ logged: await loggedTurn(session.events()) }), + hosts: [host], + } +} + +async function gap(threadId: string): Promise { + const harness = defineHarness({ + name: NAME, + adapter: scriptedModel(['Hello before the cut.']).adapter, + }) + const persistence = memoryPersistence() + const options: HarnessHostOptions = { + persistence, + log: { snapshot: { afterBytes: 1 } }, + lifecycle: { retention: { logRecords: 'since-snapshot' } }, + } + const first = createHarnessHost(options) + const session = await first.open(harness, { threadId }) + await session.send('[harness-observability] gap') + await session.close() + await first.close() + const snapshot = await persistence.stores.log.getSnapshot({ + logId: logIdOf({ logId: threadId }), + }) + if (!snapshot) throw new Error('The first host wrote no snapshot.') + const host = createHarnessHost(options) + return { + handler: createHarnessHandler({ host, harness, openAccess: true }), + report: async () => ({ head: String(snapshot.seq) }), + hosts: [host], + } +} + +const helper = defineAgent({ + name: 'helper', + description: 'Helps', + run: async () => 'helped', +}) + +/** Starts a helper of its own, so the tree has three levels. */ +const nester = defineAgent({ + name: 'nester', + description: 'Starts a helper', + run: async (ctx) => await ctx.harness.agents.run('helper'), +}) + +/** Runs a helper and a nester in a parallel group. */ +const lead = defineAgent({ + name: 'lead', + description: 'Runs a group', + run: async (ctx) => { + const results = await ctx.harness.agents.group({}, (g) => [ + g.start('helper'), + g.start('nester'), + ]) + return results.join(' ') + }, +}) + +async function tree(threadId: string): Promise { + const harness = defineHarness({ + name: NAME, + // The agents make no model call. + adapter: scriptedModel([]).adapter, + agents: [lead, nester, helper], + // The helper of the nester is two levels below the lead. + subagents: { agents: [], limits: { maxDepth: 3 } }, + }) + const host = createHarnessHost({ persistence: memoryPersistence() }) + const session = await host.open(harness, { threadId }) + const run = session.agents.lead.start() + const result = await run + return { + handler: createHarnessHandler({ host, harness, openAccess: true }), + report: async () => ({ + result, + leadRunId: run.id, + agentRuns: session.agentRuns().map((item) => ({ + runId: item.runId, + parentRunId: item.parentRunId ?? null, + })), + }), + hosts: [host], + } +} + +async function lifecycle( + threadId: string, + released: Promise, +): Promise { + const { adapter } = scriptedModel(['One.', 'Two.'], released) + const harness = defineHarness({ name: NAME, adapter }) + const host = createHarnessHost({ + persistence: { stores: { log: memoryLogStore() } }, + }) + const session = await host.open(harness, { threadId }) + const turn = session.send('[harness-observability] one', { inputId: 'in-1' }) + Promise.resolve(turn).catch(() => {}) + await waitUntil(() => session.snapshot().status === 'running') + return { + handler: createHarnessHandler({ host, harness, openAccess: true }), + report: async () => ({ status: session.snapshot().status }), + hosts: [host], + } +} + +async function tracing( + threadId: string, + released: Promise, +): Promise { + const backend = memoryLogBackend() + const ownerTracing = fakeObservability() + const otherTracing = fakeObservability() + const ownerModel = scriptedModel(['One.', 'Two.'], released) + const otherModel = scriptedModel(['Not on this host.']) + const ownerHost = createHarnessHost({ + persistence: { stores: { log: memoryLogStore({ backend }) } }, + observability: ownerTracing.observability, + }) + const otherHost = createHarnessHost({ + persistence: { stores: { log: memoryLogStore({ backend }) } }, + observability: otherTracing.observability, + }) + const owner = await ownerHost.open( + defineHarness({ name: NAME, adapter: ownerModel.adapter }), + { threadId }, + ) + const otherHarness = defineHarness({ + name: NAME, + adapter: otherModel.adapter, + }) + await otherHost.open(otherHarness, { threadId }) + const turn = owner.send('[harness-observability] go') + Promise.resolve(turn).catch(() => {}) + await turn.receipt + return { + handler: createHarnessHandler({ + host: otherHost, + harness: otherHarness, + openAccess: true, + }), + report: async () => ({ + owner: ownerTracing.read(), + other: otherTracing.read(), + ownerCalls: ownerModel.counter.calls, + otherCalls: otherModel.counter.calls, + }), + hosts: [ownerHost, otherHost], + } +} + +const scenarios = { + wire: (threadId: string) => wire(threadId), + gap: (threadId: string) => gap(threadId), + tree: (threadId: string) => tree(threadId), + lifecycle, + tracing, +} + +type Scenario = keyof typeof scenarios + +interface Run extends Setup { + release: () => void + stop: () => Promise +} + +// ponytail: the runs that did not stop yet, kept in memory by run id. +const runs = new Map() +/** A run that nobody stops ends after this time. */ +const RUN_TTL_MS = 60_000 + +const errorOf = (error: unknown) => + error instanceof Error ? error.message : String(error) + +async function start(request: Request) { + const { scenario }: { scenario: Scenario } = await request.json() + if (!(scenario in scenarios)) { + return Response.json({ error: `No scenario ${scenario}.` }) + } + const runId = crypto.randomUUID() + const threadId = `e2e-harness-observability-${scenario}-${runId}` + const release = gate() + try { + const setup = await scenarios[scenario](threadId, release.opened) + const stop = async () => { + if (!runs.delete(runId)) return + release.open() + for (const host of setup.hosts) await host.close() + } + runs.set(runId, { ...setup, release: release.open, stop }) + setTimeout(() => void stop(), RUN_TTL_MS) + return Response.json({ runId, threadId }) + } catch (error) { + return Response.json({ error: errorOf(error) }) + } +} + +export const Route = createFileRoute('/api/harness-observability/$')({ + server: { + handlers: { + POST: async ({ request }) => { + const url = new URL(request.url) + const route = url.pathname.split('/').at(-1) + if (route === 'start') return start(request) + const run = runs.get(url.searchParams.get('runId') ?? '') + if (!run) return Response.json({ error: 'No such run.' }) + if (route === 'control') return run.handler(request) + if (route === 'release') { + run.release() + return Response.json({ released: true }) + } + if (route === 'stop') { + await run.stop() + return Response.json({ stopped: true }) + } + return Response.json({ error: `No route ${route}.` }, { status: 404 }) + }, + GET: async ({ request }) => { + const url = new URL(request.url) + const route = url.pathname.split('/').at(-1) + const run = runs.get(url.searchParams.get('runId') ?? '') + if (!run) return Response.json({ error: 'No such run.' }) + if (route !== 'report') return run.handler(request) + try { + return Response.json(await run.report()) + } catch (error) { + return Response.json({ error: errorOf(error) }) + } + }, + }, + }, +}) diff --git a/testing/e2e/src/routes/api.harness-protocol.$.ts b/testing/e2e/src/routes/api.harness-protocol.$.ts index c3eaa48988..a3c60fb740 100644 --- a/testing/e2e/src/routes/api.harness-protocol.$.ts +++ b/testing/e2e/src/routes/api.harness-protocol.$.ts @@ -341,7 +341,7 @@ function handlerFor(request: Request) { harness, // Each request makes a new handler, so a fixed secret keeps a signed URL // working on the next request. - media: { secret: 'e2e-media-secret' }, + media: { secret: 'e2e-media-secret-of-32-bytes-or-more' }, }) } diff --git a/testing/e2e/src/routes/api.middleware-test.ts b/testing/e2e/src/routes/api.middleware-test.ts index f5d7cb672a..cccdabc1d5 100644 --- a/testing/e2e/src/routes/api.middleware-test.ts +++ b/testing/e2e/src/routes/api.middleware-test.ts @@ -187,6 +187,33 @@ const toolBadArgsMiddleware: ChatMiddleware = { }, } +// `tool-guard` mode: the first middleware rewrites the city, so +// onBeforeToolCall is done. The guard still runs, sees the rewritten city, +// and denies the call. +const toolRewriteMiddleware: ChatMiddleware = { + name: 'tool-rewrite', + onBeforeToolCall(_ctx, hookCtx) { + if (hookCtx.toolName !== 'get_weather') return undefined + return { type: 'transformArgs' as const, args: { city: 'Blocked City' } } + }, +} + +const toolGuardMiddleware: ChatMiddleware = { + name: 'tool-guard', + onGuardToolCall(_ctx, hookCtx) { + const args = hookCtx.args + const city = + typeof args === 'object' && args !== null && 'city' in args + ? args.city + : undefined + if (hookCtx.toolName !== 'get_weather' || city !== 'Blocked City') return + return { + type: 'skip' as const, + result: JSON.stringify({ guarded: true, city }), + } + }, +} + /** * Capability provide/consume flow (`capability` mode). * @@ -591,6 +618,8 @@ export const Route = createFileRoute('/api/middleware-test')({ middleware.push(toolReplaceMiddleware) if (middlewareMode === 'tool-bad-args') middleware.push(toolBadArgsMiddleware) + if (middlewareMode === 'tool-guard') + middleware.push(toolRewriteMiddleware, toolGuardMiddleware) if (middlewareMode === 'capability') { // Order matters: the provider's setup() must run before the // consumer reads the capability. Array order is preserved. diff --git a/testing/e2e/src/routes/middleware-test.tsx b/testing/e2e/src/routes/middleware-test.tsx index fc326cc99b..19da5b7511 100644 --- a/testing/e2e/src/routes/middleware-test.tsx +++ b/testing/e2e/src/routes/middleware-test.tsx @@ -20,6 +20,7 @@ const MIDDLEWARE_MODES = [ { id: 'tool-skip', label: 'Tool Skip (skip with custom result)' }, { id: 'tool-replace', label: 'Tool Replace (replace result after call)' }, { id: 'tool-bad-args', label: 'Tool Bad Args (invalid args rewrite)' }, + { id: 'tool-guard', label: 'Tool Guard (deny on final args)' }, { id: 'capability', label: 'Capability (provide/consume prefix)' }, { id: 'phase-recorder', label: 'Phase Recorder (capture phase + chunks)' }, { id: 'otel', label: 'OpenTelemetry (capture spans/metrics)' }, diff --git a/testing/e2e/tests/harness-access.spec.ts b/testing/e2e/tests/harness-access.spec.ts index 80fead9335..47faaa3bcf 100644 --- a/testing/e2e/tests/harness-access.spec.ts +++ b/testing/e2e/tests/harness-access.spec.ts @@ -30,8 +30,10 @@ const refusalOf = async (response: APIResponse) => ({ }) /** - * Open the session stream of a thread, read its first frame, and close the - * stream. A refused request answers with JSON, not with a stream. + * Open the session stream of a thread, read its first frame after the + * `harness.hello`, and close the stream. `: ping` comments have no `data:` + * line, so they are skipped. A refused request answers with JSON, not with + * a stream. */ async function firstEvent(url: string, headers: Record) { const response = await fetch(url, { headers }) @@ -40,17 +42,24 @@ async function firstEvent(url: string, headers: Record) { } const reader = response.body.pipeThrough(new TextDecoderStream()).getReader() let buffer = '' - while (!buffer.includes('\n\n')) { + let body: { type?: string } | null = null + while (body === null) { const read = await reader.read() if (read.done) break buffer += read.value + const blocks = buffer.split('\n\n') + buffer = blocks.pop() ?? '' + for (const block of blocks) { + const data = block.split('\n').find((line) => line.startsWith('data: ')) + if (!data) continue + const frame: { type?: string } = JSON.parse(data.slice('data: '.length)) + if (frame.type === 'harness.hello') continue + body = frame + break + } } await reader.cancel() - const data = buffer.split('\n').find((line) => line.startsWith('data: ')) - return { - status: response.status, - body: data ? JSON.parse(data.slice('data: '.length)) : null, - } + return { status: response.status, body } } /** One user of the `membersOnly()` host, by the token of the route. */ diff --git a/testing/e2e/tests/harness-forwarding.spec.ts b/testing/e2e/tests/harness-forwarding.spec.ts index eb6f54f093..ed0b976801 100644 --- a/testing/e2e/tests/harness-forwarding.spec.ts +++ b/testing/e2e/tests/harness-forwarding.spec.ts @@ -33,3 +33,49 @@ test.describe('harness: an input on a host that does not own the thread', () => ]) }) }) + +test.describe('harness: the host that does not own the thread', () => { + test('the host that does not own the thread shows running, then idle', async ({ + request, + }) => { + const res = await request.post('/api/harness-forwarding?scenario=status') + expect(res.ok()).toBe(true) + const result: { + error?: string + statuses: Array + text: string + ownerCalls: number + otherCalls: number + } = await res.json() + expect(result.error ?? null).toBeNull() + + expect(result.text).toBe('One.') + // The other host runs nothing. Its status comes from the log tail. + expect(result.statuses).toEqual(['idle', 'running', 'idle']) + expect(result.ownerCalls).toBe(1) + expect(result.otherCalls).toBe(0) + }) + + test('a cancel sent to the other host stops the turn', async ({ + request, + }) => { + const res = await request.post('/api/harness-forwarding?scenario=cancel') + expect(res.ok()).toBe(true) + const result: { + error?: string + receipt: { status: string } + outcome: { text?: string; error?: string } + turnStatus: string + ownerCalls: number + otherCalls: number + } = await res.json() + expect(result.error ?? null).toBeNull() + + // The other host stores the cancel in the log, and the owner applies it. + expect(result.receipt.status).toBe('queued') + expect(result.outcome).toEqual({ error: 'Cancelled.' }) + expect(result.turnStatus).toBe('cancelled') + expect(result.ownerCalls).toBe(1) + expect(result.otherCalls).toBe(0) + }) +}) diff --git a/testing/e2e/tests/harness-observability.spec.ts b/testing/e2e/tests/harness-observability.spec.ts new file mode 100644 index 0000000000..d2ad0a0342 --- /dev/null +++ b/testing/e2e/tests/harness-observability.spec.ts @@ -0,0 +1,382 @@ +import { encodeWsFrame } from '@tanstack/ai' +import { test, expect } from './fixtures' +import type { APIRequestContext } from '@playwright/test' +import type { AttributeValue } from '@opentelemetry/api' +import type { OperationReport, SessionEvent } from '@tanstack/ai-harness' + +type Scenario = 'wire' | 'gap' | 'tree' | 'lifecycle' | 'tracing' + +/** One frame of an SSE stream, with the `id` of its SSE event. */ +interface Read { + id: string | null + frame: { + type: string + v?: number + cursor?: string + threadId?: string + operationId?: string + event?: { + type: string + name?: string + value?: unknown + parentRunId?: string + } + } +} + +interface RecordedSpan { + name: string + attributes: Record + links: Array<{ context: unknown }> | null + hasParent: boolean + ended: boolean +} + +interface Traces { + spans: Array + reports: Array +} + +const TRACE_ID = '4bf92f3577b34da6a3ce929d0e0e4736' +const CALLER = `00-${TRACE_ID}-00f067aa0ba902b7-01` +/** The span context that `CALLER` names. */ +const CALLER_CONTEXT = { + traceId: TRACE_ID, + spanId: '00f067aa0ba902b7', + traceFlags: 1, + isRemote: true, +} + +const ROUTE = '/api/harness-observability' + +/** Start a run of `scenario`. */ +async function start(request: APIRequestContext, scenario: Scenario) { + const res = await request.post(`${ROUTE}/start`, { data: { scenario } }) + expect(res.ok()).toBe(true) + const started: { error?: string; runId: string; threadId: string } = + await res.json() + expect(started.error ?? null).toBeNull() + return started +} + +async function report(request: APIRequestContext, runId: string) { + const res = await request.get(`${ROUTE}/report?runId=${runId}`) + const body: T & { error?: string } = await res.json() + expect(body.error ?? null).toBeNull() + return body +} + +const stop = (request: APIRequestContext, runId: string) => + request.post(`${ROUTE}/stop?runId=${runId}`) + +/** + * Read `GET events` of the run until `isDone` is true for the frames so + * far. `: ping` comments have no `data:` line, so they are skipped. + */ +async function readEvents( + url: string, + isDone: (frames: Array) => boolean, + headers: Record = {}, +) { + const response = await fetch(url, { headers }) + if (!response.body) throw new Error('The session stream has no body.') + const reader = response.body.pipeThrough(new TextDecoderStream()).getReader() + const frames: Array = [] + let buffer = '' + while (!isDone(frames)) { + const read = await reader.read() + if (read.done) break + buffer += read.value + const blocks = buffer.split('\n\n') + buffer = blocks.pop() ?? '' + for (const block of blocks) { + const lines = block.split('\n') + const data = lines.find((line) => line.startsWith('data: ')) + if (!data) continue + const id = lines.find((line) => line.startsWith('id: ')) + frames.push({ + id: id ? id.slice('id: '.length) : null, + frame: JSON.parse(data.slice('data: '.length)), + }) + } + } + await reader.cancel() + return { response, frames } +} + +const isFinished = ({ frame }: Read) => + frame.event?.type === 'CUSTOM' && + frame.event.name === 'harness.operation.finished' + +/** The values of the `name` events in `frames`, in order. */ +const valuesOf = (frames: Array, name: string) => + frames.flatMap(({ frame }) => + frame.event?.type === 'CUSTOM' && frame.event.name === name + ? [frame.event.value] + : [], + ) + +test.describe('harness observability', () => { + test('every frame of GET events is a wire chunk and starts with hello', async ({ + request, + baseURL, + }) => { + const { runId, threadId } = await start(request, 'wire') + try { + const { response, frames } = await readEvents( + `${baseURL}${ROUTE}/events?runId=${runId}&threadId=${threadId}&v=2&from=0`, + (seen) => seen.some(isFinished), + { 'X-Request-Id': 'obs-wire-1' }, + ) + expect(response.headers.get('X-Request-Id')).toBe('obs-wire-1') + expect(frames[0]).toEqual({ + id: null, + frame: { type: 'harness.hello', v: 2, threadId }, + }) + + const events = frames.slice(1) + const { logged } = await report<{ logged: Array }>( + request, + runId, + ) + expect(events.map(({ frame }) => frame.type)).toEqual( + logged.map(() => 'harness.event'), + ) + // The SSE id and the frame cursor are the cursor of the logged event. + expect(events.map(({ id, frame }) => [id, frame.cursor])).toEqual( + logged.map((entry) => [entry.cursor, entry.cursor]), + ) + events.forEach(({ frame }, index) => { + const entry = logged[index]! + expect(frame.threadId).toBe(threadId) + expect(frame.operationId).toBe(entry.operationId) + expect(frame.event).toEqual( + JSON.parse(encodeWsFrame(entry.event, undefined)), + ) + }) + // At least one logged chunk has keys that the wire drops, so an + // internal chunk on the wire fails the check above. + const isLossy = logged.some( + (entry) => + encodeWsFrame(entry.event, undefined) !== JSON.stringify(entry.event), + ) + expect(isLossy).toBe(true) + } finally { + await stop(request, runId) + } + }) + + test('a cursor from before a truncation gets harness.gap', async ({ + request, + baseURL, + }) => { + const { runId, threadId } = await start(request, 'gap') + try { + const { head } = await report<{ head: string }>(request, runId) + const gap = { + id: head, + frame: { + type: 'harness.event', + cursor: head, + threadId, + operationId: 'session', + event: { + type: 'CUSTOM', + name: 'harness.gap', + value: { from: '0', resumeAt: head, reason: 'truncated' }, + timestamp: expect.any(Number), + }, + }, + } + const url = `${baseURL}${ROUTE}/events?runId=${runId}&threadId=${threadId}` + const fromQuery = await readEvents( + `${url}&from=0`, + (seen) => seen.length >= 2, + ) + expect(fromQuery.frames).toEqual([ + { id: null, frame: { type: 'harness.hello', v: 2, threadId } }, + gap, + ]) + + // A client that resumes with `Last-Event-ID` gets the same gap. + const resumed = await readEvents(url, (seen) => seen.length >= 2, { + 'Last-Event-ID': '0', + }) + expect(resumed.frames[1]).toEqual(gap) + } finally { + await stop(request, runId) + } + }) + + test('a parallel group of agents gives the whole tree in one stream', async ({ + request, + baseURL, + }) => { + const { runId, threadId } = await start(request, 'tree') + try { + const run = await report<{ + result: string + leadRunId: string + agentRuns: Array<{ runId: string; parentRunId: string | null }> + }>(request, runId) + expect(run.result).toBe('helped helped') + // The lead, its helper and nester, and the helper of the nester. + const expected = Object.fromEntries( + run.agentRuns.map((item) => [item.runId, item.parentRunId]), + ) + expect(Object.keys(expected)).toHaveLength(4) + expect(expected[run.leadRunId]).toBeNull() + const children = run.agentRuns.filter( + (item) => item.parentRunId === run.leadRunId, + ) + expect(children).toHaveLength(2) + const grandchildren = run.agentRuns.filter((item) => + children.some((child) => child.runId === item.parentRunId), + ) + expect(grandchildren).toHaveLength(1) + + const { frames } = await readEvents( + `${baseURL}${ROUTE}/events?runId=${runId}&threadId=${threadId}&from=0`, + (seen) => + seen.some( + (read) => + read.frame.operationId === run.leadRunId && isFinished(read), + ), + ) + + // The tree from the frames alone: the `parentRunId` of the AG-UI + // RUN_STARTED of each operation. + const tree = new Map() + for (const { frame } of frames.slice(1)) { + expect(frame.threadId).toBe(threadId) + if (frame.event?.type !== 'RUN_STARTED') continue + tree.set(frame.operationId ?? '', frame.event.parentRunId ?? null) + } + // The session has operations of its own, for example its inputs. + const agentTree = Object.fromEntries( + [...tree].filter(([operationId]) => operationId in expected), + ) + expect(agentTree).toEqual(expected) + } finally { + await stop(request, runId) + } + }) + + test('a new thread gets harness.ownership.changed, and a send during a turn gets harness.input.queued', async ({ + request, + baseURL, + }) => { + const { runId, threadId } = await start(request, 'lifecycle') + try { + // The first turn runs and waits for the model. + const receipt = await request.post(`${ROUTE}/control?runId=${runId}`, { + data: { + threadId, + input: { + op: 'prompt', + message: '[harness-observability] two', + inputId: 'in-2', + }, + }, + }) + expect(await receipt.json()).toMatchObject({ + status: 'queued', + inputId: 'in-2', + }) + await request.post(`${ROUTE}/release?runId=${runId}`) + + const { frames } = await readEvents( + `${baseURL}${ROUTE}/events?runId=${runId}&threadId=${threadId}&from=0`, + (seen) => seen.filter(isFinished).length >= 2, + ) + expect(valuesOf(frames, 'harness.ownership.changed')).toEqual([ + { epoch: 1, reason: 'acquired' }, + ]) + expect(valuesOf(frames, 'harness.input.queued')).toEqual([ + { inputId: 'in-2', delivery: 'queue' }, + ]) + // The host owns the log before the first turn starts. + const names = frames.map(({ frame }) => frame.event?.name) + expect(names.indexOf('harness.ownership.changed')).toBeLessThan( + names.indexOf('harness.operation.started'), + ) + } finally { + await stop(request, runId) + } + }) + + test('a traceparent on POST control links the span of the run on the host that runs it', async ({ + request, + }) => { + const { runId, threadId } = await start(request, 'tracing') + try { + // The other host stores the input in the log. The owner runs it after + // its turn. + const receipt = await request.post(`${ROUTE}/control?runId=${runId}`, { + headers: { traceparent: CALLER }, + data: { + threadId, + input: { op: 'followUp', message: '[harness-observability] then' }, + }, + }) + expect(await receipt.json()).toMatchObject({ status: 'queued' }) + await request.post(`${ROUTE}/release?runId=${runId}`) + + type Traced = { + owner: Traces + other: Traces + ownerCalls: number + otherCalls: number + } + let traced: Traced | undefined + const operations = (traces: Traces) => + traces.spans.filter((span) => span.name === 'harness.operation') + await expect + .poll( + async () => { + traced = await report(request, runId) + const ended = operations(traced.owner).filter((span) => span.ended) + return [ended.length, traced.owner.reports.length] + }, + { timeout: 20_000 }, + ) + .toEqual([2, 4]) + const { owner, other, ownerCalls, otherCalls } = traced! + + const spans = operations(owner) + for (const span of spans) { + expect(span.attributes).toMatchObject({ + 'harness.thread_id': threadId, + 'harness.operation.kind': 'chat', + 'harness.operation.status': 'completed', + }) + // A link, not a parent: the caller is on another machine. + expect(span.hasParent).toBe(false) + } + const linked = spans.filter((span) => span.links !== null) + expect(linked).toHaveLength(1) + expect(linked[0]?.links).toEqual([{ context: CALLER_CONTEXT }]) + + // `onOperation` reports each run once when it starts and once when it + // settles. + for (const span of spans) { + const runOf = owner.reports.filter( + (item) => item.operation.id === span.attributes['harness.run_id'], + ) + expect(runOf.map((item) => item.phase)).toEqual(['started', 'settled']) + expect(runOf[1]?.operation).toMatchObject({ + threadId, + kind: 'chat', + status: 'completed', + }) + } + // The host that only stored the input runs nothing, so it reports + // nothing. + expect(other).toEqual({ spans: [], reports: [] }) + expect(ownerCalls).toBe(2) + expect(otherCalls).toBe(0) + } finally { + await stop(request, runId) + } + }) +}) diff --git a/testing/e2e/tests/harness-protocol.spec.ts b/testing/e2e/tests/harness-protocol.spec.ts index 6d29f3fe4c..30923d806f 100644 --- a/testing/e2e/tests/harness-protocol.spec.ts +++ b/testing/e2e/tests/harness-protocol.spec.ts @@ -584,4 +584,34 @@ test.describe('harness protocol', () => { await reader.cancel() expect(statuses).toEqual(['idle', 'running', 'idle']) }) + + test('an unsupported v gets an error frame with a requestId', async ({ + request, + testId, + aimockPort, + }) => { + const auth = headers(testId, aimockPort) + const threadId = `version-${testId}` + await openHarnessThread(request, auth, threadId) + + // The host sends one error frame, then closes the stream. + const response = await request.get( + `/api/harness-protocol/events?threadId=${threadId}&v=9`, + { headers: { ...auth, 'x-request-id': 'e2e-version-9' } }, + ) + expect(response.headers()['x-request-id']).toBe('e2e-version-9') + const frames = (await response.text()) + .split('\n') + .filter((line) => line.startsWith('data: ')) + .map((line) => JSON.parse(line.slice('data: '.length))) + expect(frames).toEqual([ + { + type: 'harness.error', + message: + 'Protocol version 9 is not supported. This host speaks version 2.', + code: 'unsupported_version', + requestId: 'e2e-version-9', + }, + ]) + }) }) diff --git a/testing/e2e/tests/middleware.spec.ts b/testing/e2e/tests/middleware.spec.ts index 6e1d5b50c3..e0606981ff 100644 --- a/testing/e2e/tests/middleware.spec.ts +++ b/testing/e2e/tests/middleware.spec.ts @@ -167,6 +167,42 @@ test.describe('Middleware Lifecycle', () => { expect(toolResults[0].content).not.toContain('temperature') }) + test('onGuardToolCall denies on the final args after a transformArgs', async ({ + page, + testId, + aimockPort, + }) => { + const params = new URLSearchParams() + if (testId) params.set('testId', testId) + if (aimockPort) params.set('aimockPort', String(aimockPort)) + const qs = params.toString() + await page.goto(`/middleware-test${qs ? '?' + qs : ''}`) + await page.waitForTimeout(2000) + await page.locator('#mw-scenario-select').selectOption('with-tool') + await page.locator('#mw-mode-select').selectOption('tool-guard') + await page.locator('#mw-run-button').click() + + await page.waitForFunction( + () => + document + .querySelector('#mw-metadata') + ?.getAttribute('data-test-complete') === 'true', + { timeout: 10000 }, + ) + + const messagesJson = await page.locator('#mw-messages-json').textContent() + const messages = JSON.parse(messagesJson || '[]') + const toolResults = messages.flatMap((m: any) => + m.parts.filter((p: any) => p.type === 'tool-result'), + ) + // The rewrite ran first. The guard saw the new city and denied the call, + // so get_weather never ran. + expect(toolResults).toHaveLength(1) + expect(toolResults[0].content).toContain('guarded') + expect(toolResults[0].content).toContain('Blocked City') + expect(toolResults[0].content).not.toContain('temperature') + }) + test('the next model call replays the exact tool argument bytes', async ({ page, request, From ea94e8dcc15999e0f5e56da57e9a6d33692effaf Mon Sep 17 00:00:00 2001 From: Alem Tuzlak Date: Mon, 19 Oct 2026 13:54:58 +0200 Subject: [PATCH 10/17] fix(ai-harness): coding tools review fixes for many threads - The changed-since-read check keeps what each chat run read, so a lead and its tool-mode subagents do not lose each other's edits. - File locks are keyed by the sandbox (or host) and the real path, not by the root string. - A foreground command moves to the background only under maxJobs, and the log gets its process id (exec onStart), so a later host can stop it. - A job that the backend cannot check (killJob throws, or no killJob) stays in the log, and the model hears that it may still run. hostBackend.killJob throws when ps gives no start time. - workspaceRoot(base, ...parts) encodes each part of a workspace folder, so a thread id from a client cannot leave base. - host.unpin({ harness, threadId }) unpins a thread from a host that does not come back. --- .../src/first-party/coding/backend.ts | 66 +++++++++++++--- .../ai-harness/src/first-party/coding/bash.ts | 32 ++++++-- .../ai-harness/src/first-party/coding/edit.ts | 35 ++++++--- .../src/first-party/coding/index.ts | 2 +- .../src/first-party/coding/ring-buffer.ts | 4 +- .../src/first-party/coding/workspace.ts | 76 +++++++++++++----- packages/ai-harness/src/host.ts | 46 +++++++++++ packages/ai-harness/src/ownership.ts | 35 +++++++++ .../ai-harness/tests/coding-backend.test.ts | 16 ++++ packages/ai-harness/tests/coding-bash.test.ts | 1 + packages/ai-harness/tests/coding-jobs.test.ts | 48 +++++++++++- .../ai-harness/tests/coding-locks.test.ts | 77 +++++++++++++++++-- .../tests/workspace-binding.test.ts | 67 +++++++++++++++- 13 files changed, 450 insertions(+), 55 deletions(-) diff --git a/packages/ai-harness/src/first-party/coding/backend.ts b/packages/ai-harness/src/first-party/coding/backend.ts index 7b8aba4288..2659c42e13 100644 --- a/packages/ai-harness/src/first-party/coding/backend.ts +++ b/packages/ai-harness/src/first-party/coding/backend.ts @@ -78,7 +78,10 @@ export interface WorkspaceBackend { * Run a shell command and wait for it to end. `env` is added to the * environment. A command that fails resolves with its exit code. On * {@link hostBackend}, a command gets only a few variables of this - * process (see {@link commandEnv}), plus `env`. + * process (see {@link commandEnv}), plus `env`. `onStart` gets the id of + * the process when it started, like the `jobId` of `spawn`, for + * {@link killJob}: the host can move a running command to the background. + * A backend that has no id does not call it. */ exec: ( command: string, @@ -87,6 +90,7 @@ export interface WorkspaceBackend { env?: Record timeoutMs?: number signal?: AbortSignal + onStart?: (jobId: string) => void }, ) => Promise<{ exitCode: number; stdout: string; stderr: string }> /** @@ -120,8 +124,9 @@ export interface WorkspaceBackend { * Stop the job with the `jobId` that `spawn` gave, also a job of an * earlier process or machine, with every command it started. Resolves * `true` when the job ran and was stopped, and `false` when no such job - * runs. Optional: without it, a job that a crash left running is not - * stopped. + * runs. Rejects when it cannot check that the process is the job, for + * example when `ps` is missing: then the job may still run. Optional: + * without it, a job that a crash left running is not stopped. */ killJob?: (jobId: string) => Promise /** @@ -209,11 +214,46 @@ export function pathsOf(backend: WorkspaceBackend) { } } +/** + * A workspace folder in `base` for `parts`, like a tenant and a thread id: + * `base/part/part`. Each part is URL-encoded, so a `/`, a `\`, or a `..` in + * an id from a client cannot leave `base` or reach the folder of another + * thread. Throws for an empty part, `.`, or `..`, and when the result is not + * a folder in `base`. A POSIX `base` (like `/srv/work`) gives a POSIX path, + * also on Windows. + * + * @example + * ```ts + * workspaceRoot('/srv/work', 'acme', '../other/t-1') + * // '/srv/work/acme/..%2Fother%2Ft-1' + * ``` + */ +export function workspaceRoot(base: string, ...parts: Array) { + const path = base.startsWith('/') ? nodePath.posix : nodePath + const names = parts.map((part) => { + const name = encodeURIComponent(part) + if (name === '' || name === '.' || name === '..') { + throw new Error(`${JSON.stringify(part)} cannot name a workspace folder.`) + } + return name + }) + const full = path.resolve(base, ...names) + const inBase = path.relative(path.resolve(base), full) + const isIn = + inBase !== '' && + inBase !== '..' && + !inBase.startsWith(`..${path.sep}`) && + !path.isAbsolute(inBase) + if (!isIn) throw new Error(`${full} is not a folder in ${base}.`) + return full +} + /** * The variables of this process that a command gets. The others, like - * provider keys and `DATABASE_URL`, stay in the server. The - * `ai-sandbox-local-process` package has a copy of these lists and of - * {@link commandEnv} in `handle.ts`. Change both. + * provider keys and `DATABASE_URL`, stay in the server. `@tanstack/ai-sandbox` + * has a copy of these lists and of {@link commandEnv} in `command-env.ts`, + * for its local-process provider. This package does not depend on that + * one. Change both. */ const KEPT_ENV = ['PATH', 'HOME', 'LANG', 'LC_ALL', 'TMPDIR', 'TEMP', 'TMP'] /** On Windows, a command also gets these, so programs can start. */ @@ -600,6 +640,10 @@ export const hostBackend = { exec: (command, options = {}) => new Promise((done) => { const { child, kill } = startShell(command, options) + // The same id as `spawn` gives, so `killJob` stops a moved command. + if (child.pid !== undefined) { + options.onStart?.(`${child.pid}:${Date.now()}`) + } const stdout: Array = [] const stderr: Array = [] let size = 0 @@ -683,9 +727,13 @@ export const hostBackend = { // Only the process that started with the job. Another process can have // its pid now. const start = await processStart(pid) - const isJob = - start !== undefined && Math.abs(start - startedAt) <= START_SLACK_MS - if (!isJob) return false + // Not `false`: the job can be this process, and still run. + if (start === undefined) { + throw new Error( + `Cannot check that process ${pid} is the job: no start time from ps.`, + ) + } + if (Math.abs(start - startedAt) > START_SLACK_MS) return false if (process.platform === 'win32') { await taskkill(pid) return true diff --git a/packages/ai-harness/src/first-party/coding/bash.ts b/packages/ai-harness/src/first-party/coding/bash.ts index 33d90301e5..c9890a76df 100644 --- a/packages/ai-harness/src/first-party/coding/bash.ts +++ b/packages/ai-harness/src/first-party/coding/bash.ts @@ -199,8 +199,10 @@ function jobsOf(signal: AbortSignal | undefined) { * `tanstack/workspace-tools:job-ended`. * * A command that runs in the foreground supports `detach` of the tool - * context: the host can move it to the background while it runs. Then the - * end of the run does not stop it, and `signal` kills it like a job. + * context, while fewer than `maxJobs` jobs run: the host can move it to the + * background while it runs. Then the end of the run does not stop it, + * `signal` kills it like a job, and `jobs` gets the id of its process in + * the backend, when the backend gave one. */ export function bashTools( env: ToolEnv, @@ -393,6 +395,8 @@ export function bashTools( if (runSignal?.aborted) stop.abort() const job = { kill: () => stop.abort() } let isDone = false + // The id of the process, so a later host can stop a moved command. + let backendId: string | undefined const work = env.backend .exec(command, { cwd, @@ -401,6 +405,9 @@ export function bashTools( ? Math.min(requested, MAX_TIMEOUT_MS) : defaultTimeoutMs, signal: stop.signal, + onStart: (id) => { + backendId = id + }, }) .finally(() => { isDone = true @@ -411,14 +418,29 @@ export function bashTools( report(exitCode, `${stdout}${stderr ? `\nstderr:\n${stderr}` : ''}`), ) // The host can move the command to the background. The timeout still - // applies to it. - const detach = context?.detach + // applies to it. A moved command is a job, so only under the limit. + // ponytail: the limit is checked at the start. Jobs that start before + // the move can go over it by the moved commands. + const detach = jobs.running.size < maxJobs ? context?.detach : undefined if (!detach) return work const moved = detach(work).then((text) => { // The run can end now, and the command goes on as a job of the // session. The session signal kills it, not the run. runSignal?.removeEventListener('abort', onRunAbort) - if (!isDone) jobs.running.add(job) + const jobId = context?.toolCallId + if (!isDone) { + jobs.running.add(job) + // The session logged the move. Add the backend id, so a later + // host can report and stop the command. + if (jobId !== undefined && backendId !== undefined) { + const { threadId } = context ?? {} + options.jobs?.started({ + jobId, + backendId, + ...(threadId !== undefined ? { threadId } : {}), + }) + } + } // The session closed already, so its listener does not run again. if (signal?.aborted) job.kill() return text diff --git a/packages/ai-harness/src/first-party/coding/edit.ts b/packages/ai-harness/src/first-party/coding/edit.ts index 171d164650..0dfb6c8876 100644 --- a/packages/ai-harness/src/first-party/coding/edit.ts +++ b/packages/ai-harness/src/first-party/coding/edit.ts @@ -118,19 +118,36 @@ export function findMatches(content: string, find: string) { } /** - * What one workspace binding saw of each file: its size and change time - * when a tool read it or wrote it. A write checks that the file is still as - * the binding saw it, so an edit of another session or agent is never lost. - * Call each function under the lock of the path. The paths come from - * `env.reach`. + * What each chat run of one workspace binding saw of each file: its size and + * change time when a tool read it or wrote it. A write checks that the file + * is still as its own chat run saw it, so an edit of another session, of a + * subagent, or of a parallel agent is never lost. `chatOf` gives the chat run + * of the call that runs now. Without one, the calls share one record. Call + * each function under the lock of the path. The paths come from `env.reach`. * * ponytail: a change that keeps the size and the `mtimeMs` is not seen. */ -export function fileVersions(env: ToolEnv) { - const seen = new Map() +export function fileVersions( + env: ToolEnv, + chatOf: () => object | undefined = () => undefined, +) { + type Seen = Map + const shared: Seen = new Map() + const byChat = new WeakMap() + /** The record of the chat run of the call that runs now. */ + const records = () => { + const chat = chatOf() + if (!chat) return shared + const known = byChat.get(chat) + if (known) return known + const created: Seen = new Map() + byChat.set(chat, created) + return created + } return { /** Remember the file at `path` as it is now: a tool read or wrote it. */ saw: async (path: string) => { + const seen = records() const info = await env.backend.stat(path) if (info?.type === 'file') { seen.set(path, { size: info.size, mtimeMs: info.mtimeMs }) @@ -140,7 +157,7 @@ export function fileVersions(env: ToolEnv) { }, /** Forget `path`: a tool removed it. */ forget: (path: string) => { - seen.delete(path) + records().delete(path) }, /** Is something at `path`? Call it before the lock. */ exists: async (path: string) => @@ -151,7 +168,7 @@ export function fileVersions(env: ToolEnv) { * there before the lock (`missing`), and a file is there now. */ check: async (path: string, { missing }: { missing: boolean }) => { - const known = seen.get(path) + const known = records().get(path) const now = await env.backend.stat(path) const shown = env.shown(path) if (known) { diff --git a/packages/ai-harness/src/first-party/coding/index.ts b/packages/ai-harness/src/first-party/coding/index.ts index d5d0b951f5..48473f5bb9 100644 --- a/packages/ai-harness/src/first-party/coding/index.ts +++ b/packages/ai-harness/src/first-party/coding/index.ts @@ -2,7 +2,7 @@ // plugins use the file system and the shell. export { workspaceTools } from './workspace' export type { WorkspaceBinder, WorkspaceToolsOptions } from './workspace' -export { hostBackend } from './backend' +export { hostBackend, workspaceRoot } from './backend' export type { BoundWorkspace, WorkspaceBackend, diff --git a/packages/ai-harness/src/first-party/coding/ring-buffer.ts b/packages/ai-harness/src/first-party/coding/ring-buffer.ts index 815c432037..f8c8469bc7 100644 --- a/packages/ai-harness/src/first-party/coding/ring-buffer.ts +++ b/packages/ai-harness/src/first-party/coding/ring-buffer.ts @@ -9,8 +9,8 @@ const isContinuation = (byte: number) => (byte & 0xc0) === 0x80 * the oldest bytes are dropped. `text()` gives the output as UTF-8 text, and * starts with a line that tells how many bytes were dropped. * - * The `ai-sandbox` package has a copy of this function in - * `workspace-backend.ts`. Change both. + * `@tanstack/ai-sandbox` has a copy of this function in `command-env.ts`. + * This package does not depend on that one. Change both. * * @example * ```ts diff --git a/packages/ai-harness/src/first-party/coding/workspace.ts b/packages/ai-harness/src/first-party/coding/workspace.ts index 0107a465ab..f6aa445a93 100644 --- a/packages/ai-harness/src/first-party/coding/workspace.ts +++ b/packages/ai-harness/src/first-party/coding/workspace.ts @@ -196,9 +196,14 @@ export function createLock() { } /** - * The file locks of this process, by the bound workspace root and the real - * path of the file. Every session and agent on one workspace shares them, so - * two edits of one file run one after the other. + * The file locks of this process, by the place of the backend (a sandbox, or + * this host) and the real path of the file. Every session and agent on one + * file shares them, also through another root or a link, so two edits of + * one file run one after the other. Two sandboxes with the same root do not. + * + * ponytail: a backend with no `binding` counts as this host. Two such + * backends of other machines with the same paths share locks: they wait + * when they need not, but no edit is lost. * * ponytail: a lock is for this process only. Two processes or machines on * one shared sandbox do not wait for each other. The unchanged check of the @@ -242,17 +247,24 @@ export const ThreadWorkspace = createCapability()( 'tanstack/thread-workspace', ) -/** The run of the workspace tool call that runs now. */ -const callRuns = new AsyncLocalStorage() +/** + * The workspace tool call that runs now: its run, and its chat run (the + * abort signal of the call, one for each `chat()`). A tool-mode subagent + * shares the run of its turn, but has its own chat run. + */ +const calls = new AsyncLocalStorage<{ run?: RunContext; chat?: object }>() /** * The run of the workspace tool call that runs now, or `undefined` outside * a call. A `WorkspaceHooks` hook reads it to know which run wrote a file. */ export function writingRun() { - return callRuns.getStore() + return calls.getStore()?.run } +/** The chat run of the workspace tool call that runs now. See `calls`. */ +const callingChat = () => calls.getStore()?.chat + /** * The tools for one workspace: `backend`, with the boundary `root`, and what * each call touches for `PermissionResources`: the paths a call names, their @@ -343,6 +355,18 @@ function bindTools( /** `full` with every link resolved, as a path in `root` when it is in it. */ const resolved = async (full: string) => forRules(await realPath(full)) + /** + * The lock key of `path`, a path from `reach`: the place of the backend, + * and the real path (`forRules` the other way). + */ + const lockKey = async (path: string) => { + const real = within(root, path, paths.sep) + ? paths.join(await realRoot(), paths.relative(root, path)) + : path + const place = backend.binding?.() + return JSON.stringify([place?.provider, place?.sandboxId, real]) + } + const shown = (full: string) => (within(root, full, paths.sep) ? paths.relative(root, full) || '.' : full) .split(paths.sep) @@ -473,11 +497,10 @@ function bindTools( backend, root, hooks: session.hooks ?? (() => []), - // `path` is the real path from `reach`. The key has the root, so the - // fixed lock order of `patch` (by path) is the same for every session on - // this root. - lock: (path, fn) => - fileLocks(`${root}\0${path}`, async () => { + // `path` is the path from `reach`. The fixed lock order of `patch` (by + // path) is the order of the real paths too: one root is one prefix. + lock: async (path, fn) => + fileLocks(await lockKey(path), async () => { // `reach` checked `path` before the lock. A link made since then // would move the work to another file, so check it again here. const now = await resolved(path) @@ -502,9 +525,9 @@ function bindTools( }) !== 'allow' }, } - // What this binding read and wrote, so a write fails when the file changed - // since then. - const versions = fileVersions(env) + // What each chat run of this binding read and wrote, so a write fails + // when the file changed since its own chat run saw it. + const versions = fileVersions(env, callingChat) // `read_file` reads under the lock. Its lock also remembers the file, as // the read sees it. const readEnv: ToolEnv = { @@ -641,9 +664,19 @@ async function openWorkspace( if (binding.kind === 'host' && hostId) await log.pin(hostId) } // A host before this one started these jobs. Background jobs do not move - // with a session, so each one is stopped and reported. + // with a session, so each one is stopped and reported. A backend with no + // `killJob`, or one that cannot check the job (it throws), cannot tell: + // the job may still run, so the log keeps it. for (const { jobId, backendId } of log.lostJobs(threadId)) { - const isStopped = await (found.killJob?.(backendId) ?? false) + const isStopped = await ( + found.killJob?.(backendId) ?? Promise.resolve(undefined) + ).catch(() => undefined) + if (isStopped === undefined) { + await session.note?.( + `Background job ${jobId} may still run: it ran on a host before this one, and this host cannot stop it.`, + ) + continue + } await session.note?.( isStopped ? `Background job ${jobId} was stopped: it ran on a host before the session moved to this one.` @@ -838,14 +871,16 @@ export function createWorkspaceTools( options.workspace === undefined ? fixedTools(options.root, options, session) : threadTools(options.workspace, options, session) - // Each call runs with its run in `callRuns`, for `writingRun`. + // Each call runs with its run and its chat run in `calls`. const tools = source.tools.map((tool): AnyTool => { const { execute } = tool if (!execute) return tool return { ...tool, execute: (args: unknown, context?: ToolExecutionContext) => - callRuns.run(runOf(context), () => execute(args, context)), + calls.run({ run: runOf(context), chat: context?.abortSignal }, () => + execute(args, context), + ), } }) return { @@ -960,12 +995,13 @@ function withEditStyle( * @example * ```ts * workspaceTools({ root: process.cwd(), outside: 'ask' }) - * // A folder for each thread: + * // A folder for each thread. `workspaceRoot` keeps a thread id from a + * // client in its own folder: * workspaceTools({ * workspace: { * bind: async ({ run }) => ({ * ...hostBackend, - * root: `/srv/work/${run.threadId}`, + * root: workspaceRoot('/srv/work', run.tenantId ?? 'default', run.threadId), * }), * }, * }) diff --git a/packages/ai-harness/src/host.ts b/packages/ai-harness/src/host.ts index c355732975..fb37ad3b58 100644 --- a/packages/ai-harness/src/host.ts +++ b/packages/ai-harness/src/host.ts @@ -835,6 +835,26 @@ export interface HarnessHost

{ threadId: string principal: { id: string; tenantId?: string } }) => Promise + /** + * Unpin the thread `threadId` from the host that its folder binding pins + * it to, for example a host that does not come back. Then any host can + * take the thread. Its `workspace.reattach` gets the folder binding of the + * old host: give a new workspace to move the thread there, or `'pinned'` + * to pin it again. The ownership of the old host must have expired: + * while it is live, this rejects. A thread with no pin does nothing. With + * access, the user needs the `delete` action. The `tenantId` of + * `principal` picks the tenant. + * + * @example + * ```ts + * await host.unpin({ harness: coder, threadId: 't-1', principal }) + * ``` + */ + unpin: (options: { + harness: AnyHarness + threadId: string + principal?: P + }) => Promise /** * The user of `request`, from `access.identify`, or `null` when * `identify` refuses it. Network code other than the HTTP handler (an MCP @@ -2841,6 +2861,32 @@ export function createHarnessHost

( principal: undefined, }) }, + async unpin({ harness, threadId, principal }) { + if (isClosed) throw hostClosed() + await checkEntry('host.unpin', 'delete', { threadId, principal }) + const logId = logIdOf({ tenantId: principal?.tenantId, logId: threadId }) + const state = await loadLogState({ + store: logStore, + logId, + rootThread: threadId, + }) + // The pin is the host of a folder binding of a thread of the log. + const pins = new Set( + [...state.sessions.values()].flatMap(({ thread }) => + thread.workspace?.kind === 'host' && thread.workspace.hostId + ? [thread.workspace.hostId] + : [], + ), + ) + for (const pin of pins) { + if (await ownership.unpin(logId, harness.name, pin)) return + } + if (pins.size > 0) { + throw new Error( + `The host of the thread ${threadId} still owns it, so it stays pinned. Try again when its ownership expired.`, + ) + } + }, hasAccess: access !== undefined, async identify({ request }) { if (!access) { diff --git a/packages/ai-harness/src/ownership.ts b/packages/ai-harness/src/ownership.ts index 9228111787..56d0e4b713 100644 --- a/packages/ai-harness/src/ownership.ts +++ b/packages/ai-harness/src/ownership.ts @@ -189,6 +189,41 @@ export class HostOwnership { return this.store.setAffinity({ logId, epoch, affinity: hostId }) } + /** + * Unpin `logId` from the host `hostId`. The pin lets only that host + * acquire the log, so this host acquires it with that affinity, clears + * the pin, and gives the log up with a wake time of now: a sweep then + * takes the work that the old host left. `false` when a live host owns + * the log. + */ + async unpin(logId: string, harness: string, hostId: string) { + const owned = this.owned.get(logId) + if (owned !== undefined) { + await this.pin(logId, owned, null) + return true + } + await this.releasing.get(logId)?.catch(() => {}) + const got = await this.store.acquire({ + logId, + harness, + ownerId: this.ownerId, + ttlMs: this.ttlMs, + affinity: hostId, + }) + if (got === null) return false + try { + await this.pin(logId, got.epoch, null) + } finally { + await this.store.release({ + logId, + ownerId: this.ownerId, + epoch: got.epoch, + wakeInMs: 0, + }) + } + return true + } + /** Stop owning `logId` with no store call: the store refused its epoch. */ forget(logId: string) { this.owned.delete(logId) diff --git a/packages/ai-harness/tests/coding-backend.test.ts b/packages/ai-harness/tests/coding-backend.test.ts index 22e8b26401..7b478adab7 100644 --- a/packages/ai-harness/tests/coding-backend.test.ts +++ b/packages/ai-harness/tests/coding-backend.test.ts @@ -7,6 +7,7 @@ import { createHarnessHost, defineHarness, definePlugin } from '../src' import { WorkspaceHooks, hostBackend, + workspaceRoot, workspaceTools, } from '../src/first-party/coding' import { commandEnv } from '../src/first-party/coding/backend' @@ -206,6 +207,21 @@ describe('hostBackend', () => { }) }) +describe('workspaceRoot', () => { + it('keeps each thread id from a client in its own folder in base', () => { + expect(workspaceRoot('/srv/work', 'acme', '../other-tenant/t-1')).toBe( + '/srv/work/acme/..%2Fother-tenant%2Ft-1', + ) + // A child thread id has `/agent/` in it. + expect(workspaceRoot('/srv/work', 'acme', 't-1/agent/c-2')).toBe( + '/srv/work/acme/t-1%2Fagent%2Fc-2', + ) + expect(() => workspaceRoot('/srv/work', 'acme', '..')).toThrow( + 'cannot name a workspace folder', + ) + }) +}) + describe('hostBackend.readLines', () => { /** * The page of a whole-file read, as `read_file` makes it, and the line diff --git a/packages/ai-harness/tests/coding-bash.test.ts b/packages/ai-harness/tests/coding-bash.test.ts index 99118e88c6..eb3da45884 100644 --- a/packages/ai-harness/tests/coding-bash.test.ts +++ b/packages/ai-harness/tests/coding-bash.test.ts @@ -165,6 +165,7 @@ describe('bash', () => { env: { AGENT: '1' }, timeoutMs: 120_000, signal: expect.any(AbortSignal), + onStart: expect.any(Function), }, }, ]) diff --git a/packages/ai-harness/tests/coding-jobs.test.ts b/packages/ai-harness/tests/coding-jobs.test.ts index be9b63be22..bc88266b94 100644 --- a/packages/ai-harness/tests/coding-jobs.test.ts +++ b/packages/ai-harness/tests/coding-jobs.test.ts @@ -83,7 +83,10 @@ function workspace(backend: Partial = {}) { } /** What a call of the tool gets from the run: its signal and `detach`. */ -type CallContext = Pick +type CallContext = Pick< + ToolExecutionContext, + 'abortSignal' | 'detach' | 'toolCallId' | 'threadId' +> /** The `bash` tool of one plugin setup. */ function bashTool(env: ToolEnv, options?: BashOptions) { @@ -230,6 +233,49 @@ describe('bash in the foreground', () => { expect(execs[0]?.aborted).toBe(true) }) + it('cannot move a command at the job limit', async () => { + const { env } = workspace() + const bash = bashTool(env, { maxJobs: 1 }) + await startJobs(bash, 1) + const run = new AbortController() + const detach = vi.fn(async () => 'Moved to the background.') + + const result = bash( + { command: 'npm test' }, + { abortSignal: run.signal, detach }, + ) + run.abort() + + expect(await result).toBe('exit code: 1\n') + expect(detach).not.toHaveBeenCalled() + }) + + it('gives the session the backend id of a moved command', async () => { + const { env } = workspace({ + exec: (_command, options) => { + options?.onStart?.('pid-3') + return new Promise(() => {}) + }, + }) + const started = vi.fn() + const bash = bashTool(env, { jobs: { started, ended: vi.fn() } }) + + await bash( + { command: 'npm run dev' }, + { + toolCallId: 'call-1', + threadId: 't', + detach: async () => 'Moved to the background.', + }, + ) + + expect(started).toHaveBeenCalledWith({ + jobId: 'call-1', + backendId: 'pid-3', + threadId: 't', + }) + }) + it('counts a moved command as a job of the session', async () => { const { env } = workspace() const signal = new AbortController().signal diff --git a/packages/ai-harness/tests/coding-locks.test.ts b/packages/ai-harness/tests/coding-locks.test.ts index 4ea1dd558d..0767039f08 100644 --- a/packages/ai-harness/tests/coding-locks.test.ts +++ b/packages/ai-harness/tests/coding-locks.test.ts @@ -1,4 +1,11 @@ -import { mkdtemp, readFile, rm, writeFile } from 'node:fs/promises' +import { + mkdtemp, + readFile, + rm, + symlink, + unlink, + writeFile, +} from 'node:fs/promises' import { tmpdir } from 'node:os' import { join } from 'node:path' import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest' @@ -26,13 +33,19 @@ afterEach(async () => { const put = (path: string, text: string) => writeFile(join(root, path), text) const read = (path: string) => readFile(join(root, path), 'utf8') -/** The workspace tools of one session on `root`. Runs a tool by its name. */ -function openSession(backend: WorkspaceBackend = hostBackend) { - const { tools } = createWorkspaceTools({ root, backend, web: false }) - return async (name: string, args: object) => { +/** + * The workspace tools of one session on `at` (default: `root`). Runs a tool + * by its name. With `signal`, the call runs in the chat run of that signal. + */ +function openSession(backend: WorkspaceBackend = hostBackend, at = root) { + const { tools } = createWorkspaceTools({ root: at, backend, web: false }) + return async (name: string, args: object, signal?: AbortSignal) => { const tool = tools.find((each) => each.name === name) if (!tool?.execute) throw new Error(`No tool ${name}.`) - const result: unknown = await tool.execute(args) + const result: unknown = await tool.execute( + args, + signal && { emitCustomEvent: () => undefined, abortSignal: signal }, + ) return result } } @@ -114,6 +127,41 @@ describe('file locks across sessions', () => { expect(await read('a.txt')).toBe('first\n') }) + it('makes an edit through another spelling of the root wait for the lock', async () => { + await put('a.txt', 'one\n') + const link = `${root}-link` + await symlink(root, link, 'junction') + try { + const first = watchedBackend({ held: true }) + const firstCall = openSession(first.backend) + const secondCall = openSession(hostBackend, link) + await firstCall('read_file', { path: 'a.txt' }) + await secondCall('read_file', { path: 'a.txt' }) + + const firstEdit = firstCall('edit_file', { + path: 'a.txt', + old: 'one', + new: 'first', + }) + await first.reached + const secondEdit = tracked( + secondCall('edit_file', { path: 'a.txt', old: 'one', new: 'second' }), + ) + // One folder, one lock: the second edit waits for the first. + await settle() + expect(secondEdit.state.settled).toBe(false) + + first.release() + await firstEdit + await expect(secondEdit.done).rejects.toThrow( + 'a.txt changed since it was read.', + ) + expect(await read('a.txt')).toBe('first\n') + } finally { + await unlink(link) + } + }) + it('lets the second session edit after it reads the file again', async () => { await put('a.txt', 'one two\n') const firstCall = openSession() @@ -160,6 +208,23 @@ describe('file locks across sessions', () => { }) }) +describe('the unchanged check of each chat run', () => { + it('fails a write of the lead after a subagent of the session changed the file', async () => { + await put('a.txt', 'one\n') + const call = openSession() + const lead = new AbortController().signal + const subagent = new AbortController().signal + await call('read_file', { path: 'a.txt' }, lead) + await call('read_file', { path: 'a.txt' }, subagent) + await call('write_file', { path: 'a.txt', content: 'the fix\n' }, subagent) + + await expect( + call('write_file', { path: 'a.txt', content: 'one, again\n' }, lead), + ).rejects.toThrow('a.txt changed since it was read.') + expect(await read('a.txt')).toBe('the fix\n') + }) +}) + describe('the unchanged check', () => { it.each([ ['edit_file', { path: 'a.txt', old: 'one', new: 'two' }], diff --git a/packages/ai-harness/tests/workspace-binding.test.ts b/packages/ai-harness/tests/workspace-binding.test.ts index 0cd02eff7b..11a929efb0 100644 --- a/packages/ai-harness/tests/workspace-binding.test.ts +++ b/packages/ai-harness/tests/workspace-binding.test.ts @@ -315,6 +315,29 @@ describe('the workspace binding in the session log', () => { await again.host.close() }) + it('unpins a folder thread from a host that does not come back', async () => { + const persistence = memoryPersistence() + const { log } = persistence.stores + const inFolder = { root: '/ws', backend: backend() } + expect(await readOnce(persistence, inFolder, 'c0', 'pod-a')).toBe(READ) + expect(await canOtherHostTake(log, LOG)).toBe(false) + + const other = createHarnessHost({ + persistence, + ownership: { hostId: 'pod-b' }, + }) + await other.unpin({ + harness: defineHarness({ + name: 'test/workspace-binding', + adapter: readTurn('b0').adapter, + }), + threadId: THREAD, + }) + + expect(await canOtherHostTake(log, LOG)).toBe(true) + await other.close() + }) + it('refuses a folder thread of another hostId with the pinned error', async () => { const persistence = memoryPersistence() const { log } = persistence.stores @@ -358,8 +381,11 @@ describe('the workspace binding in the session log', () => { ) }) - it('stops a job that a crashed host left running, and notes it', async () => { - const persistence = memoryPersistence() + /** + * A host that starts the background job `job-7` on `persistence`, then + * stops with no close, and its ownership expires. + */ + async function leaveRunningJob(persistence: HarnessPersistence) { const spawn = () => ({ wait: () => new Promise<{ exitCode: number }>(() => {}), kill: () => {}, @@ -388,6 +414,11 @@ describe('the workspace binding in the session log', () => { ]) // The first host never closes, like a process that stopped. await ownershipExpired() + } + + it('stops a job that a crashed host left running, and notes it', async () => { + const persistence = memoryPersistence() + await leaveRunningJob(persistence) const killJob = vi.fn(async () => true) const { adapter } = readTurn('c1') const next = open({ @@ -416,4 +447,36 @@ describe('the workspace binding in the session log', () => { ]) await next.host.close() }) + + it('keeps a job that the backend cannot check, and says that it may still run', async () => { + const persistence = memoryPersistence() + await leaveRunningJob(persistence) + // Like `hostBackend` with no `ps`: it cannot tell if the job runs. + const killJob = vi.fn(async () => { + throw new Error('no ps') + }) + const { adapter } = readTurn('c1') + const next = open({ + persistence, + adapter, + plugins: [ + workspaceTools({ + root: '/ws', + backend: backend({ killJob }), + web: false, + }), + ], + }) + + await (await next.opened).send('read') + + expect(await recordsOf(persistence, 'harness.job')).toMatchObject([ + { jobId: 'bash-1', ended: false }, + ]) + const notes = await recordsOf(persistence, 'harness.note') + expect(notes.map((record) => record.note)).toEqual([ + 'Background job bash-1 may still run: it ran on a host before this one, and this host cannot stop it.', + ]) + await next.host.close() + }) }) From 61dbf9283cd483a4045872a5723d8142e7ece801 Mon Sep 17 00:00:00 2001 From: Alem Tuzlak Date: Mon, 19 Oct 2026 13:55:03 +0200 Subject: [PATCH 11/17] fix(ai-sandbox, ai-sandbox-local-process): job ids with a start time, one commandEnv, and passEnv - sandboxWorkspaceBackend gives a job id of pid and start time (ps -o lstart), and killJob stops only the process that started then. It rejects when it cannot check a process that runs. - commandEnv and createOutputTail live in @tanstack/ai-sandbox, and local-process imports commandEnv from there. - localProcessSandbox({ passEnv }) gives commands more variables of the host, for a CLI login or a proxy. --- .../ai-sandbox-local-process/src/handle.ts | 43 +------ .../ai-sandbox-local-process/src/provider.ts | 11 ++ .../tests/local-process.test.ts | 20 +++ packages/ai-sandbox/src/command-env.ts | 118 ++++++++++++++++++ packages/ai-sandbox/src/index.ts | 4 + packages/ai-sandbox/src/workspace-backend.ts | 110 +++++++--------- .../tests/workspace-backend.test.ts | 61 +++++++-- 7 files changed, 252 insertions(+), 115 deletions(-) create mode 100644 packages/ai-sandbox/src/command-env.ts diff --git a/packages/ai-sandbox-local-process/src/handle.ts b/packages/ai-sandbox-local-process/src/handle.ts index 7657dedbe0..2b29a3bae4 100644 --- a/packages/ai-sandbox-local-process/src/handle.ts +++ b/packages/ai-sandbox-local-process/src/handle.ts @@ -16,6 +16,7 @@ import * as path from 'node:path' import { DEFAULT_WORKSPACE_ROOT, UnsupportedCapabilityError, + commandEnv, createExecBackedGit, } from '@tanstack/ai-sandbox' import type { ChildProcess } from 'node:child_process' @@ -111,45 +112,6 @@ function prependShellPath(env: NodeJS.ProcessEnv): NodeJS.ProcessEnv { return env } -/** - * The variables of this process that a command gets. The others, like - * provider keys and `DATABASE_URL`, stay in the server. - */ -const KEPT_ENV = ['PATH', 'HOME', 'LANG', 'LC_ALL', 'TMPDIR', 'TEMP', 'TMP'] -/** On Windows, a command also gets these, so programs can start. */ -const KEPT_WINDOWS_ENV = [ - 'SystemRoot', - 'ComSpec', - 'PATHEXT', - 'WINDIR', - 'USERPROFILE', -] - -/** - * The environment of a command: the variables of this process in a short - * list (`PATH`, `HOME`, `LANG`, `LC_ALL`, `TMPDIR`, `TEMP`, `TMP`, and on - * Windows `SystemRoot`, `ComSpec`, `PATHEXT`, `WINDIR`, `USERPROFILE`), plus - * `env`. No other variable of this process goes to the command. On Windows - * it sets `NoDefaultCurrentDirectoryInExePath`, so a program name is found - * on the PATH only, and not in the current folder. - * - * A copy of `commandEnv` in `@tanstack/ai-harness` - * (`src/first-party/coding/backend.ts`). This package does not depend on - * that one. Change both. - */ -function commandEnv(env: Record, platform: NodeJS.Platform) { - const isWindows = platform === 'win32' - const names = isWindows ? [...KEPT_ENV, ...KEPT_WINDOWS_ENV] : KEPT_ENV - const kept: Record = {} - for (const name of names) { - // On Windows, `process.env` finds a name in any letter case. - const value = process.env[name] - if (value !== undefined) kept[name] = value - } - if (!isWindows) return { ...kept, ...env } - return { ...kept, ...env, NoDefaultCurrentDirectoryInExePath: '1' } -} - export const LOCAL_PROCESS_CAPS: SandboxCapabilities = { fs: true, exec: true, @@ -622,6 +584,8 @@ export interface LocalProcessHandleOptions { forkFactory: (sourceRoot: string) => Promise /** Env vars to delete from the env of each command (see `commandEnv`). */ scrubEnv?: Array + /** Env vars of this process that each command also gets (see `commandEnv`). */ + passEnv?: Array /** * Sink for non-fatal teardown diagnostics — currently a `killTree` that could * not confirm the process tree is gone. Teardown never throws, so without a @@ -886,6 +850,7 @@ export class LocalProcessHandle implements SandboxHandle { const env: NodeJS.ProcessEnv = commandEnv( { ...this.envVars, ...extra }, process.platform, + this.options.passEnv, ) // Drop scrubbed vars so a host CLI falls back to its own stored auth // (e.g. remove ANTHROPIC_API_KEY → Claude Code uses the logged-in diff --git a/packages/ai-sandbox-local-process/src/provider.ts b/packages/ai-sandbox-local-process/src/provider.ts index 8cccf2687f..241a851570 100644 --- a/packages/ai-sandbox-local-process/src/provider.ts +++ b/packages/ai-sandbox-local-process/src/provider.ts @@ -39,6 +39,14 @@ export interface LocalProcessSandboxConfig { * of billing the API. */ scrubEnv?: Array + /** + * More env vars of this process for each command, added to the short list + * above. Use it for what a host CLI or the network needs, for example + * `CLAUDE_CONFIG_DIR`, `CODEX_HOME`, `XDG_CONFIG_HOME`, `HTTPS_PROXY`, + * `NO_PROXY`, `NODE_EXTRA_CA_CERTS`, or `SSH_AUTH_SOCK`. `scrubEnv` still + * removes a name in both lists. + */ + passEnv?: Array /** * Sink for non-fatal teardown diagnostics — a `killTree` that could not * confirm the spawned process tree is gone. Teardown is total by construction @@ -66,6 +74,7 @@ class LocalProcessProvider implements SandboxProvider { root, removeOnDestroy: this.removeDefault(), scrubEnv: this.config.scrubEnv, + passEnv: this.config.passEnv, logger: this.config.logger, forkFactory: async (sourceRoot) => { const dest = path.join(this.baseDir(), `fork-${randomUUID()}`) @@ -74,6 +83,8 @@ class LocalProcessProvider implements SandboxProvider { return new LocalProcessHandle({ root: dest, removeOnDestroy: true, + scrubEnv: this.config.scrubEnv, + passEnv: this.config.passEnv, logger: this.config.logger, forkFactory: () => Promise.reject(new Error('nested fork unsupported')), diff --git a/packages/ai-sandbox-local-process/tests/local-process.test.ts b/packages/ai-sandbox-local-process/tests/local-process.test.ts index c516645104..fa493b1b27 100644 --- a/packages/ai-sandbox-local-process/tests/local-process.test.ts +++ b/packages/ai-sandbox-local-process/tests/local-process.test.ts @@ -133,6 +133,26 @@ describe('local-process process', () => { } }) + it('gives a command the server variables named in passEnv', async () => { + const sbx = await localProcessSandbox({ + baseDir, + removeOnDestroy: true, + passEnv: ['LP_PROXY'], + }).create({}) + process.env.LP_PROXY = 'http://proxy:8080' + process.env.LP_SERVER_SECRET = 'hunter2' + try { + const r = await sbx.process.exec( + 'echo "$LP_PROXY ${LP_SERVER_SECRET:-unset}"', + ) + expect(r.stdout.trim()).toBe('http://proxy:8080 unset') + } finally { + delete process.env.LP_PROXY + delete process.env.LP_SERVER_SECRET + await sbx.destroy() + } + }) + it('gives a command the env of env.set and of the call', async () => { const sbx = await fresh() await sbx.env.set({ LP_SET: 'from-set' }) diff --git a/packages/ai-sandbox/src/command-env.ts b/packages/ai-sandbox/src/command-env.ts new file mode 100644 index 0000000000..a0c45623f2 --- /dev/null +++ b/packages/ai-sandbox/src/command-env.ts @@ -0,0 +1,118 @@ +/** + * The variables of this process that a command gets. The others, like + * provider keys and `DATABASE_URL`, stay in the server. + */ +const KEPT_ENV = ['PATH', 'HOME', 'LANG', 'LC_ALL', 'TMPDIR', 'TEMP', 'TMP'] +/** On Windows, a command also gets these, so programs can start. */ +const KEPT_WINDOWS_ENV = [ + 'SystemRoot', + 'ComSpec', + 'PATHEXT', + 'WINDIR', + 'USERPROFILE', +] + +/** + * The environment of a command on this machine: the variables of this + * process in a short list (`PATH`, `HOME`, `LANG`, `LC_ALL`, `TMPDIR`, + * `TEMP`, `TMP`, and on Windows `SystemRoot`, `ComSpec`, `PATHEXT`, + * `WINDIR`, `USERPROFILE`), plus the variables named in `passEnv`, plus + * `env`. No other variable of this process goes to the command. On Windows + * it sets `NoDefaultCurrentDirectoryInExePath`, so a program name is found + * on the PATH only, and not in the current folder. + * + * `@tanstack/ai-harness` keeps a copy for its `hostBackend` + * (`src/first-party/coding/backend.ts`), because it does not depend on this + * package. Change both. + * + * @example + * ```ts + * spawn('sh', ['-c', command], { + * env: commandEnv({ AGENT: '1' }, process.platform, ['HTTPS_PROXY']), + * }) + * ``` + */ +export function commandEnv( + env: Record | undefined, + platform: NodeJS.Platform, + passEnv: ReadonlyArray = [], +) { + const isWindows = platform === 'win32' + const names = isWindows + ? [...KEPT_ENV, ...KEPT_WINDOWS_ENV, ...passEnv] + : [...KEPT_ENV, ...passEnv] + const kept: Record = {} + for (const name of names) { + // On Windows, `process.env` finds a name in any letter case. + const value = process.env[name] + if (value !== undefined) kept[name] = value + } + if (!isWindows) return { ...kept, ...env } + return { ...kept, ...env, NoDefaultCurrentDirectoryInExePath: '1' } +} + +const encoder = new TextEncoder() +const decoder = new TextDecoder() + +/** True for a byte in the middle of a UTF-8 character (`10xxxxxx`). */ +const isContinuation = (byte: number) => (byte & 0xc0) === 0x80 + +/** + * The newest output of a command, at most `maxBytes` bytes. When more comes, + * the oldest bytes are dropped. `text()` gives the output as UTF-8 text, and + * starts with a line that tells how many bytes were dropped. + * + * `@tanstack/ai-harness` keeps a copy (`src/first-party/coding/ring-buffer.ts`), + * because it does not depend on this package. Change both. + * + * @example + * ```ts + * const output = createOutputTail(1024 * 1024) + * child.stdout.on('data', output.push) + * output.text() // '[12 earlier bytes were dropped.]\n...' + * ``` + */ +export function createOutputTail(maxBytes: number) { + // The chunks from `start` on are the output. The ones before are dropped. + const chunks: Array = [] + let start = 0 + let size = 0 + let dropped = 0 + + const push = (chunk: Uint8Array | string) => { + const bytes = typeof chunk === 'string' ? encoder.encode(chunk) : chunk + // A copy of the end only, so a big chunk is not kept whole. + const kept = bytes.byteLength > maxBytes ? bytes.slice(-maxBytes) : bytes + dropped += bytes.byteLength - kept.byteLength + chunks.push(kept) + size += kept.byteLength + while (size > maxBytes) { + const first = chunks[start] + if (first === undefined) break + const over = size - maxBytes + const cut = Math.min(over, first.byteLength) + if (cut === first.byteLength) start += 1 + else chunks[start] = first.subarray(cut) + size -= cut + dropped += cut + } + // Remove the dropped chunks from the array now and then, not each time. + if (start > 1024 && start * 2 > chunks.length) { + chunks.splice(0, start) + start = 0 + } + } + + const text = () => { + let bytes = Buffer.concat(chunks.slice(start)) + if (dropped === 0) return decoder.decode(bytes) + // The cut can be in a character. Start at the next whole one. + let from = 0 + while (from < 3 && from < bytes.length && isContinuation(bytes[from] ?? 0)) + from += 1 + bytes = bytes.subarray(from) + return `[${dropped + from} earlier bytes were dropped.]\n${decoder.decode(bytes)}` + } + + return { push, text } +} diff --git a/packages/ai-sandbox/src/index.ts b/packages/ai-sandbox/src/index.ts index e265eba978..b2172c6ad4 100644 --- a/packages/ai-sandbox/src/index.ts +++ b/packages/ai-sandbox/src/index.ts @@ -457,6 +457,10 @@ export type { SandboxWatchHandle, } from './watch' +// The env and the output tail of a command on this machine, for providers +// that run commands as host processes. +export { commandEnv, createOutputTail } from './command-env' + // Keying export { computeSandboxKey, computeWorkspaceHash } from './key' export type { SandboxKeyInput } from './key' diff --git a/packages/ai-sandbox/src/workspace-backend.ts b/packages/ai-sandbox/src/workspace-backend.ts index 6ffe02c5f0..1fcf6cfea2 100644 --- a/packages/ai-sandbox/src/workspace-backend.ts +++ b/packages/ai-sandbox/src/workspace-backend.ts @@ -1,4 +1,5 @@ import { posix } from 'node:path' +import { createOutputTail } from './command-env' import { shellQuote } from './journal' import type { WorkspaceBackend } from '@tanstack/ai-harness/plugins/coding' import type { ExecResult, SandboxFs, SandboxHandle } from './contracts' @@ -14,66 +15,6 @@ const MAX_PAGE_BYTES = 64 * 1024 /** The line in the output of `readLines` between the head and the lines. */ const LINES_MARK = '@lines' -const encoder = new TextEncoder() -const decoder = new TextDecoder() - -/** True for a byte in the middle of a UTF-8 character (`10xxxxxx`). */ -const isContinuation = (byte: number) => (byte & 0xc0) === 0x80 - -/** - * The newest output of a command, at most `maxBytes` bytes. When more comes, - * the oldest bytes are dropped. `text()` gives the output as UTF-8 text, and - * starts with a line that tells how many bytes were dropped. - * - * A copy of `createOutputTail` in `@tanstack/ai-harness` - * (`src/first-party/coding/ring-buffer.ts`). That package does not export - * it. Change both. - */ -function createOutputTail(maxBytes: number) { - // The chunks from `start` on are the output. The ones before are dropped. - const chunks: Array = [] - let start = 0 - let size = 0 - let dropped = 0 - - const push = (chunk: Uint8Array | string) => { - const bytes = typeof chunk === 'string' ? encoder.encode(chunk) : chunk - // A copy of the end only, so a big chunk is not kept whole. - const kept = bytes.byteLength > maxBytes ? bytes.slice(-maxBytes) : bytes - dropped += bytes.byteLength - kept.byteLength - chunks.push(kept) - size += kept.byteLength - while (size > maxBytes) { - const first = chunks[start] - if (first === undefined) break - const over = size - maxBytes - const cut = Math.min(over, first.byteLength) - if (cut === first.byteLength) start += 1 - else chunks[start] = first.subarray(cut) - size -= cut - dropped += cut - } - // Remove the dropped chunks from the array now and then, not each time. - if (start > 1024 && start * 2 > chunks.length) { - chunks.splice(0, start) - start = 0 - } - } - - const text = () => { - let bytes = Buffer.concat(chunks.slice(start)) - if (dropped === 0) return decoder.decode(bytes) - // The cut can be in a character. Start at the next whole one. - let from = 0 - while (from < 3 && from < bytes.length && isContinuation(bytes[from] ?? 0)) - from += 1 - bytes = bytes.subarray(from) - return `[${dropped + from} earlier bytes were dropped.]\n${decoder.decode(bytes)}` - } - - return { push, text } -} - /** * The page in the output of the `readLines` command, for a page from line * `first`. The output is the size, the first bytes as hex from `od`, a line @@ -136,8 +77,11 @@ function parsePage(stdout: string, first: number) { * - `binding()` gives the provider and the id of the sandbox. The session * log keeps them, so `reattach` of `workspace` can resume the sandbox on * another machine. - * - The id of a job is its pid in the sandbox. `killJob` stops it, also - * from a backend over the same sandbox on another machine. + * - The id of a job is its pid in the sandbox and its start time from + * `ps`. `killJob` stops it, also from a backend over the same sandbox on + * another machine, and only when the process with that pid started at + * that time. A sandbox without a `ps -o lstart` cannot check it, so + * `killJob` rejects for a job that runs. * * @param handle - The sandbox, from `provider.create()` or `provider.resume()`. * @param options - `owned: true` when the backend owns the sandbox. Then @@ -358,8 +302,12 @@ export function sandboxWorkspaceBackend( started.then((each) => each.kill()).catch(() => undefined) }, output: output.text, + // The pid and the start time: a later process can get the same pid. jobId: started.then( - (each) => String(each.pid), + async (each) => { + const start = await processStart(String(each.pid)) + return start === undefined ? String(each.pid) : `${each.pid}:${start}` + }, () => undefined, ), } @@ -377,13 +325,41 @@ export function sandboxWorkspaceBackend( await destroyed } - // ponytail: a pid only. A sandbox that restarted can give the pid to - // another process, and this stops that process. + /** + * When the process `pid` in the sandbox started, as `ps` gives it, or + * `undefined` when `ps` gives nothing: no such process, or no `ps` that + * can tell (busybox). + */ + const processStart = async (pid: string) => { + const result = await handle.process + .exec(`LC_ALL=C ps -o lstart= -p ${pid}`) + .catch(() => undefined) + const start = result?.exitCode === 0 ? result.stdout.trim() : '' + return start === '' ? undefined : start + } + + // The id is `pid:start`. Only the process that started with the job is + // stopped: after a restart (also of the host of a local-process + // sandbox), another process can have its pid. const killJob: NonNullable = async (jobId) => { - if (!/^\d+$/.test(jobId)) return false + const match = /^(\d+)(?::(.+))?$/.exec(jobId) + const pid = match?.[1] + if (pid === undefined) return false + const start = await processStart(pid) + if (start === undefined) { + // No start time: no such process, or no `ps` to check it with. + const alive = await handle.process.exec(`kill -0 ${pid}`) + if (alive.exitCode !== 0) return false + } + if (start === undefined || match?.[2] === undefined) { + throw new Error( + `Cannot check that process ${pid} in the sandbox is the job: no start time from ps.`, + ) + } + if (start !== match[2]) return false // The process group first, so the commands of the job stop too. const { exitCode } = await handle.process.exec( - `kill -9 -- -${jobId} 2>/dev/null || kill -9 ${jobId}`, + `kill -9 -- -${pid} 2>/dev/null || kill -9 ${pid}`, ) return exitCode === 0 } diff --git a/packages/ai-sandbox/tests/workspace-backend.test.ts b/packages/ai-sandbox/tests/workspace-backend.test.ts index 36e9075f2c..c3ef73ac13 100644 --- a/packages/ai-sandbox/tests/workspace-backend.test.ts +++ b/packages/ai-sandbox/tests/workspace-backend.test.ts @@ -164,9 +164,26 @@ function fakeProcesses( const aborted: Array = [] const killed: Array = [] const commands: Array<{ command: string; cwd?: string }> = [] + // The pids that run, and the start time that `ps` gives for each. + const starts = new Map([['7', 'Mon Oct 19 13:00:00 2026']]) const process: SandboxHandle['process'] = { exec: (command, options = {}) => { commands.push({ command, cwd: options.cwd }) + const ps = /^LC_ALL=C ps -o lstart= -p (\d+)$/.exec(command) + if (ps) { + const start = starts.get(ps[1] ?? '') + return Promise.resolve( + start === undefined + ? { exitCode: 1, stdout: '', stderr: '' } + : { exitCode: 0, stdout: `${start}\n`, stderr: '' }, + ) + } + if (/^kill -0 \d+$/.test(command)) { + return Promise.resolve({ exitCode: 0, stdout: '', stderr: '' }) + } + if (command.startsWith('kill -9 ')) { + return Promise.resolve({ exitCode: 0, stdout: '', stderr: '' }) + } const page = PAGE_COMMAND.exec(command) if (page) { const [, quoted = '', first = '', end = ''] = page @@ -236,7 +253,7 @@ function fakeProcesses( } }, } - return { process, aborted, killed, commands } + return { process, aborted, killed, commands, starts } } /** @@ -252,11 +269,24 @@ function fakeSandbox( } = {}, ) { const { fs, files, dirs, links } = memoryFs(seed, options.withLstat) - const { process, aborted, killed, commands } = fakeProcesses(links, files) + const { process, aborted, killed, commands, starts } = fakeProcesses( + links, + files, + ) const sandbox = makeFakeHandle('sandbox-1', 'fake', options.caps) const handle: SandboxHandle = { ...sandbox, fs, process } const backend = sandboxWorkspaceBackend(handle, { owned: options.owned }) - return { backend, fs, files, dirs, aborted, killed, commands, sandbox } + return { + backend, + fs, + files, + dirs, + aborted, + killed, + commands, + starts, + sandbox, + } } const SEED: FsSeed = { @@ -617,19 +647,32 @@ describe('sandboxWorkspaceBackend on another machine', () => { }) }) - it('gives the pid of a job as its id, and kills a job by that id', async () => { + it('gives the pid and the start time of a job as its id, and kills only that process', async () => { const { backend, commands } = fakeSandbox() const job = backend.spawn?.('serve') - expect(await job?.jobId).toBe('7') + const jobId = await job?.jobId + expect(jobId).toBe('7:Mon Oct 19 13:00:00 2026') - await backend.killJob('7') + // After a restart, another process can have the pid. + expect(await backend.killJob('7:Sun Oct 18 09:00:00 2026')).toBe(false) + expect(await backend.killJob(jobId ?? '')).toBe(true) - expect(commands.map(({ command }) => command)).toEqual([ - 'kill -9 -- -7 2>/dev/null || kill -9 7', - ]) + const kills = commands + .map(({ command }) => command) + .filter((command) => command.startsWith('kill -9')) + expect(kills).toEqual(['kill -9 -- -7 2>/dev/null || kill -9 7']) job?.kill() }) + it('rejects for a job that runs when ps gives no start time', async () => { + const { backend, starts } = fakeSandbox() + // Like busybox: no `ps -o lstart`, but the process runs. + starts.clear() + await expect(backend.killJob('7:Mon Oct 19 13:00:00 2026')).rejects.toThrow( + 'Cannot check that process 7 in the sandbox is the job', + ) + }) + it('runs no command for a job id that is not a pid', async () => { const { backend, commands } = fakeSandbox() expect(await backend.killJob('7; rm -rf /')).toBe(false) From 3f34d56ce3b86d086616cec3830ab8d3f8776cdb Mon Sep 17 00:00:00 2001 From: Alem Tuzlak Date: Mon, 19 Oct 2026 13:55:07 +0200 Subject: [PATCH 12/17] docs(ai-harness, ai-sandbox): safe workspace roots, host.unpin, a stable hostId, no isolation claims for host commands, and passEnv --- .changeset/ai-harness-coding-tools.md | 9 +++--- .changeset/ai-sandbox-harness.md | 4 ++- .changeset/sandbox-local-process-env.md | 2 +- docs/config.json | 7 +++-- docs/harness/coding-tools-backends.md | 3 +- docs/harness/coding-tools.md | 2 +- docs/harness/coding-workspaces.md | 42 ++++++++++++++++++------- docs/sandbox/auth.md | 21 ++++++++++--- docs/sandbox/providers.md | 14 ++++++++- 9 files changed, 76 insertions(+), 28 deletions(-) diff --git a/.changeset/ai-harness-coding-tools.md b/.changeset/ai-harness-coding-tools.md index 85f989ea92..18912cfffd 100644 --- a/.changeset/ai-harness-coding-tools.md +++ b/.changeset/ai-harness-coding-tools.md @@ -6,12 +6,13 @@ Add coding tools in the new Node entry `@tanstack/ai-harness/plugins/coding`: - `workspaceTools({ root })` gives `read_file` (text pages, images, PDFs), `write_file`, `edit_file` or `patch` (`editStyle`), `list_files`, `grep`, `bash` (with `background`), `webfetch`, and `websearch` (with a `SearchProvider`). With `permissions()`, edits, `bash`, and `webfetch` ask first, and the rules also apply to the real path of a link. `outside: 'ask'` makes a path outside `root` need an approval. `permissions()` decides outside paths and the real paths of links before the call runs, so a tool never asks while it runs. The `cwd` thread setting gives each thread its own working folder. - `workspaceTools({ workspace: { bind, reattach? } })` gives each thread its own workspace. The root of the workspace is the boundary of every tool. The session log keeps where the workspace is (`harness.workspace.bound`). On another host, `reattach` gives the workspace back (for example a resumed sandbox), or `'pinned'` or `'gone'`. A gone workspace sends `tanstack/workspace-tools:stranded`. With this plugin first, `snapshots()` and `formatter()` use the workspace of each thread, and their `root` is optional. -- `createHarnessHost({ ownership: { hostId } })` gives a host a stable id. Then a thread that uses a folder on that host is pinned to it, and other hosts refuse its tool calls. Without `hostId`, threads are not pinned. +- `createHarnessHost({ ownership: { hostId } })` gives a host a stable id. Then a thread that uses a folder on that host is pinned to it, and other hosts refuse its tool calls. Without `hostId`, threads are not pinned. `host.unpin({ harness, threadId })` unpins a thread from a host that does not come back, once its ownership expired. +- `workspaceRoot(base, ...parts)` gives a folder in `base` for a tenant and a thread id. It encodes each part, so a thread id from a client cannot leave `base` or reach another thread. - Commands, `rg`, `git`, and the formatters get only a short list of variables of the server process, plus the `env` of the backend call. Git runs without the hooks and the `fsmonitor` of the project. -- The tools resolve links with `realpath` under the file lock, and fail when a backend cannot resolve a path. A write fails when the file changed after the tool read it. The file locks are shared by the sessions of a host. -- Background jobs: at most 8 per session (`maxJobs`), the newest 1 MiB of output per job, and the events `tanstack/workspace-tools:job-started`, `job-output`, and `job-ended`. A job that an earlier host left running is stopped at the next tool call, with a note to the model. The wake turn of a job runs as the user who started it. The `bash` spill file has `[redacted]` in place of each secret of the session. +- The tools resolve links with `realpath` under the file lock, and fail when a backend cannot resolve a path. A write fails when the file changed after the tool read it in the same chat run, so a lead and its subagents do not lose each other's edits. The file locks are shared by the sessions of a host, by the sandbox (or host) and the real path of the file. +- Background jobs: at most 8 per session (`maxJobs`), the newest 1 MiB of output per job, and the events `tanstack/workspace-tools:job-started`, `job-output`, and `job-ended`. A running command that `session.background()` moves counts as a job, and its process id goes to the log. A job that an earlier host left running is stopped at the next tool call, with a note to the model. When the backend cannot check the job (for example with no `ps`), the log keeps it, and the note says that it can still run. The wake turn of a job runs as the user who started it. The `bash` spill file has `[redacted]` in place of each secret of the session. - `read_file` sends images and PDFs up to 256 KB. `webfetch` gives the model the first 256 KB of a page. -- A `WorkspaceBackend` runs the tools. `hostBackend` is the default. `realpath` is required. `spawn` can give a `jobId`, and `killJob`, `binding`, and `close` are optional. On Windows, commands do not run a program from the working folder by its name. +- A `WorkspaceBackend` runs the tools. `hostBackend` is the default. `realpath` is required. `spawn` can give a `jobId`, `exec` can give one with `onStart`, and `killJob`, `binding`, and `close` are optional. `killJob` rejects when it cannot check that the process is the job. On Windows, commands do not run a program from the working folder by its name. - `snapshots()` adds `/undo` and `/redo`, which restore only the files of the turn (git 2.26 or newer), file restore for `session.revert()`, and a diff for each model step. Each thread of each tenant has its own shadow repository, which never stores `.env` files. A step keeps only the files of its own run. The steps and the undo point are in the session log, so another host with the same repository can undo. A revert uses the last 500 steps. `host.sessions.delete` removes the repository of the thread, and `purgeSnapshots()` does it from your own code. A failed snapshot sends `tanstack/snapshots:failed`. - `formatter()` formats each written file with the formatter of the project. With `permissions()`, the rules decide the formatter with the call that writes, before that call runs: one approval covers both, and the formatter runs after the write. An `always` answer is tied to a hash of the formatter config files. A write of a formatter config below the root, or under `node_modules`, asks each time, and an `always` answer saves no rule for that formatter run. - `WorkspaceHooks` lets other plugins run code after a read or a write. diff --git a/.changeset/ai-sandbox-harness.md b/.changeset/ai-sandbox-harness.md index ea295c6dcf..aed6253d07 100644 --- a/.changeset/ai-sandbox-harness.md +++ b/.changeset/ai-sandbox-harness.md @@ -6,6 +6,8 @@ Add `@tanstack/ai-sandbox/harness` for `@tanstack/ai-harness`: - `codingAgents({ sandbox, agents, workspace? })` is a harness plugin. The lead model gets one tool per coding agent (Claude Code, Codex, Grok Build, or any ACP agent). Each agent works in the sandbox, keeps its own session per thread (also after a restart), and starts read-only in the `plan` mode of `permissions()`. `/fresh [agent]` starts a new session. `workspace: 'shared'` (default) runs the agents of a thread one at a time in one sandbox. `'per-agent'` gives each agent its own sandbox. - `codingAgents` runs agents only on the host that owns the log of the thread, so `'shared'` runs one agent at a time also with many hosts. -- `sandboxWorkspaceBackend(sandbox, { owned? })` runs the coding tools of `@tanstack/ai-harness/plugins/coding` (`workspaceTools()`, `snapshots()`, `formatter()`) in a sandbox. It resolves links in the sandbox, reads one page of a file in the sandbox, and stops background jobs by their id, also from another host. `binding()` gives the provider and the sandbox id, so another host can resume the sandbox. `close()` stops the jobs of the backend. A session close closes it with `'moved'`, so an owned sandbox stays for the next host. `host.sessions.delete` closes it with `'closed'`, and with `owned: true` that destroys the sandbox. +- `sandboxWorkspaceBackend(sandbox, { owned? })` runs the coding tools of `@tanstack/ai-harness/plugins/coding` (`workspaceTools()`, `snapshots()`, `formatter()`) in a sandbox. It resolves links in the sandbox, reads one page of a file in the sandbox, and stops background jobs by their id (the pid and its start time), also from another host. It never stops another process that got the pid of a job. `binding()` gives the provider and the sandbox id, so another host can resume the sandbox. `close()` stops the jobs of the backend. A session close closes it with `'moved'`, so an owned sandbox stays for the next host. `host.sessions.delete` closes it with `'closed'`, and with `owned: true` that destroys the sandbox. + +`commandEnv` and `createOutputTail` are exported from `@tanstack/ai-sandbox`, for providers that run commands as host processes. `@tanstack/ai-harness` is an optional peer dependency. diff --git a/.changeset/sandbox-local-process-env.md b/.changeset/sandbox-local-process-env.md index eb3957d401..ca17bde66c 100644 --- a/.changeset/sandbox-local-process-env.md +++ b/.changeset/sandbox-local-process-env.md @@ -2,4 +2,4 @@ '@tanstack/ai-sandbox-local-process': minor --- -A command in `localProcessSandbox` gets only a few variables of the host process: `PATH`, `HOME`, `LANG`, `LC_ALL`, `TMPDIR`, `TEMP`, `TMP`, and on Windows `SystemRoot`, `ComSpec`, `PATHEXT`, `WINDIR`, and `USERPROFILE`. The workspace secrets and the env of the call still reach it. Provider keys and other secrets of the server do not. If a command needs another variable of the host, set it as a workspace secret or in the env of the call. `scrubEnv` still removes a variable from the env of each command. +A command in `localProcessSandbox` gets only a few variables of the host process: `PATH`, `HOME`, `LANG`, `LC_ALL`, `TMPDIR`, `TEMP`, `TMP`, and on Windows `SystemRoot`, `ComSpec`, `PATHEXT`, `WINDIR`, and `USERPROFILE`. The workspace secrets and the env of the call still reach it. Provider keys and other secrets of the server do not. If a command needs another variable of the host, set it as a workspace secret or in the env of the call. `passEnv` names more variables of the host for each command, for example `HTTPS_PROXY` or `CLAUDE_CONFIG_DIR`. `scrubEnv` still removes a variable from the env of each command. diff --git a/docs/config.json b/docs/config.json index fc516be022..babe7aee31 100644 --- a/docs/config.json +++ b/docs/config.json @@ -1074,7 +1074,7 @@ "label": "Sandbox and extend the coding tools", "to": "harness/coding-tools-backends", "addedAt": "2026-10-13", - "updatedAt": "2026-10-18" + "updatedAt": "2026-10-19" }, { "label": "Run the coding tools for many threads", @@ -1199,7 +1199,7 @@ "label": "Providers", "to": "sandbox/providers", "addedAt": "2026-06-29", - "updatedAt": "2026-10-07" + "updatedAt": "2026-10-19" }, { "label": "Harnesses", @@ -1210,7 +1210,8 @@ { "label": "Harness Auth", "to": "sandbox/auth", - "addedAt": "2026-08-18" + "addedAt": "2026-08-18", + "updatedAt": "2026-10-19" }, { "label": "Workspace", diff --git a/docs/harness/coding-tools-backends.md b/docs/harness/coding-tools-backends.md index ca6b849f48..7b8f43e3de 100644 --- a/docs/harness/coding-tools-backends.md +++ b/docs/harness/coding-tools-backends.md @@ -105,7 +105,8 @@ For background jobs and for threads that move to another server: | Member | What it does | Without it | | --- | --- | --- | | `spawn(command, options)` | Starts a command in the background. Gives `{ wait, kill, output, jobId? }`. `jobId` resolves with the id of the process, for `killJob`. | `background: true` gives an error. | -| `killJob(jobId)` | Stops a job that `spawn` started, also on an earlier process or machine. | A job that a crash left running is not stopped. | +| `killJob(jobId)` | Stops a job that `spawn` started, also on an earlier process or machine. Resolves `false` when no such job runs. Rejects when it cannot check that the process is the job: then the log keeps the job. | A job that a crash left running is not stopped. | +| `exec` option `onStart(jobId)` | Gives the id of a foreground command when it starts, so a command that moves to the background can be stopped by a later host too. | A moved command has no id in the log. | | `binding()` | Gives `{ kind: 'sandbox', provider, sandboxId }` for a sandbox that another machine can resume. | The workspace is a folder on this host. | | `close({ reason })` | Releases the backend. `'closed'`: the session closed. `'moved'`: the session continues on another machine, so keep what it can resume. | Nothing to release. | diff --git a/docs/harness/coding-tools.md b/docs/harness/coding-tools.md index a00faa5463..224c1410fd 100644 --- a/docs/harness/coding-tools.md +++ b/docs/harness/coding-tools.md @@ -174,7 +174,7 @@ For `bash`: - The model gets the exit code and the last 1000 lines (20 KiB) of the output. - A command stops after 2 minutes, with exit code 124. Every process that it started stops too. Set another limit with `bashTimeoutMs`. The model can ask for up to 10 minutes. - With `spillDir`, long output goes to a file in that folder, and the model gets the path. `spillDir` is relative to `root`, or absolute. The file has `[redacted]` in place of each secret of the session, the same as the tool result. See [Keep tokens out of the log and the model](./auth#keep-tokens-out-of-the-log-and-the-model). -- A command gets only a few variables of your process, like `PATH` and `HOME`. Your provider keys stay out of it. See [Keep server secrets out of commands](./coding-workspaces#2-keep-server-secrets-out-of-commands). +- A command gets only a few variables of your process, like `PATH` and `HOME`. Your provider keys are not in its env. On the host this is not isolation: a command can still read the files of your server. See [Keep server secrets out of commands](./coding-workspaces#2-keep-server-secrets-out-of-commands). With `background: true`, the call returns at once with a job id. When the job ends, the model gets a note with the output, and a new turn starts. That turn runs as the user whose run started the job. diff --git a/docs/harness/coding-workspaces.md b/docs/harness/coding-workspaces.md index 7fb9eba82f..0a1a636582 100644 --- a/docs/harness/coding-workspaces.md +++ b/docs/harness/coding-workspaces.md @@ -23,10 +23,13 @@ Pass `workspace` to `workspaceTools()` in place of `root`. `bind` gets the run, ```ts group=harness-coding-workspaces import { mkdir } from 'node:fs/promises' -import { join } from 'node:path' import { defineHarness } from '@tanstack/ai-harness' import { permissions } from '@tanstack/ai-harness/plugins' -import { hostBackend, workspaceTools } from '@tanstack/ai-harness/plugins/coding' +import { + hostBackend, + workspaceRoot, + workspaceTools, +} from '@tanstack/ai-harness/plugins/coding' import { openaiText } from '@tanstack/ai-openai' export const coder = defineHarness({ @@ -37,7 +40,7 @@ export const coder = defineHarness({ workspaceTools({ workspace: { bind: async ({ run }) => { - const root = join('/srv/workspaces', run.tenantId ?? 'default', run.threadId) + const root = workspaceRoot('/srv/workspaces', run.tenantId ?? 'default', run.threadId) await mkdir(root, { recursive: true }) return { ...hostBackend, root } }, @@ -48,12 +51,13 @@ export const coder = defineHarness({ ``` - The tools call `bind` once for each thread, before the first turn or at the first tool call. If `bind` throws, the next turn calls it again. -- The `root` that `bind` gives is the boundary of every tool. A path outside it is refused. With `outside: 'ask'`, a rule that allows the path, or `bypass` mode, lets a call use it. No approval can cover it, because the root is not known before the call. The `cwd` thread setting must stay inside it. +- A client sends the thread id. Do not put it in a path with `join` or a template string: an id like `../other-tenant/t-1` leads into another folder. `workspaceRoot` encodes each part, so each id gets its own folder in `/srv/workspaces`, and it throws for a path that leaves it. +- The `root` that `bind` gives is the boundary of the file tools. A path outside it is refused. With `outside: 'ask'`, a rule that allows the path, or `bypass` mode, lets a call use it. No approval can cover it, because the root is not known before the call. The `cwd` thread setting must stay inside it. - The tools resolve every link with `realpath` before they read or write, while they hold the lock on the file. So a link in one thread cannot lead into the folder of another thread. - The session log keeps a `harness.workspace.bound` record with the place of the workspace. [Move a thread to another server](#5-move-a-thread-to-another-server) uses it. - Put `workspaceTools()` before [`snapshots()`](./snapshots) and [`formatter()`](./coding-tools-backends#format-each-file-after-a-write). Then they work in the workspace of each thread, and their `root` is optional. -For a sandbox per thread, create it in `bind`, and give the backend `owned: true`. Then the thread owns the sandbox: +A folder on the server is not isolation: `bash` runs as the user of your server (see [Keep server secrets out of commands](#2-keep-server-secrets-out-of-commands)). When many users share the server, give each thread a sandbox. Create it in `bind`, and give the backend `owned: true`. Then the thread owns the sandbox: ```ts group=harness-coding-workspaces import { sandboxWorkspaceBackend } from '@tanstack/ai-sandbox/harness' @@ -82,7 +86,14 @@ A command gets only a few variables of the server process: - On Windows also `SystemRoot`, `ComSpec`, `PATHEXT`, `WINDIR`, and `USERPROFILE`. - `AGENT=1`, so a tool can see that an agent runs it. -Provider keys, `DATABASE_URL`, and every other variable stay in the server. The same rule applies to `bash`, `rg`, `git`, and the formatters. Git also runs without the hooks and the `fsmonitor` of the project. +Provider keys, `DATABASE_URL`, and the other variables of the server are not in the env of a command. The same rule applies to `bash`, `rg`, `git`, and the formatters. Git also runs without the hooks and the `fsmonitor` of the project. + +With `hostBackend`, and with the local-process sandbox, this short env and the `root` are defense in depth. They are not isolation: + +- A command runs as the user of your server. `bash` can read every file that this user can read, for example the `.env` file of the server, or the folder of another thread. +- `bash` can also read the env of each process of that user, for example `/proc//environ` of the server. + +For users that you do not trust, use [a sandbox for each thread](#1-give-each-thread-its-own-workspace). To give commands a variable, add it to the `env` of the backend calls: @@ -111,15 +122,15 @@ export const withTokenTools = workspaceTools({ root: '/srv/repo', backend: withT Two agents, or two sessions, can work in the same folder. No edit is lost: -- A write holds a lock on the file. The lock is shared by every session on this server, and it is keyed by the root and the real path of the file. -- `write_file`, `edit_file`, and `patch` fail when the file changed after the tool read it. The model gets `src/app.ts changed since it was read. Read it again before you edit it.`, reads the file again, and tries again. +- A write holds a lock on the file. The lock is shared by every session on this server, and it is keyed by the sandbox (or this server) and the real path of the file. Two roots that lead to one folder share the lock. Two sandboxes with the same root do not. +- `write_file`, `edit_file`, and `patch` fail when the file changed after the tool read it in the same agent run. A subagent and its lead each check their own read. The model gets `src/app.ts changed since it was read. Read it again before you edit it.`, reads the file again, and tries again. - The check compares the size and the time of the last change. A sandbox gives no change time, so there only a change of size is seen. ## 4. Watch background jobs With `background: true`, `bash` returns at once with a job id like `bash-1`. Two limits keep jobs small: -- At most 8 background jobs run at a time in a session. One more gives the model a tool error. Set another limit with `maxJobs`. +- At most 8 background jobs run at a time in a session. One more gives the model a tool error. Set another limit with `maxJobs`. A running command that `session.background()` moves counts too, and at the limit it does not move. - Each job keeps the newest 1 MiB of its output. When a job ends, the model gets a note, and the wake turn runs as the user whose run started the job. @@ -212,25 +223,32 @@ import { memoryPersistence } from '@tanstack/ai-persistence' const host = createHarnessHost({ persistence: memoryPersistence(), - ownership: { hostId: process.env.HOSTNAME ?? 'web-1' }, + ownership: { hostId: process.env.POD_NAME ?? 'web-1' }, }) ``` - With `hostId`, a thread that uses a local folder is pinned to this server. Workers on other servers do not take it. - Another server that opens the thread refuses its tool calls with `The workspace of this thread is on the host web-1. Only that host can continue the thread.` -- Keep the id the same after a restart, for example the machine or pod name. +- Keep the id the same after a restart. On Kubernetes, use the pod name of a StatefulSet. Do not use `HOSTNAME` of a Deployment pod: it changes at each restart, and the threads of the old pod stay pinned to it. - Without `hostId`, threads are not pinned, and every server calls `bind` again. Use this when every server sees the same folder, for example on a shared disk. +When a server does not come back, unpin its threads with `host.unpin`. Its ownership must have expired first. Then the next server calls `reattach` with the old binding. Give a new workspace there, for example a fresh clone, to move the thread: + +```ts group=harness-coding-workspaces +await host.unpin({ harness: coder, threadId: 'fix-login' }) +``` + ### After a move - Background jobs do not move. At the next tool call of the thread, the new server stops each job that the old server left running. The model gets the note `Background job bash-1 was stopped: it ran on a host before the session moved to this one.` +- A job that the backend cannot check stays in the log, and the model gets `Background job bash-1 may still run: it ran on a host before this one, and this host cannot stop it.` For example, `hostBackend` needs `ps` to check a job. - `/undo`, `/redo`, and a revert work on the new server when it reaches the same snapshot repository. See [Undo a turn](./snapshots#on-more-than-one-server). - [Coding agents](./coding-agents) run only on the server that owns the log. So two servers never run the agents of one shared workspace at the same time. ## What you have now - A folder or a sandbox for each thread, with its root as the boundary of every tool. -- Commands that never see the secrets of your server. +- Commands that get only a short env, and a sandbox for users that you do not trust. - Agents that share a folder and never lose an edit. - Events for every background job, and a limit on how many run. - Threads that move to another server with their sandbox, or stay on the server that has their folder. diff --git a/docs/sandbox/auth.md b/docs/sandbox/auth.md index ddf0204be8..2a69352af6 100644 --- a/docs/sandbox/auth.md +++ b/docs/sandbox/auth.md @@ -53,7 +53,19 @@ const stream = chat({ `scrubEnv` removes keys from the env of each command, for example a key that a workspace secret sets. If the CLI sees `XAI_API_KEY`, it can prefer that key -over your login. The keys of the host process never reach the command. +over your login. The keys of the host process do not reach the command. + +The command gets only a few variables of the host process. If your login or +your network needs more, for example `CLAUDE_CONFIG_DIR`, `CODEX_HOME`, +`HTTPS_PROXY`, or `SSH_AUTH_SOCK`, name them in `passEnv`: + +```ts +import { localProcessSandbox } from '@tanstack/ai-sandbox-local-process' + +const provider = localProcessSandbox({ + passEnv: ['CODEX_HOME', 'HTTPS_PROXY'], +}) +``` ## API key @@ -101,8 +113,9 @@ const stream = chat({ | [Codex](../adapters/codex) | `CODEX_API_KEY` | `codex login` | | [ACP-Compatible](../adapters/acp-compatible) | set `authMethodId` (for Grok, `xai.api_key`) | skip ACP `authenticate` | -OpenCode still reads `OPENAI_API_KEY` from the process env. It has no `authMode` -flag. +OpenCode reads `OPENAI_API_KEY` from its env. It has no `authMode` flag. On +local-process, the key of the host process does not reach it: set the key as a +workspace secret or with `env.set`. ## Client and server @@ -163,4 +176,4 @@ export async function POST(request: Request) { ``` See [Harnesses](./harnesses) for which adapter to pick, and [Providers](./providers) -for `scrubEnv` on local-process. +for `scrubEnv` and `passEnv` on local-process. diff --git a/docs/sandbox/providers.md b/docs/sandbox/providers.md index 9faff0a499..821d9d27e3 100644 --- a/docs/sandbox/providers.md +++ b/docs/sandbox/providers.md @@ -91,7 +91,7 @@ const dev = localProcessSandbox() re-creates and re-bootstraps under the same identity. The snapshot step is skipped silently (see [Capabilities](#capabilities)). -### Host login vs API key (`scrubEnv`) +### Host login vs API key (`scrubEnv`, `passEnv`) The provider is where the agent runs, not how it signs in. The default `authMode` is `'api-key'`. Set `'host'` when the machine already has a CLI @@ -110,6 +110,18 @@ const hostLogin = localProcessSandbox({ }) ``` +A command gets only a few variables of the host process: `PATH`, `HOME`, +`LANG`, `LC_ALL`, `TMPDIR`, `TEMP`, and `TMP` (and a few more on Windows). +When a CLI login or the network needs more, name them in `passEnv`: + +```ts +import { localProcessSandbox } from '@tanstack/ai-sandbox-local-process' + +const behindProxy = localProcessSandbox({ + passEnv: ['CLAUDE_CONFIG_DIR', 'HTTPS_PROXY', 'NO_PROXY', 'NODE_EXTRA_CA_CERTS'], +}) +``` + If the same local-process sandbox runs on a CI machine, set `authMode: 'api-key'`. Then inject the key as a workspace secret. Isolated and cloud providers have no host CLI login. Use `authMode: 'api-key'` and From 468e72cac7825bcdd12472efd4e0815cd8031a6d Mon Sep 17 00:00:00 2001 From: Alem Tuzlak Date: Mon, 19 Oct 2026 14:30:24 +0200 Subject: [PATCH 13/17] refactor(ai-utils, ai-harness, ai-sandbox, ai-sandbox-local-process): one commandEnv and createOutputTail in @tanstack/ai-utils The three packages import both helpers from @tanstack/ai-utils, a new direct dependency of each. No more copies. commandEnv reads process.env only when there is a process, and createOutputTail needs no Buffer. --- .changeset/ai-sandbox-harness.md | 2 - .changeset/ai-utils-command-env.md | 5 ++ packages/ai-harness/package.json | 1 + .../src/first-party/coding/backend.ts | 46 +------------ .../ai-harness/src/first-party/coding/bash.ts | 2 +- .../src/first-party/coding/ring-buffer.ts | 65 ------------------ .../ai-harness/tests/coding-backend.test.ts | 27 -------- packages/ai-harness/tests/coding-jobs.test.ts | 31 --------- .../ai-sandbox-local-process/package.json | 3 + .../ai-sandbox-local-process/src/handle.ts | 2 +- packages/ai-sandbox/package.json | 3 +- packages/ai-sandbox/src/index.ts | 4 -- packages/ai-sandbox/src/workspace-backend.ts | 2 +- .../src/command.ts} | 37 +++++----- packages/ai-utils/src/index.ts | 1 + packages/ai-utils/tests/command.test.ts | 68 +++++++++++++++++++ pnpm-lock.yaml | 12 +++- 17 files changed, 115 insertions(+), 196 deletions(-) create mode 100644 .changeset/ai-utils-command-env.md delete mode 100644 packages/ai-harness/src/first-party/coding/ring-buffer.ts rename packages/{ai-sandbox/src/command-env.ts => ai-utils/src/command.ts} (78%) create mode 100644 packages/ai-utils/tests/command.test.ts diff --git a/.changeset/ai-sandbox-harness.md b/.changeset/ai-sandbox-harness.md index aed6253d07..a837603ec3 100644 --- a/.changeset/ai-sandbox-harness.md +++ b/.changeset/ai-sandbox-harness.md @@ -8,6 +8,4 @@ Add `@tanstack/ai-sandbox/harness` for `@tanstack/ai-harness`: - `codingAgents` runs agents only on the host that owns the log of the thread, so `'shared'` runs one agent at a time also with many hosts. - `sandboxWorkspaceBackend(sandbox, { owned? })` runs the coding tools of `@tanstack/ai-harness/plugins/coding` (`workspaceTools()`, `snapshots()`, `formatter()`) in a sandbox. It resolves links in the sandbox, reads one page of a file in the sandbox, and stops background jobs by their id (the pid and its start time), also from another host. It never stops another process that got the pid of a job. `binding()` gives the provider and the sandbox id, so another host can resume the sandbox. `close()` stops the jobs of the backend. A session close closes it with `'moved'`, so an owned sandbox stays for the next host. `host.sessions.delete` closes it with `'closed'`, and with `owned: true` that destroys the sandbox. -`commandEnv` and `createOutputTail` are exported from `@tanstack/ai-sandbox`, for providers that run commands as host processes. - `@tanstack/ai-harness` is an optional peer dependency. diff --git a/.changeset/ai-utils-command-env.md b/.changeset/ai-utils-command-env.md new file mode 100644 index 0000000000..86516cd423 --- /dev/null +++ b/.changeset/ai-utils-command-env.md @@ -0,0 +1,5 @@ +--- +'@tanstack/ai-utils': minor +--- + +Add `commandEnv(env, platform, passEnv?)`, the short env of a command that a process starts (no provider keys or other secrets of the server), and `createOutputTail(maxBytes)`, the newest output of a command. `@tanstack/ai-harness`, `@tanstack/ai-sandbox`, and `@tanstack/ai-sandbox-local-process` use them. diff --git a/packages/ai-harness/package.json b/packages/ai-harness/package.json index 19f5f00cf2..871933d675 100644 --- a/packages/ai-harness/package.json +++ b/packages/ai-harness/package.json @@ -85,6 +85,7 @@ "test:types": "tsc" }, "dependencies": { + "@tanstack/ai-utils": "workspace:^", "@tanstack/store": "^0.11.1" }, "devDependencies": { diff --git a/packages/ai-harness/src/first-party/coding/backend.ts b/packages/ai-harness/src/first-party/coding/backend.ts index 2659c42e13..39660b7535 100644 --- a/packages/ai-harness/src/first-party/coding/backend.ts +++ b/packages/ai-harness/src/first-party/coding/backend.ts @@ -1,8 +1,8 @@ import * as childProcess from 'node:child_process' import * as fs from 'node:fs/promises' import * as nodePath from 'node:path' +import { commandEnv, createOutputTail } from '@tanstack/ai-utils' import { isRecord } from '../../utils' -import { createOutputTail } from './ring-buffer' import type { RunContext } from '../../types' import type { WorkspaceHooks } from '../workspace-hooks' @@ -78,7 +78,7 @@ export interface WorkspaceBackend { * Run a shell command and wait for it to end. `env` is added to the * environment. A command that fails resolves with its exit code. On * {@link hostBackend}, a command gets only a few variables of this - * process (see {@link commandEnv}), plus `env`. `onStart` gets the id of + * process (see `commandEnv` of `@tanstack/ai-utils`), plus `env`. `onStart` gets the id of * the process when it started, like the `jobId` of `spawn`, for * {@link killJob}: the host can move a running command to the background. * A backend that has no id does not call it. @@ -248,48 +248,6 @@ export function workspaceRoot(base: string, ...parts: Array) { return full } -/** - * The variables of this process that a command gets. The others, like - * provider keys and `DATABASE_URL`, stay in the server. `@tanstack/ai-sandbox` - * has a copy of these lists and of {@link commandEnv} in `command-env.ts`, - * for its local-process provider. This package does not depend on that - * one. Change both. - */ -const KEPT_ENV = ['PATH', 'HOME', 'LANG', 'LC_ALL', 'TMPDIR', 'TEMP', 'TMP'] -/** On Windows, a command also gets these, so programs can start. */ -const KEPT_WINDOWS_ENV = [ - 'SystemRoot', - 'ComSpec', - 'PATHEXT', - 'WINDIR', - 'USERPROFILE', -] - -/** - * The environment of a command: the variables of this process in a short - * list (`PATH`, `HOME`, `LANG`, `LC_ALL`, `TMPDIR`, `TEMP`, `TMP`, and on - * Windows `SystemRoot`, `ComSpec`, `PATHEXT`, `WINDIR`, `USERPROFILE`), plus - * `env`. No other variable of this process goes to the command. On Windows - * it sets `NoDefaultCurrentDirectoryInExePath`. Then `cmd.exe` takes a - * program name like `rg` from the PATH only, and not from the current - * folder, where a cloned repo can put an `rg.cmd`. - */ -export function commandEnv( - env: Record | undefined, - platform: NodeJS.Platform, -) { - const isWindows = platform === 'win32' - const names = isWindows ? [...KEPT_ENV, ...KEPT_WINDOWS_ENV] : KEPT_ENV - const kept: Record = {} - for (const name of names) { - // On Windows, `process.env` finds a name in any letter case. - const value = process.env[name] - if (value !== undefined) kept[name] = value - } - if (!isWindows) return { ...kept, ...env } - return { ...kept, ...env, NoDefaultCurrentDirectoryInExePath: '1' } -} - /** True for the error of a path that is not there. */ const isMissing = (error: unknown) => error instanceof Error && 'code' in error && error.code === 'ENOENT' diff --git a/packages/ai-harness/src/first-party/coding/bash.ts b/packages/ai-harness/src/first-party/coding/bash.ts index c9890a76df..1fdb589483 100644 --- a/packages/ai-harness/src/first-party/coding/bash.ts +++ b/packages/ai-harness/src/first-party/coding/bash.ts @@ -1,9 +1,9 @@ import { toolDefinition } from '@tanstack/ai' +import { createOutputTail } from '@tanstack/ai-utils' import { isRecord } from '../../utils' import { boundText } from '../bound-output' import { pathsOf, runOf, stringArg } from './backend' import { JobEnded, JobOutput, JobStarted } from './events' -import { createOutputTail } from './ring-buffer' import type { PluginEvent } from '../../extensions' import type { BackgroundJobsApi } from '../../plugins' import type { RunContext } from '../../types' diff --git a/packages/ai-harness/src/first-party/coding/ring-buffer.ts b/packages/ai-harness/src/first-party/coding/ring-buffer.ts deleted file mode 100644 index f8c8469bc7..0000000000 --- a/packages/ai-harness/src/first-party/coding/ring-buffer.ts +++ /dev/null @@ -1,65 +0,0 @@ -const encoder = new TextEncoder() -const decoder = new TextDecoder() - -/** True for a byte in the middle of a UTF-8 character (`10xxxxxx`). */ -const isContinuation = (byte: number) => (byte & 0xc0) === 0x80 - -/** - * The newest output of a command, at most `maxBytes` bytes. When more comes, - * the oldest bytes are dropped. `text()` gives the output as UTF-8 text, and - * starts with a line that tells how many bytes were dropped. - * - * `@tanstack/ai-sandbox` has a copy of this function in `command-env.ts`. - * This package does not depend on that one. Change both. - * - * @example - * ```ts - * const output = createOutputTail(1024 * 1024) - * child.stdout.on('data', output.push) - * output.text() // '[12 earlier bytes were dropped.]\n...' - * ``` - */ -export function createOutputTail(maxBytes: number) { - // The chunks from `start` on are the output. The ones before are dropped. - const chunks: Array = [] - let start = 0 - let size = 0 - let dropped = 0 - - const push = (chunk: Uint8Array | string) => { - const bytes = typeof chunk === 'string' ? encoder.encode(chunk) : chunk - // A copy of the end only, so a big chunk is not kept whole. - const kept = bytes.byteLength > maxBytes ? bytes.slice(-maxBytes) : bytes - dropped += bytes.byteLength - kept.byteLength - chunks.push(kept) - size += kept.byteLength - while (size > maxBytes) { - const first = chunks[start] - if (first === undefined) break - const over = size - maxBytes - const cut = Math.min(over, first.byteLength) - if (cut === first.byteLength) start += 1 - else chunks[start] = first.subarray(cut) - size -= cut - dropped += cut - } - // Remove the dropped chunks from the array now and then, not each time. - if (start > 1024 && start * 2 > chunks.length) { - chunks.splice(0, start) - start = 0 - } - } - - const text = () => { - let bytes = Buffer.concat(chunks.slice(start)) - if (dropped === 0) return decoder.decode(bytes) - // The cut can be in a character. Start at the next whole one. - let from = 0 - while (from < 3 && from < bytes.length && isContinuation(bytes[from] ?? 0)) - from += 1 - bytes = bytes.subarray(from) - return `[${dropped + from} earlier bytes were dropped.]\n${decoder.decode(bytes)}` - } - - return { push, text } -} diff --git a/packages/ai-harness/tests/coding-backend.test.ts b/packages/ai-harness/tests/coding-backend.test.ts index 7b478adab7..c3f16999fd 100644 --- a/packages/ai-harness/tests/coding-backend.test.ts +++ b/packages/ai-harness/tests/coding-backend.test.ts @@ -10,7 +10,6 @@ import { workspaceRoot, workspaceTools, } from '../src/first-party/coding' -import { commandEnv } from '../src/first-party/coding/backend' import { mockAdapter, text, toolCall } from './helpers' let dir = '' @@ -179,32 +178,6 @@ describe('hostBackend', () => { expect(job.output()).toBe('unset\n') }) }) - - it('gives a command only the listed variables of this process, plus env', async () => { - const server = { - SECRET_TOKEN: 'hunter2', - LANG: 'C.UTF-8', - SystemRoot: 'C:\\Windows', - } - await withEnv(server, async () => { - const linux = commandEnv({ AGENT: '1' }, 'linux') - expect(linux).toMatchObject({ LANG: 'C.UTF-8', AGENT: '1' }) - expect(linux).not.toHaveProperty('SECRET_TOKEN') - // The Windows variables and setting are for Windows only. - expect(linux).not.toHaveProperty('SystemRoot') - expect(linux).not.toHaveProperty('NoDefaultCurrentDirectoryInExePath') - const windows = commandEnv( - { AGENT: '1', NoDefaultCurrentDirectoryInExePath: '' }, - 'win32', - ) - expect(windows).toMatchObject({ - SystemRoot: 'C:\\Windows', - AGENT: '1', - NoDefaultCurrentDirectoryInExePath: '1', - }) - expect(windows).not.toHaveProperty('SECRET_TOKEN') - }) - }) }) describe('workspaceRoot', () => { diff --git a/packages/ai-harness/tests/coding-jobs.test.ts b/packages/ai-harness/tests/coding-jobs.test.ts index bc88266b94..707ce98f83 100644 --- a/packages/ai-harness/tests/coding-jobs.test.ts +++ b/packages/ai-harness/tests/coding-jobs.test.ts @@ -4,7 +4,6 @@ import { resolve } from 'node:path' import { afterEach, describe, expect, it, vi } from 'vitest' import { hostBackend } from '../src/first-party/coding/backend' import { bashTools } from '../src/first-party/coding/bash' -import { createOutputTail } from '../src/first-party/coding/ring-buffer' import { createWorkspaceTools } from '../src/first-party/coding/workspace' import type { ToolExecutionContext } from '@tanstack/ai' import type { PluginEvent } from '../src/extensions' @@ -107,36 +106,6 @@ async function startJobs(bash: ReturnType, count: number) { /** Wait until the promises that are ready now have run. */ const settle = () => new Promise((done) => setTimeout(done, 0)) -describe('createOutputTail', () => { - it('gives all output under the limit, with no note', () => { - const output = createOutputTail(10) - output.push('ab') - output.push(new TextEncoder().encode('cd')) - expect(output.text()).toBe('abcd') - }) - - it('keeps the newest bytes over the limit, and tells how many were dropped', () => { - const output = createOutputTail(4) - output.push('abc') - output.push('def') - expect(output.text()).toBe('[2 earlier bytes were dropped.]\ncdef') - }) - - it('keeps the end of one chunk that is bigger than the limit', () => { - const output = createOutputTail(3) - output.push('abcdef') - expect(output.text()).toBe('[3 earlier bytes were dropped.]\ndef') - }) - - it('starts at a whole character when the cut is in one', () => { - const output = createOutputTail(4) - // `é` is 2 bytes. The cut drops its first byte, so it drops all of it. - output.push('é') - output.push('abc') - expect(output.text()).toBe('[2 earlier bytes were dropped.]\nabc') - }) -}) - describe('bash job limit', () => { it('refuses a 9th background job, and names the limit', async () => { const { env } = workspace() diff --git a/packages/ai-sandbox-local-process/package.json b/packages/ai-sandbox-local-process/package.json index 58b08f5e15..f47989a778 100644 --- a/packages/ai-sandbox-local-process/package.json +++ b/packages/ai-sandbox-local-process/package.json @@ -41,6 +41,9 @@ "test:lib:dev": "pnpm test:lib --watch", "test:types": "tsc" }, + "dependencies": { + "@tanstack/ai-utils": "workspace:^" + }, "peerDependencies": { "@tanstack/ai-sandbox": "workspace:^" }, diff --git a/packages/ai-sandbox-local-process/src/handle.ts b/packages/ai-sandbox-local-process/src/handle.ts index 2b29a3bae4..31b2f6144d 100644 --- a/packages/ai-sandbox-local-process/src/handle.ts +++ b/packages/ai-sandbox-local-process/src/handle.ts @@ -16,9 +16,9 @@ import * as path from 'node:path' import { DEFAULT_WORKSPACE_ROOT, UnsupportedCapabilityError, - commandEnv, createExecBackedGit, } from '@tanstack/ai-sandbox' +import { commandEnv } from '@tanstack/ai-utils' import type { ChildProcess } from 'node:child_process' import type { Readable } from 'node:stream' import type { diff --git a/packages/ai-sandbox/package.json b/packages/ai-sandbox/package.json index 0dea764a2c..9514066ef6 100644 --- a/packages/ai-sandbox/package.json +++ b/packages/ai-sandbox/package.json @@ -83,7 +83,8 @@ }, "dependencies": { "@modelcontextprotocol/sdk": "^1.29.0", - "@tanstack/ai-skills": "workspace:^" + "@tanstack/ai-skills": "workspace:^", + "@tanstack/ai-utils": "workspace:^" }, "peerDependencies": { "@ngrok/ngrok": "^1.0.0", diff --git a/packages/ai-sandbox/src/index.ts b/packages/ai-sandbox/src/index.ts index b2172c6ad4..e265eba978 100644 --- a/packages/ai-sandbox/src/index.ts +++ b/packages/ai-sandbox/src/index.ts @@ -457,10 +457,6 @@ export type { SandboxWatchHandle, } from './watch' -// The env and the output tail of a command on this machine, for providers -// that run commands as host processes. -export { commandEnv, createOutputTail } from './command-env' - // Keying export { computeSandboxKey, computeWorkspaceHash } from './key' export type { SandboxKeyInput } from './key' diff --git a/packages/ai-sandbox/src/workspace-backend.ts b/packages/ai-sandbox/src/workspace-backend.ts index 1fcf6cfea2..0587929cd1 100644 --- a/packages/ai-sandbox/src/workspace-backend.ts +++ b/packages/ai-sandbox/src/workspace-backend.ts @@ -1,5 +1,5 @@ import { posix } from 'node:path' -import { createOutputTail } from './command-env' +import { createOutputTail } from '@tanstack/ai-utils' import { shellQuote } from './journal' import type { WorkspaceBackend } from '@tanstack/ai-harness/plugins/coding' import type { ExecResult, SandboxFs, SandboxHandle } from './contracts' diff --git a/packages/ai-sandbox/src/command-env.ts b/packages/ai-utils/src/command.ts similarity index 78% rename from packages/ai-sandbox/src/command-env.ts rename to packages/ai-utils/src/command.ts index a0c45623f2..b70bbf91c3 100644 --- a/packages/ai-sandbox/src/command-env.ts +++ b/packages/ai-utils/src/command.ts @@ -13,38 +13,37 @@ const KEPT_WINDOWS_ENV = [ ] /** - * The environment of a command on this machine: the variables of this - * process in a short list (`PATH`, `HOME`, `LANG`, `LC_ALL`, `TMPDIR`, - * `TEMP`, `TMP`, and on Windows `SystemRoot`, `ComSpec`, `PATHEXT`, - * `WINDIR`, `USERPROFILE`), plus the variables named in `passEnv`, plus - * `env`. No other variable of this process goes to the command. On Windows - * it sets `NoDefaultCurrentDirectoryInExePath`, so a program name is found - * on the PATH only, and not in the current folder. - * - * `@tanstack/ai-harness` keeps a copy for its `hostBackend` - * (`src/first-party/coding/backend.ts`), because it does not depend on this - * package. Change both. + * The environment of a command that this process starts: the variables of + * this process in a short list (`PATH`, `HOME`, `LANG`, `LC_ALL`, `TMPDIR`, + * `TEMP`, `TMP`, and on Windows `SystemRoot`, `ComSpec`, `PATHEXT`, `WINDIR`, + * `USERPROFILE`), plus the variables named in `passEnv`, plus `env`. No + * other variable of this process goes to the command. On Windows it sets + * `NoDefaultCurrentDirectoryInExePath`, so `cmd.exe` takes a program name + * from the PATH only, and not from the current folder, where a cloned repo + * can put an `rg.cmd`. Without `process` (a browser), only `env`. * * @example * ```ts - * spawn('sh', ['-c', command], { + * spawn(command, { + * shell: true, * env: commandEnv({ AGENT: '1' }, process.platform, ['HTTPS_PROXY']), * }) * ``` */ export function commandEnv( env: Record | undefined, - platform: NodeJS.Platform, + platform: string, passEnv: ReadonlyArray = [], ) { const isWindows = platform === 'win32' const names = isWindows ? [...KEPT_ENV, ...KEPT_WINDOWS_ENV, ...passEnv] : [...KEPT_ENV, ...passEnv] + const own = typeof process !== 'undefined' ? process.env : {} const kept: Record = {} for (const name of names) { // On Windows, `process.env` finds a name in any letter case. - const value = process.env[name] + const value = own[name] if (value !== undefined) kept[name] = value } if (!isWindows) return { ...kept, ...env } @@ -62,9 +61,6 @@ const isContinuation = (byte: number) => (byte & 0xc0) === 0x80 * the oldest bytes are dropped. `text()` gives the output as UTF-8 text, and * starts with a line that tells how many bytes were dropped. * - * `@tanstack/ai-harness` keeps a copy (`src/first-party/coding/ring-buffer.ts`), - * because it does not depend on this package. Change both. - * * @example * ```ts * const output = createOutputTail(1024 * 1024) @@ -104,7 +100,12 @@ export function createOutputTail(maxBytes: number) { } const text = () => { - let bytes = Buffer.concat(chunks.slice(start)) + let bytes = new Uint8Array(size) + let at = 0 + for (const chunk of chunks.slice(start)) { + bytes.set(chunk, at) + at += chunk.byteLength + } if (dropped === 0) return decoder.decode(bytes) // The cut can be in a character. Start at the next whole one. let from = 0 diff --git a/packages/ai-utils/src/index.ts b/packages/ai-utils/src/index.ts index abedd7a63d..06799eb649 100644 --- a/packages/ai-utils/src/index.ts +++ b/packages/ai-utils/src/index.ts @@ -1,5 +1,6 @@ export { generateId } from './id' export { getApiKeyFromEnv } from './env' +export { commandEnv, createOutputTail } from './command' export { transformNullsToUndefined, undoNullWidening } from './transforms' export type { NullWideningMap } from './transforms' export { diff --git a/packages/ai-utils/tests/command.test.ts b/packages/ai-utils/tests/command.test.ts new file mode 100644 index 0000000000..c74530bb4a --- /dev/null +++ b/packages/ai-utils/tests/command.test.ts @@ -0,0 +1,68 @@ +import { afterEach, describe, expect, it, vi } from 'vitest' +import { commandEnv, createOutputTail } from '../src/command' + +describe('commandEnv', () => { + afterEach(() => { + vi.unstubAllEnvs() + }) + + it('gives a command only the listed variables of this process, plus env', () => { + vi.stubEnv('SECRET_TOKEN', 'hunter2') + vi.stubEnv('LANG', 'C.UTF-8') + vi.stubEnv('SystemRoot', 'C:\\Windows') + vi.stubEnv('HTTPS_PROXY', 'http://proxy:8080') + + const linux = commandEnv({ AGENT: '1' }, 'linux') + expect(linux).toMatchObject({ LANG: 'C.UTF-8', AGENT: '1' }) + expect(linux).not.toHaveProperty('SECRET_TOKEN') + // The Windows variables and setting are for Windows only. + expect(linux).not.toHaveProperty('SystemRoot') + expect(linux).not.toHaveProperty('NoDefaultCurrentDirectoryInExePath') + expect(linux).not.toHaveProperty('HTTPS_PROXY') + + expect(commandEnv(undefined, 'linux', ['HTTPS_PROXY'])).toMatchObject({ + HTTPS_PROXY: 'http://proxy:8080', + }) + + const windows = commandEnv( + { AGENT: '1', NoDefaultCurrentDirectoryInExePath: '' }, + 'win32', + ) + expect(windows).toMatchObject({ + SystemRoot: 'C:\\Windows', + AGENT: '1', + NoDefaultCurrentDirectoryInExePath: '1', + }) + expect(windows).not.toHaveProperty('SECRET_TOKEN') + }) +}) + +describe('createOutputTail', () => { + it('gives all output under the limit, with no note', () => { + const output = createOutputTail(10) + output.push('ab') + output.push(new TextEncoder().encode('cd')) + expect(output.text()).toBe('abcd') + }) + + it('keeps the newest bytes over the limit, and tells how many were dropped', () => { + const output = createOutputTail(4) + output.push('abc') + output.push('def') + expect(output.text()).toBe('[2 earlier bytes were dropped.]\ncdef') + }) + + it('keeps the end of one chunk that is bigger than the limit', () => { + const output = createOutputTail(3) + output.push('abcdef') + expect(output.text()).toBe('[3 earlier bytes were dropped.]\ndef') + }) + + it('starts at a whole character when the cut is in one', () => { + const output = createOutputTail(4) + // `é` is 2 bytes. The cut drops its first byte, so it drops all of it. + output.push('é') + output.push('abc') + expect(output.text()).toBe('[2 earlier bytes were dropped.]\nabc') + }) +}) diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index 34292dadc1..7b7cfea333 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -2641,6 +2641,9 @@ importers: packages/ai-harness: dependencies: + '@tanstack/ai-utils': + specifier: workspace:^ + version: link:../ai-utils '@tanstack/store': specifier: ^0.11.1 version: 0.11.2 @@ -3218,6 +3221,9 @@ importers: '@tanstack/ai-skills': specifier: workspace:^ version: link:../ai-skills + '@tanstack/ai-utils': + specifier: workspace:^ + version: link:../ai-utils devDependencies: '@ngrok/ngrok': specifier: ^1.7.0 @@ -3335,6 +3341,10 @@ importers: version: 4.1.10(vitest@4.1.11) packages/ai-sandbox-local-process: + dependencies: + '@tanstack/ai-utils': + specifier: workspace:^ + version: link:../ai-utils devDependencies: '@tanstack/ai': specifier: workspace:* @@ -3570,7 +3580,7 @@ importers: devDependencies: '@google/genai': specifier: ^2.10.0 - version: 2.10.0(@modelcontextprotocol/sdk@1.29.0(@cfworker/json-schema@4.1.1)(zod@4.3.6)) + version: 2.10.0(@modelcontextprotocol/sdk@1.29.0(@cfworker/json-schema@4.1.1)(zod@4.6.5)) '@tanstack/ai': specifier: workspace:* version: link:../ai From f3507198127fc0a7f4e38b0fda2a34841d81d0bc Mon Sep 17 00:00:00 2001 From: Alem Tuzlak Date: Mon, 19 Oct 2026 14:30:25 +0200 Subject: [PATCH 14/17] feat(ai-persistence, ai-harness): host heartbeats, so a pin lapses only when its host stopped - LogStore.heartbeat({ hostId, ttlMs }). acquire and listOrphaned ignore the affinity of a host with no live heartbeat. The memory store, the conformance suite, and the SQL recipe have it. - A host with ownership.hostId sends a heartbeat each renewMs, also when it owns no log. A live host keeps its pins on idle threads. When it stops, another host reattaches its threads. - Docs for the workspaces, deploy, durable sessions, and store pages. --- .changeset/ai-harness-coding-tools.md | 2 +- .changeset/ai-persistence-host-heartbeat.md | 5 ++ docs/config.json | 2 +- docs/harness/coding-workspaces.md | 6 +- docs/harness/deploy.md | 2 +- docs/harness/durable-sessions.md | 2 +- docs/persistence/store-reference.md | 30 +++++++- packages/ai-harness/src/host.ts | 4 +- packages/ai-harness/src/ownership.ts | 30 ++++++-- packages/ai-harness/tests/helpers.ts | 1 + .../tests/workspace-binding.test.ts | 74 +++++++++++++++++++ packages/ai-persistence/src/memory.ts | 31 ++++++-- .../ai-persistence/src/testkit/conformance.ts | 32 ++++++++ packages/ai-persistence/src/types.ts | 14 +++- 14 files changed, 214 insertions(+), 21 deletions(-) create mode 100644 .changeset/ai-persistence-host-heartbeat.md diff --git a/.changeset/ai-harness-coding-tools.md b/.changeset/ai-harness-coding-tools.md index 18912cfffd..8a999a64d9 100644 --- a/.changeset/ai-harness-coding-tools.md +++ b/.changeset/ai-harness-coding-tools.md @@ -6,7 +6,7 @@ Add coding tools in the new Node entry `@tanstack/ai-harness/plugins/coding`: - `workspaceTools({ root })` gives `read_file` (text pages, images, PDFs), `write_file`, `edit_file` or `patch` (`editStyle`), `list_files`, `grep`, `bash` (with `background`), `webfetch`, and `websearch` (with a `SearchProvider`). With `permissions()`, edits, `bash`, and `webfetch` ask first, and the rules also apply to the real path of a link. `outside: 'ask'` makes a path outside `root` need an approval. `permissions()` decides outside paths and the real paths of links before the call runs, so a tool never asks while it runs. The `cwd` thread setting gives each thread its own working folder. - `workspaceTools({ workspace: { bind, reattach? } })` gives each thread its own workspace. The root of the workspace is the boundary of every tool. The session log keeps where the workspace is (`harness.workspace.bound`). On another host, `reattach` gives the workspace back (for example a resumed sandbox), or `'pinned'` or `'gone'`. A gone workspace sends `tanstack/workspace-tools:stranded`. With this plugin first, `snapshots()` and `formatter()` use the workspace of each thread, and their `root` is optional. -- `createHarnessHost({ ownership: { hostId } })` gives a host a stable id. Then a thread that uses a folder on that host is pinned to it, and other hosts refuse its tool calls. Without `hostId`, threads are not pinned. `host.unpin({ harness, threadId })` unpins a thread from a host that does not come back, once its ownership expired. +- `createHarnessHost({ ownership: { hostId } })` gives a host a stable id. Then a thread that uses a folder on that host is pinned to it, and other hosts refuse its tool calls. Without `hostId`, threads are not pinned. A host with `hostId` sends a heartbeat to the log store, and its pins lapse when no heartbeat came for `ttlMs`: then another host reattaches its threads. `host.unpin({ harness, threadId })` unpins a thread from a host that does not come back, once its ownership expired. - `workspaceRoot(base, ...parts)` gives a folder in `base` for a tenant and a thread id. It encodes each part, so a thread id from a client cannot leave `base` or reach another thread. - Commands, `rg`, `git`, and the formatters get only a short list of variables of the server process, plus the `env` of the backend call. Git runs without the hooks and the `fsmonitor` of the project. - The tools resolve links with `realpath` under the file lock, and fail when a backend cannot resolve a path. A write fails when the file changed after the tool read it in the same chat run, so a lead and its subagents do not lose each other's edits. The file locks are shared by the sessions of a host, by the sandbox (or host) and the real path of the file. diff --git a/.changeset/ai-persistence-host-heartbeat.md b/.changeset/ai-persistence-host-heartbeat.md new file mode 100644 index 0000000000..a997320915 --- /dev/null +++ b/.changeset/ai-persistence-host-heartbeat.md @@ -0,0 +1,5 @@ +--- +'@tanstack/ai-persistence': minor +--- + +`LogStore` gets `heartbeat({ hostId, ttlMs })`. A pin (`setAffinity`) now counts only while its host is alive: `acquire` and `listOrphaned` ignore the affinity of a host with no live heartbeat, so the logs of a host that stopped go to other hosts. A host that is alive keeps its pins, also on the logs that it released. A custom store must add `heartbeat` and this rule. `memoryLogStore()` and the conformance suite have both. diff --git a/docs/config.json b/docs/config.json index babe7aee31..cf87ed8cac 100644 --- a/docs/config.json +++ b/docs/config.json @@ -980,7 +980,7 @@ "label": "Durable sessions", "to": "harness/durable-sessions", "addedAt": "2026-10-13", - "updatedAt": "2026-10-18" + "updatedAt": "2026-10-19" }, { "label": "Run background work on workers", diff --git a/docs/harness/coding-workspaces.md b/docs/harness/coding-workspaces.md index 0a1a636582..671915fab3 100644 --- a/docs/harness/coding-workspaces.md +++ b/docs/harness/coding-workspaces.md @@ -227,12 +227,14 @@ const host = createHarnessHost({ }) ``` -- With `hostId`, a thread that uses a local folder is pinned to this server. Workers on other servers do not take it. +- With `hostId`, a thread that uses a local folder is pinned to this server. Workers on other servers do not take it, also while the thread is idle. - Another server that opens the thread refuses its tool calls with `The workspace of this thread is on the host web-1. Only that host can continue the thread.` - Keep the id the same after a restart. On Kubernetes, use the pod name of a StatefulSet. Do not use `HOSTNAME` of a Deployment pod: it changes at each restart, and the threads of the old pod stay pinned to it. - Without `hostId`, threads are not pinned, and every server calls `bind` again. Use this when every server sees the same folder, for example on a shared disk. -When a server does not come back, unpin its threads with `host.unpin`. Its ownership must have expired first. Then the next server calls `reattach` with the old binding. Give a new workspace there, for example a fresh clone, to move the thread: +The server sends a heartbeat to the log store each `ownership.renewMs`. When no heartbeat came for `ownership.ttlMs` (default 30 seconds), its pins do not count. Then the next server that gets work for the thread calls `reattach` with the old binding. Give a new workspace there, for example a fresh clone, to move the thread. Keep a restart shorter than `ttlMs`, or raise `ttlMs`. + +To move the threads of a server at once, unpin them with `host.unpin`. Its ownership must have expired first: ```ts group=harness-coding-workspaces await host.unpin({ harness: coder, threadId: 'fix-login' }) diff --git a/docs/harness/deploy.md b/docs/harness/deploy.md index b5793fcf9b..c5de208505 100644 --- a/docs/harness/deploy.md +++ b/docs/harness/deploy.md @@ -29,7 +29,7 @@ Give every server the same log store. Then any server can take a request for any - A [worker](./workers) runs the threads that no server has open: submitted inputs, schedules, and the work that a deploy cut. - Set the same `media: { secret }` on the handler of each server, so each server checks the media URLs of the others. See [Serve the session over HTTP](./connect#serve-the-session-over-http). - Stop each server with `host.close({ recoverable: true })`. A worker takes its cut work at once. See [Stop a host for a deploy](./durable-sessions#stop-a-host-for-a-deploy). -- If your agents use the [coding tools](./coding-tools) on the disk of a server, set `ownership: { hostId }` on each host. Then a thread stays on the server that has its folder. See [Move a thread to another server](./coding-workspaces#5-move-a-thread-to-another-server). +- If your agents use the [coding tools](./coding-tools) on the disk of a server, set `ownership: { hostId }` on each host. Then a thread stays on the server that has its folder, while that server sends heartbeats. Use an id that stays the same after a restart, like the pod name of a StatefulSet, and keep a restart shorter than `ownership.ttlMs`. See [Move a thread to another server](./coding-workspaces#5-move-a-thread-to-another-server). ## Build a worker artifact diff --git a/docs/harness/durable-sessions.md b/docs/harness/durable-sessions.md index 7e0295b9fb..593ed23e6e 100644 --- a/docs/harness/durable-sessions.md +++ b/docs/harness/durable-sessions.md @@ -293,7 +293,7 @@ const ownedHost = createHarnessHost({ - `ttlMs`: how long the ownership lasts after each renewal. The default is 30 000. After it, another host can take the thread. - `renewMs`: how often the timer renews. The default is 10 000. When you set only `ttlMs`, the default is a third of `ttlMs`. It must be less than `ttlMs`, else `createHarnessHost` throws a `RangeError`. -- `hostId`: a stable id for this host, for example the machine or pod name. With it, a thread whose coding tools use a folder on this host stays on this host. Other hosts refuse its tool calls. See [A folder on one server](./coding-workspaces#a-folder-on-one-server). +- `hostId`: a stable id for this host, for example the machine or pod name. With it, a thread whose coding tools use a folder on this host stays on this host. Other hosts refuse its tool calls. The host sends a heartbeat to the log store each `renewMs`. When no heartbeat came for `ttlMs`, other hosts can take its threads. See [A folder on one server](./coding-workspaces#a-folder-on-one-server). An input that waits for the owner keeps its JSON overrides (`model`, `modelOptions`, `reasoning`, and `promptCache`). The `adapter` and `tools` overrides are objects, so the owner runs the input without them. diff --git a/docs/persistence/store-reference.md b/docs/persistence/store-reference.md index e017a2fa66..f72ea4ce51 100644 --- a/docs/persistence/store-reference.md +++ b/docs/persistence/store-reference.md @@ -690,6 +690,7 @@ interface LogStore { epoch: number affinity: string | null }): Promise + heartbeat(beat: { hostId: string; ttlMs: number }): Promise listOrphaned(query: { limit: number harnesses?: ReadonlyArray @@ -799,7 +800,8 @@ or a `wakeInMs` less than 0, rejects with `RangeError`. - It returns `null` when another owner holds a live record: not released, and not expired. -- It returns `null` when the log has an affinity that is not `claim.affinity`. +- It returns `null` when the log has an affinity that is not `claim.affinity`, and + the host of that affinity is alive. - For the owner of a live record, it extends the expiry and returns the same epoch. - Else it adds 1 to the epoch, stores `harness`, clears the wake time, and returns `{ epoch }`. @@ -817,6 +819,10 @@ For the other ownership methods: does nothing. - `setAffinity` pins the log to the workers with that affinity, or unpins it with `null`. It is fenced. +- `heartbeat` marks the host `hostId` alive for `ttlMs` from now. `hostId` is the + affinity of its pins. A pin counts only while its host is alive, so the logs of a + host that stopped go to other hosts. A host that is alive keeps its pins, also on + the logs that it released. `listOrphaned` finds the logs that have work and no live owner. A worker sweep ([`createHarnessWorker`](../harness/workers) or `host.resumePending`) calls it, then @@ -830,7 +836,8 @@ applies: Then: - With `harnesses`, the store lists only the logs with one of those harness names. -- A log with an affinity is listed only when the query has the same `affinity`. +- A log with an affinity is listed only when the query has the same `affinity`, or + when the host of the affinity is not alive. - The log that is due first comes first, and the store lists at most `limit` logs. `limit: 0` returns `[]`. @@ -903,6 +910,11 @@ CREATE TABLE log_owner ( wake_at TIMESTAMPTZ ); +CREATE TABLE log_hosts ( + host_id TEXT PRIMARY KEY, + expires_at TIMESTAMPTZ NOT NULL -- the heartbeat expires +); + CREATE TABLE log_offers ( id BIGSERIAL PRIMARY KEY, -- the offer order log_id TEXT NOT NULL, @@ -945,6 +957,20 @@ For the other methods: only when the rules for `acquire` allow it. - Use the database clock (`now()`) for each expiry, wake time, and `offeredAt`. The clocks of the hosts can disagree. +- `heartbeat` is one upsert into `log_hosts`. In `acquire` and `listOrphaned`, an + affinity counts only when its row has `expires_at > now()`: + +```sql +INSERT INTO log_hosts (host_id, expires_at) +VALUES (:host_id, now() + :ttl_ms * interval '1 millisecond') +ON CONFLICT (host_id) DO UPDATE SET expires_at = EXCLUDED.expires_at; + +-- The pin of a log that `acquire` and `listOrphaned` respect. +SELECT o.affinity FROM log_owner o +JOIN log_hosts h ON h.host_id = o.affinity AND h.expires_at > now() +WHERE o.log_id = :log_id; +``` + - Keep the snapshot in one more table, with one row per log. - `deleteLog` deletes the rows of the log in the other tables, and resets its `log_owner` row: `head = 0`, `first_seq = 1`, `released = TRUE`, and no `wake_at` diff --git a/packages/ai-harness/src/host.ts b/packages/ai-harness/src/host.ts index fb37ad3b58..ad986ffd5d 100644 --- a/packages/ai-harness/src/host.ts +++ b/packages/ai-harness/src/host.ts @@ -176,7 +176,9 @@ export interface HarnessHostOptions

{ * thread that the host owns (default 10000 ms). `hostId`: a stable id * for this host, for example the machine or pod name. With it, a thread * that uses a local folder is pinned to this host, and another host - * refuses the thread. Leave it out for the defaults: no pins. + * refuses the thread. The host sends a heartbeat each `renewMs`. When no + * heartbeat came for `ttlMs`, the pin does not count, and another host + * can reattach the thread. Leave it out for the defaults: no pins. */ ownership?: OwnershipOptions /** How the sessions of this host live. Leave it out for the defaults. */ diff --git a/packages/ai-harness/src/ownership.ts b/packages/ai-harness/src/ownership.ts index 56d0e4b713..8c86e91ead 100644 --- a/packages/ai-harness/src/ownership.ts +++ b/packages/ai-harness/src/ownership.ts @@ -21,6 +21,10 @@ export interface OwnershipOptions { * folder is pinned to this host, and another host refuses the thread with * the pinned error. When you do not set it, threads are not pinned, and * every host can continue every thread. + * + * A host with an id sends a heartbeat to the log store each `renewMs`, + * also when it owns no log. A pin counts while its host is alive: when no + * heartbeat came for `ttlMs`, other hosts take its threads. */ hostId?: string /** @@ -94,7 +98,8 @@ export class LogOwnedElsewhereError extends Error { /** * The logs that one host owns, with their epochs. The host acquires a log * the first time it must write to it, and releases it when it has no work - * on it. One timer renews every owned log. With no owned log, there is no + * on it. One timer renews every owned log, and sends the heartbeat of a + * host with an `affinity`. With no owned log and no affinity, there is no * timer and no store call. */ export class HostOwnership { @@ -144,6 +149,20 @@ export class HostOwnership { const { ttlMs, renewMs } = ownershipTiming(options.timing) this.ttlMs = ttlMs this.renewMs = renewMs + // The pins of this host count only while it is alive. + if (this.affinity !== undefined) { + void this.beat() + this.startTimer() + } + } + + /** Tell the store that this host is alive. Tries again at the next tick. */ + private beat() { + const hostId = this.affinity + if (hostId === undefined) return Promise.resolve() + return Promise.resolve() + .then(() => this.store.heartbeat({ hostId, ttlMs: this.ttlMs })) + .catch(() => {}) } /** The epoch of `logId` while this host owns it. */ @@ -228,7 +247,7 @@ export class HostOwnership { forget(logId: string) { this.owned.delete(logId) this.takenOver.delete(logId) - if (this.owned.size === 0) this.stopTimer() + if (this.owned.size === 0 && this.affinity === undefined) this.stopTimer() } /** @@ -312,9 +331,10 @@ export class HostOwnership { this.renewing = true try { const owned = [...this.owned] - await Promise.all( - owned.map(([logId, epoch]) => this.renewOne(logId, epoch)), - ) + await Promise.all([ + this.beat(), + ...owned.map(([logId, epoch]) => this.renewOne(logId, epoch)), + ]) } finally { this.renewing = false } diff --git a/packages/ai-harness/tests/helpers.ts b/packages/ai-harness/tests/helpers.ts index 284878fc3a..bc714ca8db 100644 --- a/packages/ai-harness/tests/helpers.ts +++ b/packages/ai-harness/tests/helpers.ts @@ -35,6 +35,7 @@ export function wrapLogStore( renew: overrides.renew ?? inner.renew.bind(inner), release: overrides.release ?? inner.release.bind(inner), setAffinity: overrides.setAffinity ?? inner.setAffinity.bind(inner), + heartbeat: overrides.heartbeat ?? inner.heartbeat.bind(inner), listOrphaned: overrides.listOrphaned ?? inner.listOrphaned.bind(inner), offer: overrides.offer ?? inner.offer.bind(inner), offered: overrides.offered ?? inner.offered.bind(inner), diff --git a/packages/ai-harness/tests/workspace-binding.test.ts b/packages/ai-harness/tests/workspace-binding.test.ts index 11a929efb0..ec6d339aee 100644 --- a/packages/ai-harness/tests/workspace-binding.test.ts +++ b/packages/ai-harness/tests/workspace-binding.test.ts @@ -249,6 +249,8 @@ describe('the workspace binding in the session log', () => { const persistence = memoryPersistence() const { log } = persistence.stores const binding = { kind: 'host', hostId: 'host-a', root: '/ws' } + // host-a runs: its pins count. + await log.heartbeat({ hostId: 'host-a', ttlMs: 60_000 }) await appendToLog(log, { threadId: THREAD }, [ { type: 'harness.workspace.bound', binding }, ]) @@ -315,6 +317,77 @@ describe('the workspace binding in the session log', () => { await again.host.close() }) + it('keeps the pin of a live host on its idle thread, after its ownership expired', async () => { + const persistence = memoryPersistence() + const { log } = persistence.stores + const { adapter } = readTurn('c0') + const live = createHarnessHost({ + persistence, + ownership: { hostId: 'pod-a', ttlMs: 300, renewMs: 100 }, + }) + const harness = defineHarness({ + name: 'test/workspace-binding', + adapter, + plugins: () => [ + workspaceTools({ root: '/ws', backend: backend(), web: false }), + ], + }) + await (await live.open(harness, { threadId: THREAD })).send('read') + // Past the ttl: pod-a released the idle log, and it still runs. + await new Promise((resolve) => setTimeout(resolve, 600)) + const asPodA = { logId: LOG, harness: 'x', ownerId: 'x', ttlMs: 1_000 } + const released = await log.acquire({ ...asPodA, affinity: 'pod-a' }) + expect(released).not.toBeNull() + await log.release({ logId: LOG, ownerId: 'x', epoch: released?.epoch ?? 0 }) + + expect(await canOtherHostTake(log, LOG)).toBe(false) + await live.close() + }) + + it('reattaches the thread of a pinned host that stopped', async () => { + const persistence = memoryPersistence() + const workspace = { ...backend(), root: '/ws' } + const first = readTurn('c0') + const crashed = open({ + persistence, + adapter: first.adapter, + crashes: true, + hostId: 'pod-a', + plugins: [ + workspaceTools({ + workspace: { bind: async () => workspace }, + web: false, + }), + ], + }) + await (await crashed.opened).send('read') + // pod-a never closes, and sends no more heartbeats, like a stopped host. + await ownershipExpired() + const reattach = vi.fn(async () => workspace) + const { adapter, calls } = readTurn('c1') + const next = open({ + persistence, + adapter, + hostId: 'pod-b', + plugins: [ + workspaceTools({ + workspace: { bind: async () => workspace, reattach }, + web: false, + }), + ], + }) + + await (await next.opened).send('read') + + expect(reattach).toHaveBeenCalledWith( + expect.objectContaining({ + binding: { kind: 'host', hostId: 'pod-a', root: '/ws' }, + }), + ) + expect(resultOf(calls, 'c1')).toBe(READ) + await next.host.close() + }) + it('unpins a folder thread from a host that does not come back', async () => { const persistence = memoryPersistence() const { log } = persistence.stores @@ -343,6 +416,7 @@ describe('the workspace binding in the session log', () => { const { log } = persistence.stores const inFolder = { root: '/ws', backend: backend() } // A binding to pod-a, and no pin yet: host B gets the log. + await log.heartbeat({ hostId: 'pod-a', ttlMs: 60_000 }) await appendToLog(log, { threadId: THREAD }, [ { type: 'harness.workspace.bound', diff --git a/packages/ai-persistence/src/memory.ts b/packages/ai-persistence/src/memory.ts index c169ac4c84..9334617314 100644 --- a/packages/ai-persistence/src/memory.ts +++ b/packages/ai-persistence/src/memory.ts @@ -706,6 +706,8 @@ class LogBackendState { // Kept apart from `owners`, because a log that no host acquired can have a // pin too (`setAffinity` with UNOWNED_EPOCH). readonly affinities = new Map() + /** When the heartbeat of each host expires, by its host id. */ + readonly hosts = new Map() // Kept apart from `owners`, because any machine can offer to a log that no // host acquired. readonly lanes = new Map() @@ -859,8 +861,7 @@ class MemoryLogStore implements LogStore { const isLive = owner !== undefined && !owner.released && owner.expiresAt > now if (isLive && owner.ownerId !== ownerId) return null - const pin = this.backend.affinities.get(logId) - if (pin !== undefined && pin !== affinity) return null + if (this.isPinnedElsewhere(logId, affinity, now)) return null // The caller holds the log already: extend it at the same epoch. if (isLive) { owner.expiresAt = now + ttlMs @@ -909,6 +910,26 @@ class MemoryLogStore implements LogStore { if (wakeInMs !== undefined) owner.wakeAt = this.backend.now() + wakeInMs } + async heartbeat({ hostId, ttlMs }: Parameters[0]) { + assertMs({ name: 'ttlMs', value: ttlMs, allowZero: false }) + this.backend.hosts.set(hostId, this.backend.now() + ttlMs) + } + + /** + * True when `logId` is pinned to a host that is not `affinity`, and that + * host is alive: its heartbeat did not expire. The pin of a host with no + * live heartbeat does not count. + */ + private isPinnedElsewhere( + logId: string, + affinity: string | undefined, + now: number, + ) { + const pin = this.backend.affinities.get(logId) + if (pin === undefined || pin === affinity) return false + return (this.backend.hosts.get(pin) ?? 0) > now + } + async setAffinity({ logId, epoch, @@ -936,11 +957,11 @@ class MemoryLogStore implements LogStore { for (const logId of logIds) { const found = this.orphanOf(logId, now) if (found === undefined) continue - const pin = this.backend.affinities.get(logId) const isOtherHarness = harnesses !== undefined && !harnesses.includes(found.orphan.harness) - const isOtherAffinity = pin !== undefined && pin !== affinity - if (isOtherHarness || isOtherAffinity) continue + if (isOtherHarness || this.isPinnedElsewhere(logId, affinity, now)) { + continue + } due.push(found) } // The log that is due first comes first. diff --git a/packages/ai-persistence/src/testkit/conformance.ts b/packages/ai-persistence/src/testkit/conformance.ts index 9a0b8454f0..2db2ff7008 100644 --- a/packages/ai-persistence/src/testkit/conformance.ts +++ b/packages/ai-persistence/src/testkit/conformance.ts @@ -2144,6 +2144,8 @@ export function runPersistenceConformance( const logId = newLog() const harness = newHarness() const claim = { ...claimOf(logId, 'a'), harness } + // A pin counts while its host is alive. + await store.heartbeat({ hostId: 'host-1', ttlMs: LIVE_MS }) await store.acquire(claim) await store.setAffinity({ logId, epoch: 1, affinity: 'host-1' }) await store.release({ logId, ownerId: 'a', epoch: 1, wakeInMs: 0 }) @@ -2168,6 +2170,31 @@ export function runPersistenceConformance( ).toEqual({ epoch: 2 }) }) + it('lets other hosts take a pinned log when the heartbeat of its host expired', async (ctx) => { + const store = persistence.stores.log + if (!store) return ctx.skip('log store not provided') + const logId = newLog() + const harness = newHarness() + const host = `host-${logId}` + await store.heartbeat({ hostId: host, ttlMs: SHORT_MS }) + await store.acquire({ ...claimOf(logId, 'a'), harness, affinity: host }) + await store.setAffinity({ logId, epoch: 1, affinity: host }) + await store.release({ logId, ownerId: 'a', epoch: 1, wakeInMs: 0 }) + const other = { ...claimOf(logId, 'b'), harness } + + // The host is alive: its idle log stays with it. + expect( + await store.listOrphaned({ limit: 10, harnesses: [harness] }), + ).toEqual([]) + expect(await store.acquire(other)).toBeNull() + + await waitPastShort() + expect( + await store.listOrphaned({ limit: 10, harnesses: [harness] }), + ).toEqual([{ logId, harness, reason: 'wake' }]) + expect(await store.acquire(other)).toEqual({ epoch: 2 }) + }) + it('fences setAffinity', async (ctx) => { const store = persistence.stores.log if (!store) return ctx.skip('log store not provided') @@ -2202,6 +2229,9 @@ export function runPersistenceConformance( await expect( store.release({ logId, ownerId: 'a', epoch: 1, wakeInMs: -1 }), ).rejects.toBeInstanceOf(RangeError) + await expect( + store.heartbeat({ hostId: 'host-1', ttlMs: 0 }), + ).rejects.toBeInstanceOf(RangeError) // The refused release did not release the log. expect(await store.acquire(claimOf(logId, 'b'))).toBeNull() }) @@ -2431,6 +2461,7 @@ export function runPersistenceConformance( // A pin applies to a log that no host acquired too. const pinned = newLog() const pinnedHarness = newHarness() + await store.heartbeat({ hostId: 'host-1', ttlMs: 10_000 }) await store.setAffinity({ logId: pinned, epoch: UNOWNED_EPOCH, @@ -2642,6 +2673,7 @@ export function runPersistenceConformance( const harness = newHarness() const claim = { logId, harness, ttlMs: LIVE_MS } const input = { logId, harness, inputId: 'i1', record: { type: 'in' } } + await store.heartbeat({ hostId: 'host-1', ttlMs: LIVE_MS }) await store.acquire({ ...claim, ownerId: 'a' }) await store.setAffinity({ logId, epoch: 1, affinity: 'host-1' }) await store.append({ logId, seq: 1, epoch: 1, records: recordsOf(2) }) diff --git a/packages/ai-persistence/src/types.ts b/packages/ai-persistence/src/types.ts index 89d98d5438..acefb1d3a9 100644 --- a/packages/ai-persistence/src/types.ts +++ b/packages/ai-persistence/src/types.ts @@ -845,7 +845,8 @@ export interface LogStore { * - Returns `null` when another owner holds the log (not released, and not * expired). * - Returns `null` when the log has an affinity that is not - * `claim.affinity`. + * `claim.affinity`, and the host of that affinity is alive (see + * `heartbeat`). The pin of a host with no live heartbeat does not count. * - When `ownerId` already holds the log, extends the expiry and returns * the same epoch. * - Else increments the epoch, stores `harness`, clears the wake time, and @@ -900,6 +901,14 @@ export interface LogStore { epoch: number affinity: string | null }) => Promise + /** + * Mark the host `hostId` alive for `ttlMs` milliseconds of the store + * clock. `hostId` is the affinity of its pins. A pin counts only while its + * host is alive, so the logs of a host that stopped go to other hosts. + * + * Rejects with a `RangeError` when `ttlMs` is not more than 0. + */ + heartbeat: (beat: { hostId: string; ttlMs: number }) => Promise /** * The logs that have work and no live owner, at most `limit` of them. A log * is listed for the first reason that applies: @@ -909,7 +918,8 @@ export interface LogStore { * 3. `'wake'`: released, and the wake time is due. * * Only logs of `harnesses` are listed, when `harnesses` is given. A log with - * an affinity is listed only when `affinity` is equal. The log that is due + * an affinity is listed only when `affinity` is equal, or the host of the + * affinity is not alive (see `heartbeat`). The log that is due * first comes first. `limit: 0` gives `[]`. */ listOrphaned: (query: { From da013f28ba85f961a246682c859f584f1f1602d4 Mon Sep 17 00:00:00 2001 From: Alem Tuzlak Date: Mon, 19 Oct 2026 15:16:02 +0200 Subject: [PATCH 15/17] fix(ai-utils): type the process env of commandEnv, so the React Native smoke typechecks Without process, commandEnv reads {}. The union of process.env and {} cannot be indexed by a string under the React Native smoke tsconfig (TS7053 at command.ts:46), and the Test job stopped there. --- packages/ai-utils/src/command.ts | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/packages/ai-utils/src/command.ts b/packages/ai-utils/src/command.ts index b70bbf91c3..c6ec427ab6 100644 --- a/packages/ai-utils/src/command.ts +++ b/packages/ai-utils/src/command.ts @@ -39,7 +39,8 @@ export function commandEnv( const names = isWindows ? [...KEPT_ENV, ...KEPT_WINDOWS_ENV, ...passEnv] : [...KEPT_ENV, ...passEnv] - const own = typeof process !== 'undefined' ? process.env : {} + const own: Record = + typeof process !== 'undefined' ? process.env : {} const kept: Record = {} for (const name of names) { // On Windows, `process.env` finds a name in any letter case. From 3889b0c2003760bb0be91e40a8dfabc79ca35739 Mon Sep 17 00:00:00 2001 From: Alem Tuzlak Date: Mon, 19 Oct 2026 15:16:02 +0200 Subject: [PATCH 16/17] fix(ai-vertex): give ai-vertex the same @google/genai copy as ai-gemini again The lockfile moved the @google/genai devDependency of ai-vertex to the zod@4.6.5 peer copy, while ai-gemini keeps the zod@4.3.6 copy. Then the vi.mock('@google/genai') of the vertex tests mocks another file than the one that ai-gemini imports, and tests/factories.test.ts fails. --- pnpm-lock.yaml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index 7b7cfea333..57c8f09b7b 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -3580,7 +3580,7 @@ importers: devDependencies: '@google/genai': specifier: ^2.10.0 - version: 2.10.0(@modelcontextprotocol/sdk@1.29.0(@cfworker/json-schema@4.1.1)(zod@4.6.5)) + version: 2.10.0(@modelcontextprotocol/sdk@1.29.0(@cfworker/json-schema@4.1.1)(zod@4.3.6)) '@tanstack/ai': specifier: workspace:* version: link:../ai From f8170e1e9c8230f8ad8c6c9b86a9bd33adf7797b Mon Sep 17 00:00:00 2001 From: Alem Tuzlak Date: Mon, 19 Oct 2026 15:16:03 +0200 Subject: [PATCH 17/17] test(ai-harness): give the live-host pin test a 1 s heartbeat With a 300 ms ttl and a 100 ms heartbeat, a CI run that holds the event loop for 200 ms lets the heartbeat expire, and the pin stops counting. --- packages/ai-harness/tests/workspace-binding.test.ts | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/packages/ai-harness/tests/workspace-binding.test.ts b/packages/ai-harness/tests/workspace-binding.test.ts index ec6d339aee..7581e06b28 100644 --- a/packages/ai-harness/tests/workspace-binding.test.ts +++ b/packages/ai-harness/tests/workspace-binding.test.ts @@ -321,9 +321,11 @@ describe('the workspace binding in the session log', () => { const persistence = memoryPersistence() const { log } = persistence.stores const { adapter } = readTurn('c0') + // A heartbeat each 100 ms lasts 1 s. So a slow CI run that holds the + // event loop for some hundred ms does not let it expire. const live = createHarnessHost({ persistence, - ownership: { hostId: 'pod-a', ttlMs: 300, renewMs: 100 }, + ownership: { hostId: 'pod-a', ttlMs: 1_000, renewMs: 100 }, }) const harness = defineHarness({ name: 'test/workspace-binding', @@ -334,7 +336,7 @@ describe('the workspace binding in the session log', () => { }) await (await live.open(harness, { threadId: THREAD })).send('read') // Past the ttl: pod-a released the idle log, and it still runs. - await new Promise((resolve) => setTimeout(resolve, 600)) + await new Promise((resolve) => setTimeout(resolve, 1_200)) const asPodA = { logId: LOG, harness: 'x', ownerId: 'x', ttlMs: 1_000 } const released = await log.acquire({ ...asPodA, affinity: 'pod-a' }) expect(released).not.toBeNull()