From d3e5cf7cd9087f02d13b56bdaf3fb31eb7b77d7c Mon Sep 17 00:00:00 2001 From: Alem Tuzlak Date: Sat, 17 Oct 2026 12:46:41 +0200 Subject: [PATCH 01/12] fix(ai-harness, ai-mcp): OAuth flows and commands stop on their signal --- packages/ai-harness/src/connectors.ts | 16 +- packages/ai-harness/src/oauth.ts | 270 ++++++++++++------ packages/ai-harness/src/session/turn.ts | 27 +- packages/ai-harness/tests/auth.test.ts | 42 ++- packages/ai-harness/tests/oauth-edges.test.ts | 169 ++++++++--- .../ai-harness/tests/oauth-loopback.test.ts | 43 ++- packages/ai-mcp/src/connector.ts | 15 +- packages/ai-mcp/tests/connector-edges.test.ts | 35 +++ 8 files changed, 461 insertions(+), 156 deletions(-) diff --git a/packages/ai-harness/src/connectors.ts b/packages/ai-harness/src/connectors.ts index b70675e211..bad1d790ed 100644 --- a/packages/ai-harness/src/connectors.ts +++ b/packages/ai-harness/src/connectors.ts @@ -77,7 +77,10 @@ export function oauthConnector(options: OAuthConnectorOptions) { let credential = await credentials.require(id, { wait }) if (credential.type === 'api_key') return credential.value if (isExpired(credential) && credential.refreshToken) { - credential = await refreshCredential(oauth, credential, options.fetch) + credential = await refreshCredential(oauth, credential, { + fetch: options.fetch, + signal: run.signal, + }) await credentials.set({ id, credential }) } if (credential.type === 'api_key') return credential.value @@ -88,8 +91,9 @@ export function oauthConnector(options: OAuthConnectorOptions) { commands: { [`connect:${id}`]: defineCommand({ description: `Sign in to ${label}`, - // The sign-in belongs to the user who runs the command. - run: async (_input, { credentials }) => { + // The sign-in belongs to the user who runs the command. A cancel + // or a session close stops it at once. + run: async (_input, { credentials, signal }) => { const credential = options.login === 'device' ? await deviceLogin(oauth, { @@ -99,12 +103,14 @@ export function oauthConnector(options: OAuthConnectorOptions) { url: verificationUri, userCode, }), - ...(options.fetch ? { fetch: options.fetch } : {}), + fetch: options.fetch, + signal, }) : await loopbackLogin(oauth, { onUrl: (url) => ctx.session.authRequired({ connector: id, url }), - ...(options.fetch ? { fetch: options.fetch } : {}), + fetch: options.fetch, + signal, }) await credentials.set({ id, credential }) return `Connected to ${label}.` diff --git a/packages/ai-harness/src/oauth.ts b/packages/ai-harness/src/oauth.ts index cccf37f653..00e83e2fa3 100644 --- a/packages/ai-harness/src/oauth.ts +++ b/packages/ai-harness/src/oauth.ts @@ -11,21 +11,41 @@ export interface OAuthConfig { /** Only for confidential clients. A CLI is a public client and has none. */ clientSecret?: string scopes?: ReadonlyArray + /** + * The issuer of the authorization server. When the server sends `iss` with + * the code (RFC 9207), `loopbackLogin` refuses a code from another issuer. + */ + issuer?: string } type Fetch = typeof fetch const encoder = new TextEncoder() -function randomString(bytes = 32): string { +function randomString(bytes = 32) { return base64url(crypto.getRandomValues(new Uint8Array(bytes))) } +/** + * Resolve after `ms`, or at once when `signal` aborts. The timer is not + * unref'd: it is the wait of a device sign-in, and a script that only waits + * for that sign-in must not exit. + */ +function wait(ms: number, signal?: AbortSignal) { + return new Promise((resolve) => { + 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 }) + }) +} + /** A PKCE pair with the S256 method (RFC 7636). */ -export async function createPkce(): Promise<{ - verifier: string - challenge: string -}> { +export async function createPkce() { const verifier = randomString(32) const digest = await crypto.subtle.digest('SHA-256', encoder.encode(verifier)) return { verifier, challenge: base64url(new Uint8Array(digest)) } @@ -35,7 +55,7 @@ export async function createPkce(): Promise<{ export function buildAuthorizationUrl( config: OAuthConfig, options: { redirectUri: string; state: string; challenge: string }, -): string { +) { const url = new URL(config.authorizationUrl) url.searchParams.set('response_type', 'code') url.searchParams.set('client_id', config.clientId) @@ -51,8 +71,9 @@ export function buildAuthorizationUrl( async function tokenRequest( config: OAuthConfig, body: Record, - doFetch: Fetch, -): Promise> { + options: { fetch?: Fetch; signal?: AbortSignal }, +) { + const doFetch = options.fetch ?? fetch const response = await doFetch(config.tokenUrl, { method: 'POST', headers: { @@ -64,6 +85,7 @@ async function tokenRequest( ...(config.clientSecret ? { client_secret: config.clientSecret } : {}), ...body, }), + signal: options.signal, }) const parsed: unknown = await response.json().catch(() => ({})) if (typeof parsed !== 'object' || parsed === null) { @@ -101,7 +123,10 @@ function credentialFrom( } } -/** Trade an authorization code for tokens. */ +/** + * Trade an authorization code for tokens. When `signal` aborts, the request + * stops and the call rejects with the reason of the signal. + */ export async function exchangeCode( config: OAuthConfig, options: { @@ -109,8 +134,9 @@ export async function exchangeCode( verifier: string redirectUri: string fetch?: Fetch + signal?: AbortSignal }, -): Promise { +) { const response = await tokenRequest( config, { @@ -119,13 +145,13 @@ export async function exchangeCode( code_verifier: options.verifier, redirect_uri: options.redirectUri, }, - options.fetch ?? fetch, + options, ) return credentialFrom(response) } /** True when an OAuth credential expires within `skewMs`. */ -export function isExpired(credential: Credential, skewMs = 60_000): boolean { +export function isExpired(credential: Credential, skewMs = 60_000) { return ( credential.type === 'oauth' && credential.expiresAt !== undefined && @@ -133,19 +159,28 @@ export function isExpired(credential: Credential, skewMs = 60_000): boolean { ) } -/** Get a new access token with the refresh token. */ +/** + * Get a new access token with the refresh token. `fetch` replaces the global + * `fetch`. When `signal` aborts, the request stops and the call rejects with + * the reason of the signal. + * + * @example + * ```ts + * const fresh = await refreshCredential(config, credential, { signal }) + * ``` + */ export async function refreshCredential( config: OAuthConfig, credential: Credential, - doFetch: Fetch = fetch, -): Promise { + options: { fetch?: Fetch; signal?: AbortSignal } = {}, +) { if (credential.type !== 'oauth' || !credential.refreshToken) { throw new Error('This credential has no refresh token.') } const response = await tokenRequest( config, { grant_type: 'refresh_token', refresh_token: credential.refreshToken }, - doFetch, + options, ) return credentialFrom(response, credential) } @@ -162,59 +197,97 @@ const DONE_PAGE = * `waitForCode` resolves with the code and the `iss` the server sent with it * (RFC 9207). Pass `iss` on to the library: a server that sends it can refuse * a code without it. + * + * A request to the callback path with no `code` and no `state` gets a 400, + * and the sign-in goes on. `close()` stops listening, and a `waitForCode` + * that waits rejects. When `signal` aborts, the receiver closes, and + * `waitForCode` rejects with the reason of the signal. + * + * @example + * ```ts + * const receiver = await startLoopbackReceiver({ signal }) + * try { + * // Register receiver.redirectUri, then send the user to the sign-in page. + * const { code, iss } = await receiver.waitForCode(state) + * } finally { + * receiver.close() + * } + * ``` */ export async function startLoopbackReceiver( - options: { timeoutMs?: number } = {}, -): Promise<{ - redirectUri: string - waitForCode: (state: string) => Promise<{ code: string; iss?: string }> - close: () => void -}> { + options: { timeoutMs?: number; signal?: AbortSignal } = {}, +) { + const { signal } = options const { createServer } = await import('node:http') const server = createServer() await new Promise((resolve, reject) => { server.once('error', reject) server.listen(0, '127.0.0.1', () => resolve()) }) + if (signal?.aborted) { + server.close() + throw signal.reason + } const address = server.address() const port = typeof address === 'object' && address ? address.port : 0 const redirectUri = `http://127.0.0.1:${port}/callback` - const close = () => { + let timer: ReturnType | undefined + /** Rejects the `waitForCode` that waits. */ + let fail: (error: unknown) => void = () => {} + /** Stop listening. With `error`, a `waitForCode` that waits rejects. */ + const stop = (error?: unknown) => { + clearTimeout(timer) + signal?.removeEventListener('abort', onAbort) server.closeAllConnections() server.close() + if (error !== undefined) fail(error) } + const onAbort = () => stop(signal?.reason) + signal?.addEventListener('abort', onAbort, { once: true }) return { redirectUri, - close, - waitForCode: (state) => - new Promise((resolve, reject) => { - const timer = setTimeout( - () => { - close() - reject(new Error('Sign-in timed out.')) - }, + close: () => stop(new Error('Sign-in cancelled.')), + waitForCode: (state: string) => + new Promise<{ code: string; iss?: string }>((resolve, reject) => { + if (!server.listening) { + reject( + signal?.aborted ? signal.reason : new Error('Sign-in cancelled.'), + ) + return + } + fail = reject + timer = setTimeout( + () => stop(new Error('Sign-in timed out.')), options.timeoutMs ?? 10 * 60_000, ) + // The deadline must not keep a CLI or a test process alive. + if (typeof timer === 'object' && 'unref' in timer) timer.unref() server.on('request', (req, res) => { const url = new URL(req.url ?? '/', redirectUri) if (url.pathname !== '/callback') { res.writeHead(404).end() return } - clearTimeout(timer) + const params = url.searchParams + const code = params.get('code') + const returnedState = params.get('state') + // Not an answer of the server, for example a reload or a probe. + if (code === null && returnedState === null) { + res.writeHead(400).end() + return + } res.writeHead(200, { 'Content-Type': 'text/html' }).end(DONE_PAGE) - close() - const code = url.searchParams.get('code') - const iss = url.searchParams.get('iss') - if (url.searchParams.get('state') !== state) { - reject(new Error('Sign-in failed: the state does not match.')) + const iss = params.get('iss') + if (returnedState !== state) { + stop(new Error('Sign-in failed: the state does not match.')) } else if (!code) { - reject( - new Error( - `Sign-in failed: ${url.searchParams.get('error') ?? 'no code'}`, - ), + stop( + new Error(`Sign-in failed: ${params.get('error') ?? 'no code'}`), ) - } else resolve(iss ? { code, iss } : { code }) + } else { + resolve(iss ? { code, iss } : { code }) + stop() + } }) }), } @@ -224,84 +297,87 @@ export async function startLoopbackReceiver( * Sign in through the browser with a loopback redirect (RFC 8252 + PKCE). * Listens on `127.0.0.1` on a random port, for one callback only. Calls * `onUrl` with the URL to open. Resolves with the tokens. + * + * When the server sends `iss` with the code (RFC 9207) and `config.issuer` + * is set, they must be the same. When `signal` aborts, the sign-in stops at + * once, stops listening, and rejects with the reason of the signal. + * + * @example + * ```ts + * const credential = await loopbackLogin(config, { + * onUrl: (url) => console.log(`Open ${url}`), + * signal, + * }) + * ``` */ export async function loopbackLogin( config: OAuthConfig, - options: { onUrl: (url: string) => void; fetch?: Fetch; timeoutMs?: number }, -): Promise { - const { createServer } = await import('node:http') + options: { + onUrl: (url: string) => void + fetch?: Fetch + timeoutMs?: number + signal?: AbortSignal + }, +) { const { verifier, challenge } = await createPkce() const state = randomString(16) - const server = createServer() - await new Promise((resolve, reject) => { - server.once('error', reject) - server.listen(0, '127.0.0.1', () => resolve()) - }) - const address = server.address() - const port = typeof address === 'object' && address ? address.port : 0 - const redirectUri = `http://127.0.0.1:${port}/callback` - + const receiver = await startLoopbackReceiver(options) + const { redirectUri } = receiver try { - const code = await new Promise((resolve, reject) => { - const timer = setTimeout( - () => reject(new Error('Sign-in timed out.')), - options.timeoutMs ?? 10 * 60_000, - ) - server.on('request', (req, res) => { - const url = new URL(req.url ?? '/', redirectUri) - if (url.pathname !== '/callback') { - res.writeHead(404).end() - return - } - clearTimeout(timer) - const returnedState = url.searchParams.get('state') - const returnedCode = url.searchParams.get('code') - res.writeHead(200, { 'Content-Type': 'text/html' }).end(DONE_PAGE) - if (returnedState !== state) - reject(new Error('Sign-in failed: the state does not match.')) - else if (!returnedCode) { - reject( - new Error( - `Sign-in failed: ${url.searchParams.get('error') ?? 'no code'}`, - ), - ) - } else resolve(returnedCode) - }) - options.onUrl( - buildAuthorizationUrl(config, { redirectUri, state, challenge }), - ) - }) + options.onUrl( + buildAuthorizationUrl(config, { redirectUri, state, challenge }), + ) + const { code, iss } = await receiver.waitForCode(state) + const isOtherIssuer = + iss !== undefined && config.issuer !== undefined && iss !== config.issuer + if (isOtherIssuer) { + throw new Error('Sign-in failed: the issuer does not match.') + } return await exchangeCode(config, { code, verifier, redirectUri, - ...(options.fetch ? { fetch: options.fetch } : {}), + fetch: options.fetch, + signal: options.signal, }) } finally { - server.closeAllConnections() - server.close() + receiver.close() } } /** * Sign in with a device code (RFC 8628), for SSH sessions, containers, and * CI. Calls `onCode` with the code and the page to enter it on, then polls. + * It waits the interval the server asks for, at least 1 second. When + * `signal` aborts, it stops polling at once and rejects with the reason of + * the signal. + * + * @example + * ```ts + * const credential = await deviceLogin(config, { + * onCode: ({ userCode, verificationUri }) => + * console.log(`Enter ${userCode} at ${verificationUri}`), + * signal, + * }) + * ``` */ export async function deviceLogin( config: OAuthConfig, options: { onCode: (info: { userCode: string; verificationUri: string }) => void fetch?: Fetch - /** Test hook. Default waits the interval the server asks for. */ - sleep?: (ms: number) => Promise + signal?: AbortSignal + /** + * Test hook. Default waits `ms`, and ends early when `signal` aborts. + */ + sleep?: (ms: number, signal?: AbortSignal) => Promise }, -): Promise { +) { if (!config.deviceUrl) throw new Error('This OAuth app has no device endpoint.') + const { signal } = options const doFetch = options.fetch ?? fetch - const sleep = - options.sleep ?? - ((ms: number) => new Promise((resolve) => setTimeout(resolve, ms))) + const sleep = options.sleep ?? wait const started = await doFetch(config.deviceUrl, { method: 'POST', headers: { @@ -312,6 +388,7 @@ export async function deviceLogin( client_id: config.clientId, ...(config.scopes?.length ? { scope: config.scopes.join(' ') } : {}), }), + signal, }) const device: unknown = await started.json() if ( @@ -328,9 +405,11 @@ export async function deviceLogin( 'verification_uri' in device && typeof device.verification_uri === 'string' ? device.verification_uri : config.authorizationUrl + // 5 seconds when the server names no interval (RFC 8628 section 3.5). At + // least 1 second, so a server that sends 0 gets no tight loop. let interval = 'interval' in device && typeof device.interval === 'number' - ? device.interval * 1000 + ? Math.max(device.interval, 1) * 1000 : 5000 const expiresAt = Date.now() + @@ -340,14 +419,15 @@ export async function deviceLogin( options.onCode({ userCode: device.user_code, verificationUri }) while (Date.now() < expiresAt) { - await sleep(interval) + await sleep(interval, signal) + signal?.throwIfAborted() const response = await tokenRequest( config, { grant_type: 'urn:ietf:params:oauth:grant-type:device_code', device_code: device.device_code, }, - doFetch, + options, ) if (response.error === 'authorization_pending') continue if (response.error === 'slow_down') { diff --git a/packages/ai-harness/src/session/turn.ts b/packages/ai-harness/src/session/turn.ts index 93490427fa..581a290772 100644 --- a/packages/ai-harness/src/session/turn.ts +++ b/packages/ai-harness/src/session/turn.ts @@ -160,14 +160,25 @@ export abstract class SessionTurn< signal: operation.abortController.signal, }) try { - const result: unknown = await entry.command.run(checked, { - signal: operation.abortController.signal, - // The command acts for its sender, also while another user's turn runs. - session: this.pluginApi(principal, undefined, operation.id), - ...(principal ? { principal } : {}), - credentials: run.credentials, - run, - }) + const { signal } = operation.abortController + // A cancel or a session close ends the command at once. A handler that + // does not stop on its signal runs on alone, and its result is dropped. + const result: unknown = await Promise.race([ + entry.command.run(checked, { + signal, + // The command acts for its sender, also while another user's turn runs. + session: this.pluginApi(principal, undefined, operation.id), + ...(principal ? { principal } : {}), + credentials: run.credentials, + run, + }), + new Promise((_, reject) => { + if (signal.aborted) reject(signal.reason) + signal.addEventListener('abort', () => reject(signal.reason), { + once: true, + }) + }), + ]) operation.publish( customEvent('harness.command.result', { name, diff --git a/packages/ai-harness/tests/auth.test.ts b/packages/ai-harness/tests/auth.test.ts index 19b19a53c9..04703afed3 100644 --- a/packages/ai-harness/tests/auth.test.ts +++ b/packages/ai-harness/tests/auth.test.ts @@ -96,7 +96,7 @@ describe('token exchange', () => { }) expect(bodies[0]?.get('code_verifier')).toBe('v1') - const refreshed = await refreshCredential(oauth, first, fake) + const refreshed = await refreshCredential(oauth, first, { fetch: fake }) // The old refresh token stays when the server sends no new one. expect(refreshed).toMatchObject({ accessToken: 'a2', refreshToken: 'r1' }) expect(bodies[1]?.get('grant_type')).toBe('refresh_token') @@ -276,6 +276,46 @@ describe('oauthConnector', () => { }) }) +describe('commands on close', () => { + it('ends a command that ignores its signal when the session closes', async () => { + let started: () => void = () => {} + const running = new Promise((resolve) => { + started = resolve + }) + const plugin = definePlugin({ + name: 'test/stuck', + setup: () => ({ + commands: { + stuck: defineCommand({ + description: 'Waits for ever', + run: () => { + started() + return new Promise(() => {}) + }, + }), + }, + }), + }) + const host = createHarnessHost({ + persistence: memoryPersistence(), + credentials: memoryCredentialStore(), + }) + const session = await host.open( + defineHarness({ + name: 'test/stuck', + adapter: mockAdapter(() => text('ok')).adapter, + plugins: () => [plugin], + }), + { threadId: 't1' }, + ) + const command = session.command({ name: 'stuck' }) + await running + await session.close() + expect(command.status()).toBe('cancelled') + await host.close() + }) +}) + describe('scrubSecrets', () => { it('removes known secrets from text', () => { expect(scrubSecrets('token gh-abcdef1234 failed', ['gh-abcdef1234'])).toBe( diff --git a/packages/ai-harness/tests/oauth-edges.test.ts b/packages/ai-harness/tests/oauth-edges.test.ts index 58c13735dc..427787762c 100644 --- a/packages/ai-harness/tests/oauth-edges.test.ts +++ b/packages/ai-harness/tests/oauth-edges.test.ts @@ -1,5 +1,5 @@ import { memoryCredentialStore } from '@tanstack/ai/byok/server' -import { describe, expect, it } from 'vitest' +import { afterEach, describe, expect, it, vi } from 'vitest' import { toolDefinition } from '@tanstack/ai' import { memoryPersistence } from '@tanstack/ai-persistence' import { @@ -18,7 +18,7 @@ import { refreshCredential, } from '../src/oauth' import { mockAdapter, text, toolCall } from './helpers' -import type { HarnessRunContext } from '../src' +import type { HarnessRunContext, HarnessSession } from '../src' import type { OAuthConfig } from '../src/oauth' const oauth: OAuthConfig = { @@ -28,6 +28,35 @@ const oauth: OAuthConfig = { clientId: 'client-1', } +afterEach(() => { + vi.useRealTimers() +}) + +/** A device-flow `sleep` that returns at once, and records each wait. */ +function recordedSleep() { + const waits: Array = [] + const sleep = async (ms: number) => { + waits.push(ms) + } + return { waits, sleep } +} + +/** The first sign-in URL that `session` asks the user to open. */ +async function nextSignInUrl(session: HarnessSession) { + for await (const entry of session.events()) { + const value = entry.event.type === 'CUSTOM' ? entry.event.value : undefined + if ( + typeof value === 'object' && + value !== null && + 'url' in value && + typeof value.url === 'string' + ) { + return new URL(value.url) + } + } + throw new Error('The session closed before it asked for a sign-in.') +} + /** A fetch that answers each call with the next body, and records the forms. */ function replies(bodies: Array, status = 200) { const forms: Array = [] @@ -52,7 +81,7 @@ describe('token requests', () => { const refreshed = await refreshCredential( { ...oauth, clientSecret: 'shh' }, { type: 'oauth', accessToken: 'a1', refreshToken: 'r1' }, - fake, + { fetch: fake }, ) expect(forms[0]?.get('client_secret')).toBe('shh') expect(refreshed).toMatchObject({ @@ -118,12 +147,11 @@ describe('token requests', () => { describe('loopbackLogin edges', () => { async function login( callback: (redirect: URL, state: string) => string, - options: { timeoutMs?: number } = {}, + config = oauth, ) { const { fake } = replies([{ access_token: 'token' }]) - return loopbackLogin(oauth, { + return loopbackLogin(config, { fetch: fake, - ...options, onUrl: (value) => { const url = new URL(value) const redirect = new URL(url.searchParams.get('redirect_uri') ?? '') @@ -153,6 +181,34 @@ describe('loopbackLogin edges', () => { loopbackLogin(oauth, { onUrl: () => {}, timeoutMs: 10 }), ).rejects.toThrow('Sign-in timed out.') }) + + it('refuses a code from another issuer (RFC 9207)', async () => { + const config = { ...oauth, issuer: 'https://auth.example' } + const sentBy = (iss: string) => (redirect: URL, state: string) => + `${redirect}?code=c&state=${state}&iss=${encodeURIComponent(iss)}` + await expect( + login(sentBy('https://other.example'), config), + ).rejects.toThrow('Sign-in failed: the issuer does not match.') + await expect( + login(sentBy('https://auth.example'), config), + ).resolves.toMatchObject({ accessToken: 'token' }) + }) + + it('stops listening and leaves no timer when onUrl throws', async () => { + vi.useFakeTimers({ toFake: ['setTimeout', 'clearTimeout'] }) + let redirect = '' + await expect( + loopbackLogin(oauth, { + onUrl: (value) => { + redirect = new URL(value).searchParams.get('redirect_uri') ?? '' + throw new Error('No browser here.') + }, + }), + ).rejects.toThrow('No browser here.') + expect(vi.getTimerCount()).toBe(0) + vi.useRealTimers() + await expect(fetch(redirect)).rejects.toThrow() + }) }) describe('deviceLogin edges', () => { @@ -172,7 +228,7 @@ describe('deviceLogin edges', () => { }) it('slows down when asked, sends scopes, and falls back to the authorization URL', async () => { - const waits: Array = [] + const { waits, sleep } = recordedSleep() const shown: Array = [] const { fake, forms } = replies([ { device_code: 'd', user_code: 'CODE' }, @@ -186,9 +242,7 @@ describe('deviceLogin edges', () => { onCode: (info) => shown.push(`${info.userCode} ${info.verificationUri}`), fetch: fake, - sleep: async (ms) => { - waits.push(ms) - }, + sleep, }, ) expect(forms[0]?.get('scope')).toBe('repo') @@ -197,6 +251,40 @@ describe('deviceLogin edges', () => { expect(credential).toMatchObject({ accessToken: 'granted' }) }) + it('waits at least one second between polls', async () => { + const { waits, sleep } = recordedSleep() + const { fake } = replies([ + { device_code: 'd', user_code: 'C', interval: 0 }, + { access_token: 'granted' }, + ]) + await deviceLogin(oauth, { onCode: () => {}, fetch: fake, sleep }) + expect(waits).toEqual([1000]) + }) + + it('stops polling at once when the signal aborts', async () => { + vi.useFakeTimers({ toFake: ['setTimeout', 'clearTimeout'] }) + const controller = new AbortController() + const reason = new Error('The user left.') + const { fake, forms } = replies([{ device_code: 'd', user_code: 'C' }]) + let showCode: () => void = () => {} + const shown = new Promise((resolve) => { + showCode = resolve + }) + const login = deviceLogin(oauth, { + onCode: () => showCode(), + fetch: fake, + signal: controller.signal, + }) + await shown + // The first poll waits its 5 seconds. + expect(vi.getTimerCount()).toBe(1) + controller.abort(reason) + await expect(login).rejects.toBe(reason) + // Only the device request went out, and no timer is left. + expect(forms).toHaveLength(1) + expect(vi.getTimerCount()).toBe(0) + }) + it('stops when the device code expires', async () => { const { fake } = replies([ { device_code: 'd', user_code: 'C', expires_in: 0, interval: 0 }, @@ -278,45 +366,36 @@ describe('oauthConnector edges', () => { await host.close() }) - it('signs in through the loopback and signs out', async () => { - const persistence = memoryPersistence() + /** A session of thread `t` with the loopback connector `svc`. */ + async function openLoopback(fetchFn: typeof fetch) { const credentials = memoryCredentialStore() - const { fake } = replies([{ access_token: 'loop-token' }]) - const host = createHarnessHost({ persistence, credentials }) + const host = createHarnessHost({ + persistence: memoryPersistence(), + credentials, + }) const session = await host.open( defineHarness({ name: 'test/loopback-connector', adapter: mockAdapter([]).adapter, - plugins: () => [connector(fake, [])], + plugins: () => [connector(fetchFn, [])], }), { threadId: 't', }, ) - const controller = new AbortController() - const signIn = (async () => { - for await (const entry of session.events({ signal: controller.signal })) { - const value = - entry.event.type === 'CUSTOM' ? entry.event.value : undefined - if ( - typeof value === 'object' && - value !== null && - 'url' in value && - typeof value.url === 'string' - ) { - const url = new URL(value.url) - const redirect = url.searchParams.get('redirect_uri') ?? '' - const state = url.searchParams.get('state') ?? '' - await fetch(`${redirect}?code=c&state=${state}`) - return - } - } - })() - expect(await session.command({ name: 'connect:svc' })).toBe( - 'Connected to Service.', - ) - controller.abort() - await signIn + return { credentials, host, session } + } + + it('signs in through the loopback and signs out', async () => { + const { fake } = replies([{ access_token: 'loop-token' }]) + const { credentials, host, session } = await openLoopback(fake) + const asked = nextSignInUrl(session) + const connecting = session.command({ name: 'connect:svc' }) + const url = await asked + const redirect = url.searchParams.get('redirect_uri') ?? '' + const state = url.searchParams.get('state') ?? '' + await fetch(`${redirect}?code=c&state=${state}`) + expect(await connecting).toBe('Connected to Service.') expect( await credentials.get({ scope: { threadId: 't' }, @@ -334,6 +413,18 @@ describe('oauthConnector edges', () => { ).toBeNull() await host.close() }) + + it('closes the session at once while /connect waits for the browser', async () => { + const { host, session } = await openLoopback(replies([]).fake) + const asked = nextSignInUrl(session) + const connecting = session.command({ name: 'connect:svc' }) + const redirect = (await asked).searchParams.get('redirect_uri') ?? '' + await session.close() + expect(connecting.status()).toBe('cancelled') + // The sign-in stopped too: nothing listens for the redirect now. + await expect(fetch(redirect)).rejects.toThrow() + await host.close() + }) }) describe('auth helpers', () => { diff --git a/packages/ai-harness/tests/oauth-loopback.test.ts b/packages/ai-harness/tests/oauth-loopback.test.ts index ef4670d907..96d50497c8 100644 --- a/packages/ai-harness/tests/oauth-loopback.test.ts +++ b/packages/ai-harness/tests/oauth-loopback.test.ts @@ -1,8 +1,12 @@ -import { describe, expect, it } from 'vitest' +import { afterEach, describe, expect, it, vi } from 'vitest' import { startLoopbackReceiver } from '../src/oauth' +afterEach(() => { + vi.restoreAllMocks() +}) + describe('startLoopbackReceiver', () => { - it('returns the code for the matching state and ignores other paths', async () => { + it('returns the code for the matching state and ignores other requests', async () => { const receiver = await startLoopbackReceiver() expect(receiver.redirectUri).toMatch( /^http:\/\/127\.0\.0\.1:\d+\/callback$/, @@ -12,6 +16,9 @@ describe('startLoopbackReceiver', () => { receiver.redirectUri.replace('/callback', '/favicon.ico'), ) expect(other.status).toBe(404) + // A callback with no code and no state is not the answer of the server. + const stray = await fetch(receiver.redirectUri) + expect(stray.status).toBe(400) const done = await fetch(`${receiver.redirectUri}?code=abc&state=s1`) expect(done.status).toBe(200) expect(await done.text()).toContain('signed in') @@ -58,4 +65,36 @@ describe('startLoopbackReceiver', () => { await expect(fetch(receiver.redirectUri)).rejects.toThrow() receiver.close() }) + + it('rejects the waiting waitForCode on close, and stops listening', async () => { + const receiver = await startLoopbackReceiver() + const code = receiver.waitForCode('s') + receiver.close() + await expect(code).rejects.toThrow('Sign-in cancelled.') + await expect(fetch(receiver.redirectUri)).rejects.toThrow() + }) + + it('stops at once with the reason of the signal', async () => { + const controller = new AbortController() + const reason = new Error('The user left.') + const receiver = await startLoopbackReceiver({ signal: controller.signal }) + const code = receiver.waitForCode('s') + controller.abort(reason) + await expect(code).rejects.toBe(reason) + await expect(fetch(receiver.redirectUri)).rejects.toThrow() + // A signal that aborted before the start opens no receiver. + await expect( + startLoopbackReceiver({ signal: controller.signal }), + ).rejects.toBe(reason) + }) + + it('does not let its deadline keep the process alive', async () => { + const timers = vi.spyOn(globalThis, 'setTimeout') + const receiver = await startLoopbackReceiver() + const code = receiver.waitForCode('s') + const deadline = timers.mock.results.at(-1)?.value + expect(deadline?.hasRef()).toBe(false) + receiver.close() + await expect(code).rejects.toThrow('Sign-in cancelled.') + }) }) diff --git a/packages/ai-mcp/src/connector.ts b/packages/ai-mcp/src/connector.ts index 728925fc0e..ab95324195 100644 --- a/packages/ai-mcp/src/connector.ts +++ b/packages/ai-mcp/src/connector.ts @@ -1,5 +1,5 @@ import { randomBytes } from 'node:crypto' -import { auth } from '@modelcontextprotocol/client' +import { auth, createFetchWithInit } from '@modelcontextprotocol/client' import { AuthRequiredError, defineCommand, @@ -333,9 +333,12 @@ export function mcpConnector(options: McpConnectorOptions) { commands: { [`connect:${id}`]: defineCommand({ description: `Sign in to ${label}`, - // The sign-in belongs to the user who runs the command. + // The sign-in belongs to the user who runs the command. A cancel + // or a session close stops it at once: the loopback stops + // listening, and the OAuth requests abort. run: async (_input, { run }) => { - const receiver = await startLoopbackReceiver() + const { signal } = run + const receiver = await startLoopbackReceiver({ signal }) const state = randomBytes(16).toString('base64url') try { const provider = credentialProvider(id, run.credentials, { @@ -349,10 +352,10 @@ export function mcpConnector(options: McpConnectorOptions) { url: authorizationUrl.href, }), }) - const fetchFn = options.fetch ? { fetchFn: options.fetch } : {} + const fetchFn = createFetchWithInit(options.fetch, { signal }) // A new sign-in has no tokens, so this always ends in a // redirect to the browser. - await auth(provider, { serverUrl: url, ...fetchFn }) + await auth(provider, { serverUrl: url, fetchFn }) // The error names the service: a sign-in can end minutes later. const { code, iss } = await receiver .waitForCode(state) @@ -367,7 +370,7 @@ export function mcpConnector(options: McpConnectorOptions) { serverUrl: url, authorizationCode: code, ...(iss ? { iss } : {}), - ...fetchFn, + fetchFn, }) } finally { receiver.close() diff --git a/packages/ai-mcp/tests/connector-edges.test.ts b/packages/ai-mcp/tests/connector-edges.test.ts index c0c02c3d8c..07e72c544d 100644 --- a/packages/ai-mcp/tests/connector-edges.test.ts +++ b/packages/ai-mcp/tests/connector-edges.test.ts @@ -362,6 +362,41 @@ describe('mcpConnector /connect', () => { ) }) + it('stops its OAuth requests when the command is cancelled', async () => { + const { server, host } = await setup() + let requested: () => void = () => {} + const firstRequest = new Promise((resolve) => { + requested = resolve + }) + const session = await host.open( + defineHarness({ + name: 'test/cancel', + adapter: recorder().adapter, + plugins: () => [ + mcpConnector({ + id: 'demo', + label: 'Demo', + url: server.url, + // The server never answers: only the signal ends a request. + fetch: (_input, init) => + new Promise((_, reject) => { + requested() + const signal = init?.signal + if (signal?.aborted) return reject(signal.reason) + signal?.addEventListener('abort', () => reject(signal.reason)) + }), + }), + ], + }), + { threadId: 't', principal: { id: 'user-1' } }, + ) + const connecting = session.command({ name: 'connect:demo' }) + await firstRequest + await connecting.cancel() + await expect(connecting).rejects.toBeDefined() + expect(connecting.status()).toBe('cancelled') + }) + it('passes the iss of the callback on, for a server that sends one (RFC 9207)', async () => { const { server, credentials, host } = await setup(undefined, { issParameter: true, From bfbf9be7de8290a0a8d4615df1d34fe97a0e61ad Mon Sep 17 00:00:00 2001 From: Alem Tuzlak Date: Sat, 17 Oct 2026 12:52:26 +0200 Subject: [PATCH 02/12] fix(ai-harness-cli): SIGINT and SIGTERM cancel the turn, then stop, in every mode --- packages/ai-harness-cli/src/commands.ts | 2 +- packages/ai-harness-cli/src/index.ts | 128 +++++--- packages/ai-harness-cli/src/lines.ts | 28 +- packages/ai-harness-cli/src/optional.ts | 1 + packages/ai-harness-cli/src/signals.ts | 83 +++++ .../ai-harness-cli/tests/child-view.test.ts | 6 +- .../ai-harness-cli/tests/cli-media.test.ts | 111 +------ packages/ai-harness-cli/tests/cli.test.ts | 107 +------ packages/ai-harness-cli/tests/helpers.ts | 117 +++++++ packages/ai-harness-cli/tests/signals.test.ts | 299 ++++++++++++++++++ 10 files changed, 619 insertions(+), 263 deletions(-) create mode 100644 packages/ai-harness-cli/src/signals.ts create mode 100644 packages/ai-harness-cli/tests/helpers.ts create mode 100644 packages/ai-harness-cli/tests/signals.test.ts diff --git a/packages/ai-harness-cli/src/commands.ts b/packages/ai-harness-cli/src/commands.ts index ddd90d332d..773f9c8cb5 100644 --- a/packages/ai-harness-cli/src/commands.ts +++ b/packages/ai-harness-cli/src/commands.ts @@ -8,7 +8,7 @@ export const HELP_TEXT = [ ' /agent [json] run an agent in the background', ' /config [key value] show or change a setting', ' /connect sign in to a connector (and /disconnect )', - ' /cancel cancel the running turn (or press Esc)', + ' /cancel cancel the running turn', ' /status show the session status', ' /exit quit', ].join('\n') diff --git a/packages/ai-harness-cli/src/index.ts b/packages/ai-harness-cli/src/index.ts index 1f07d97974..ba91ba39d0 100644 --- a/packages/ai-harness-cli/src/index.ts +++ b/packages/ai-harness-cli/src/index.ts @@ -6,6 +6,7 @@ import { runLines } from './lines' import { loadOptional } from './optional' import { EXIT, runPrint } from './print' import { createToken, serve } from './serve' +import { cancelTurn, stopOnSignals } from './signals' import type { AnyHarness, HarnessHost, @@ -27,6 +28,7 @@ import type { McpHarnessModule, McpStdioModule, } from './optional' +import type { SignalSource } from './signals' /** The error for a flag whose optional package is missing or too old. */ function needs(flag: string, name: string): Error { @@ -70,6 +72,36 @@ export interface RunCliOptions

{ * `view.send(text)` sends each `@path` file in `text`, as line mode does. */ ui?: (view: SessionView) => Promise | void + /** + * Where the CLI hears SIGINT and SIGTERM. Default: `process`. A test can + * give an `EventEmitter` and emit the signals. + * + * In every mode, a signal cancels the running turn. A second signal, or a + * signal while no turn runs, closes what the CLI opened, and `runCli` + * resolves to 130 (SIGINT) or 143 (SIGTERM). A host that the CLI made + * closes with `recoverable: true`, so a durable host runs the cut work + * again. Your own `host` stays open: the CLI closes only its session. + */ + signals?: SignalSource +} + +/** + * Show `ui` on a view of `session`. Resolves when the ui ends, or when a + * signal stops the CLI (`stopped` resolves). + */ +async function showUi( + session: HarnessSession, + ui: NonNullable, + stopped: Promise, +) { + const view = createSessionView(attachingSource(session)) + try { + await view.ready + await Promise.race([ui(view), stopped]) + } finally { + view.dispose() + } + return EXIT.ok } /** @@ -112,7 +144,8 @@ function attachingSource(session: HarnessSession) { } /** - * Run a harness from the terminal. Resolves to the process exit code. + * Run a harness from the terminal. Resolves to the process exit code: 130 + * or 143 when SIGINT or SIGTERM stops it (see `signals`). * * @example * ```ts @@ -157,6 +190,10 @@ export async function runCli

( }) const { principal } = options const approvals = args.yes ? 'auto' : 'ask' + // The session modes cancel the turn of their session. The other modes + // have no turn of their own, so a signal stops them at once. + let cancel = () => false + const stop = stopOnSignals(options.signals ?? process, () => cancel()) try { if (args.acp) { const acp = await loadOptional( @@ -165,8 +202,16 @@ export async function runCli

( ) if (!acp) throw needs('--acp', '@tanstack/ai-acp') const connection = acp.serveAcp({ host, harness, principal }) - await connection.closed - return EXIT.ok + const code = await Promise.race([ + connection.closed.then(() => EXIT.ok), + stop.stopped, + ]) + connection.close() + // ponytail: serveAcp reads process.stdin and does not stop at close, + // and a stdin that flows keeps the process alive. Pause it here until + // the ACP stdio stream stops reading at close. + if (stop.signal.aborted) process.stdin.pause() + return code } if (args.mcp) { @@ -198,13 +243,13 @@ export async function runCli

( filePaths: [process.cwd()], }) // stdout carries only MCP messages. The server stops when stdin ends. - const ended = new Promise((resolve) => { - process.stdin.once('end', resolve) + const ended = new Promise((resolve) => { + stdin.once('end', () => resolve(EXIT.ok)) }) const served = stdio.serveMCPStdio(server) - await ended + const code = await Promise.race([ended, stop.stopped]) await served.close() - return EXIT.ok + return code } if (args.dashboard) { @@ -232,12 +277,9 @@ export async function runCli

( ), }) stderr.write(`Connected to ${args.dashboard}. Press Ctrl+C to stop.\n`) - await new Promise((resolve) => { - process.once('SIGINT', resolve) - process.once('SIGTERM', resolve) - }) + const code = await stop.stopped connection.close() - return EXIT.ok + return code } if (args.serve) { @@ -254,49 +296,53 @@ export async function runCli

( }) stderr.write(`Serving ${harness.name} at ${server.url}\n`) if (!args.token && !env.HARNESS_TOKEN) stderr.write(`Token: ${token}\n`) - await new Promise((resolve) => { - process.once('SIGINT', resolve) - process.once('SIGTERM', resolve) - }) + const code = await stop.stopped await server.close() - return EXIT.ok + return code } const session = await host.open(harness, { threadId: args.thread, ...(principal ? { principal } : {}), }) + cancel = cancelTurn(session) // Relative to the working folder, as a relative --media-dir is. const mediaDir = args.mediaDir ?? defaultMediaDir(harness.name) - if (args.print !== undefined) { - return await runPrint(session, args.print, { - output: args.output, - stdout, - stderr, - mediaDir, - }) + const running = + args.print !== undefined + ? runPrint(session, args.print, { + output: args.output, + stdout, + stderr, + mediaDir, + }) + : stdin.isTTY && options.ui + ? showUi(session, options.ui, stop.stopped) + : runLines(session, stdin, stdout, { + openSignIns: Boolean(stdin.isTTY), + mediaDir, + signal: stop.signal, + onCtrlC: () => stop.receive('SIGINT'), + }).then(() => EXIT.ok) + const code = await Promise.race([running, stop.stopped]) + if (!stop.signal.aborted) { + // A signal that cancelled the turn of --print gives its exit code. + return code === EXIT.cancelled ? (stop.exitCode() ?? code) : code } - if (stdin.isTTY && options.ui) { - const view = createSessionView(attachingSource(session)) - try { - await view.ready - await options.ui(view) - } finally { - view.dispose() - } - } else { - await runLines(session, stdin, stdout, { - openSignIns: Boolean(stdin.isTTY), - mediaDir, - }) - } - return EXIT.ok + // The mode ends when its session closes. Work that the close cuts can + // fail, and the signal gives the exit code. + await session.close({ recoverable: true }) + await running.catch(() => {}) + return await stop.stopped } catch (error) { stderr.write(`${error instanceof Error ? error.message : String(error)}\n`) return EXIT.failed } finally { - // A host the caller gave stays open: the caller owns it. - if (!options.host) await host.close() + stop.dispose() + // A host the caller gave stays open: the caller owns it. After a + // signal, its running work stops as on a crash, so a durable host runs + // it again. + if (!options.host) await host.close({ recoverable: stop.signal.aborted }) } } diff --git a/packages/ai-harness-cli/src/lines.ts b/packages/ai-harness-cli/src/lines.ts index e23916cbee..b72f763001 100644 --- a/packages/ai-harness-cli/src/lines.ts +++ b/packages/ai-harness-cli/src/lines.ts @@ -31,16 +31,16 @@ function typedText(raw: string) { /** * Hide what the user types in a terminal while a secret question (a key) - * waits. Raw mode stops the echo. It also stops Ctrl+C, so this sends - * SIGINT itself. Does nothing for a pipe. + * waits. Raw mode stops the echo. It also makes Ctrl+C a typed character + * and not SIGINT, so this calls `onCtrlC`. Does nothing for a pipe. */ -function secretTyping(input: NodeJS.ReadableStream) { +function secretTyping(input: NodeJS.ReadableStream, onCtrlC: () => void) { const terminal = isTerminal(input) ? input : undefined let hidden = false const onKeys = (chunk: unknown) => { if (!String(chunk).includes('\u0003')) return show() - process.kill(process.pid, 'SIGINT') + onCtrlC() } const show = () => { if (!terminal || !hidden) return @@ -71,16 +71,24 @@ function secretTyping(input: NodeJS.ReadableStream) { * `openSignIns`: open sign-in links in the browser (for an interactive terminal). * `mediaDir`: the folder for the media a turn makes. Each new file is saved * there, with one `[image saved: ]` line. + * `signal`: stop reading lines when it aborts. + * `onCtrlC`: runs on Ctrl+C while typing is hidden, as the terminal then + * sends no SIGINT. */ export async function runLines( session: HarnessSession, input: NodeJS.ReadableStream, stdout: Output, - options: { openSignIns?: boolean; mediaDir: string }, + options: { + openSignIns?: boolean + mediaDir: string + signal: AbortSignal + onCtrlC: () => void + }, ) { const view = createSessionView(session) const saves: Array> = [] - const typing = secretTyping(input) + const typing = secretTyping(input, options.onCtrlC) const printer = createPrinter((text) => stdout.write(text), { onMedia: (part) => { const saved = saveMediaLine(session, part, options.mediaDir) @@ -99,10 +107,16 @@ export async function runLines( : () => {} const line = printer.line - const lines = createInterface({ input, crlfDelay: Infinity }) + const lines = createInterface({ + input, + crlfDelay: Infinity, + signal: options.signal, + }) const pendingLater: Array> = [] try { for await (const raw of lines) { + // Lines that came before the stop do not run. + if (options.signal.aborted) break const typed = typing.isHidden() ? typedText(raw) : raw typing.show() const snapshot = session.snapshot() diff --git a/packages/ai-harness-cli/src/optional.ts b/packages/ai-harness-cli/src/optional.ts index 5c4584ba97..003219fd37 100644 --- a/packages/ai-harness-cli/src/optional.ts +++ b/packages/ai-harness-cli/src/optional.ts @@ -63,6 +63,7 @@ export interface AcpAgentModule { principal?: P }) => { closed: Promise + close: () => void } } diff --git a/packages/ai-harness-cli/src/signals.ts b/packages/ai-harness-cli/src/signals.ts new file mode 100644 index 0000000000..25695d0507 --- /dev/null +++ b/packages/ai-harness-cli/src/signals.ts @@ -0,0 +1,83 @@ +import { EXIT } from './print' +import type { HarnessSession } from '@tanstack/ai-harness' + +/** The signals that stop the CLI. */ +export type StopSignal = 'SIGINT' | 'SIGTERM' + +/** + * Where the CLI hears SIGINT and SIGTERM. `process` is one. An + * `EventEmitter` is one too, so a test can emit the signals. + */ +export interface SignalSource { + on: (signal: StopSignal, handler: () => void) => unknown + off: (signal: StopSignal, handler: () => void) => unknown +} + +/** The exit code of a process that a signal stopped: 128 plus its number. */ +const EXIT_CODES = { SIGINT: EXIT.cancelled, SIGTERM: 143 } + +/** + * Listen for SIGINT and SIGTERM on `source` while a CLI mode runs. + * + * A signal first calls `cancel`, which stops the running turn and returns + * `true`. When `cancel` returns `false` (no turn runs, or a signal cancelled + * this turn already), the signal stops the mode: the listeners go, `signal` + * aborts, and `stopped` resolves with the exit code, 130 for SIGINT or 143 + * for SIGTERM. A signal after that gets the default action of `process`, so + * a close that does not end can still be killed. + */ +export function stopOnSignals(source: SignalSource, cancel: () => boolean) { + const controller = new AbortController() + let resolveStopped: (code: number) => void = () => {} + const stopped = new Promise((resolve) => { + resolveStopped = resolve + }) + let last: StopSignal | undefined + const receive = (signal: StopSignal) => { + if (controller.signal.aborted) return + last = signal + if (cancel()) return + dispose() + controller.abort() + resolveStopped(EXIT_CODES[signal]) + } + const onInterrupt = () => receive('SIGINT') + const onTerminate = () => receive('SIGTERM') + function dispose() { + source.off('SIGINT', onInterrupt) + source.off('SIGTERM', onTerminate) + } + source.on('SIGINT', onInterrupt) + source.on('SIGTERM', onTerminate) + return { + /** Aborts when a signal stops the mode. */ + signal: controller.signal, + /** Resolves with the exit code when a signal stops the mode. */ + stopped, + /** The exit code of the last signal. `undefined` before a signal. */ + exitCode: () => (last === undefined ? undefined : EXIT_CODES[last]), + /** Handle `signal` as if `source` sent it. */ + receive, + /** Stop listening. */ + dispose, + } +} + +/** + * The `cancel` of `stopOnSignals` for a session: cancel its running turn. + * Returns `false` when no turn runs, or when it cancelled this turn + * already, so a second signal stops a turn that does not end. + */ +export function cancelTurn(session: HarnessSession) { + let cancelled: string | undefined + return () => { + const turn = session + .snapshot() + .activeOperations.find((operation) => operation.kind === 'chat') + if (turn === undefined || turn.id === cancelled) return false + cancelled = turn.id + // When the cancel fails, the turn runs on, and the next signal closes. + void session.cancel({ operationId: turn.id }).catch(() => {}) + return true + } +} diff --git a/packages/ai-harness-cli/tests/child-view.test.ts b/packages/ai-harness-cli/tests/child-view.test.ts index c62dcd7568..996148a962 100644 --- a/packages/ai-harness-cli/tests/child-view.test.ts +++ b/packages/ai-harness-cli/tests/child-view.test.ts @@ -128,7 +128,11 @@ describe('line mode with a child agent', () => { session, input, { write: (text: string) => (out += text) }, - { mediaDir: join(tmpdir(), 'child-view-media') }, + { + mediaDir: join(tmpdir(), 'child-view-media'), + signal: new AbortController().signal, + onCtrlC: () => {}, + }, ) input.end('go\n') await running diff --git a/packages/ai-harness-cli/tests/cli-media.test.ts b/packages/ai-harness-cli/tests/cli-media.test.ts index 5cce9b5142..f7f97d032f 100644 --- a/packages/ai-harness-cli/tests/cli-media.test.ts +++ b/packages/ai-harness-cli/tests/cli-media.test.ts @@ -1,118 +1,13 @@ import { mkdtemp, readFile, readdir, rm, writeFile } from 'node:fs/promises' import { tmpdir } from 'node:os' import { join } from 'node:path' -import { Readable } from 'node:stream' import { afterAll, beforeAll, describe, expect, it, vi } from 'vitest' -import { EventType, defineAgent } from '@tanstack/ai' +import { defineAgent } from '@tanstack/ai' import { defineHarness } from '@tanstack/ai-harness' import { memoryPersistence } from '@tanstack/ai-persistence' import { EXIT, runCli } from '../src' -import type { AnyTextAdapter, ImageAdapter, StreamChunk } from '@tanstack/ai' - -// ponytail: `textTurn`, `toolTurn`, `scripted`, `capture`, and `stdinFrom` -// repeat the ones in cli.test.ts. Move them to a helpers file when a third -// test file needs them. -const now = () => Date.now() -const textTurn = (text: string): Array => [ - { type: EventType.RUN_STARTED, runId: 'r', threadId: 't', timestamp: now() }, - { - type: EventType.TEXT_MESSAGE_START, - messageId: 'm', - role: 'assistant', - timestamp: now(), - }, - { - type: EventType.TEXT_MESSAGE_CONTENT, - messageId: 'm', - delta: text, - timestamp: now(), - }, - { type: EventType.TEXT_MESSAGE_END, messageId: 'm', timestamp: now() }, - { - type: EventType.RUN_FINISHED, - runId: 'r', - threadId: 't', - timestamp: now(), - metadata: { tanstack: { finishReason: 'stop' } }, - }, -] -const toolTurn = (name: string, args: string): Array => [ - { type: EventType.RUN_STARTED, runId: 'r', threadId: 't', timestamp: now() }, - { - type: EventType.TOOL_CALL_START, - toolCallId: 'call_1', - toolCallName: name, - timestamp: now(), - }, - { - type: EventType.TOOL_CALL_ARGS, - toolCallId: 'call_1', - delta: args, - timestamp: now(), - }, - { type: EventType.TOOL_CALL_END, toolCallId: 'call_1', timestamp: now() }, - { - type: EventType.RUN_FINISHED, - runId: 'r', - threadId: 't', - timestamp: now(), - metadata: { tanstack: { finishReason: 'tool_calls' } }, - }, -] - -function scripted(turns: Array>) { - let call = 0 - const seen: Array> = [] - const adapter: AnyTextAdapter = { - kind: 'text', - name: 'mock', - model: 'test-model', - // `~types` holds types only. Its values are never read, so they are casts. - '~types': { - providerOptions: {} as Record, - inputModalities: ['text'] as readonly ['text'], - messageMetadataByModality: { - text: undefined as unknown, - image: undefined as unknown, - audio: undefined as unknown, - video: undefined as unknown, - document: undefined as unknown, - }, - toolCapabilities: [] as ReadonlyArray, - toolCallMetadata: undefined as unknown, - systemPromptMetadata: undefined as never, - }, - chatStream: (options) => { - seen.push(options.messages) - const chunks = turns[call] ?? textTurn('') - call += 1 - return (async function* () { - yield* chunks - })() - }, - structuredOutput: async () => ({ data: {}, rawText: '{}' }), - } - return { adapter, seen } -} - -function capture() { - let text = '' - return { - write: (chunk: string) => (text += chunk), - get text() { - return text - }, - } -} - -/** A stdin that reads `lines`, as a terminal or as a pipe. */ -function stdinFrom(lines: Array, { isTTY }: { isTTY: boolean }) { - // NodeJS.ReadStream is a TTY socket. A test cannot make one without a real - // terminal, and runCli only reads the lines and `isTTY`, so this casts. - return Object.assign(Readable.from(lines), { - isTTY, - }) as unknown as NodeJS.ReadStream -} +import { capture, scripted, stdinFrom, textTurn, toolTurn } from './helpers' +import type { ImageAdapter } from '@tanstack/ai' /** An image model that paints the bytes "hello". */ function imageAdapter(): ImageAdapter { diff --git a/packages/ai-harness-cli/tests/cli.test.ts b/packages/ai-harness-cli/tests/cli.test.ts index 20aced812d..9fa10f3565 100644 --- a/packages/ai-harness-cli/tests/cli.test.ts +++ b/packages/ai-harness-cli/tests/cli.test.ts @@ -1,7 +1,6 @@ import { memoryCredentialStore } from '@tanstack/ai/byok/server' -import { Readable } from 'node:stream' import { describe, expect, it, vi } from 'vitest' -import { EventType, defineAgent, toolDefinition } from '@tanstack/ai' +import { defineAgent, toolDefinition } from '@tanstack/ai' import { z } from 'zod' import { configOption, @@ -15,90 +14,7 @@ import { createMCPClient } from '@tanstack/ai-mcp' import { memoryPersistence } from '@tanstack/ai-persistence' import { EXIT, parseCliArgs, runCli } from '../src' import { serve } from '../src/serve' -import type { AnyTextAdapter, StreamChunk } from '@tanstack/ai' - -const now = () => Date.now() -const textTurn = (text: string): Array => [ - { type: EventType.RUN_STARTED, runId: 'r', threadId: 't', timestamp: now() }, - { - type: EventType.TEXT_MESSAGE_START, - messageId: 'm', - role: 'assistant', - timestamp: now(), - }, - { - type: EventType.TEXT_MESSAGE_CONTENT, - messageId: 'm', - delta: text, - timestamp: now(), - }, - { type: EventType.TEXT_MESSAGE_END, messageId: 'm', timestamp: now() }, - { - type: EventType.RUN_FINISHED, - runId: 'r', - threadId: 't', - timestamp: now(), - metadata: { tanstack: { finishReason: 'stop' } }, - }, -] -const toolTurn = (name: string, args: string): Array => [ - { type: EventType.RUN_STARTED, runId: 'r', threadId: 't', timestamp: now() }, - { - type: EventType.TOOL_CALL_START, - toolCallId: 'call_1', - toolCallName: name, - timestamp: now(), - }, - { - type: EventType.TOOL_CALL_ARGS, - toolCallId: 'call_1', - delta: args, - timestamp: now(), - }, - { type: EventType.TOOL_CALL_END, toolCallId: 'call_1', timestamp: now() }, - { - type: EventType.RUN_FINISHED, - runId: 'r', - threadId: 't', - timestamp: now(), - metadata: { tanstack: { finishReason: 'tool_calls' } }, - }, -] - -function scripted(turns: Array>) { - let call = 0 - const seen: Array> = [] - const adapter: AnyTextAdapter = { - kind: 'text', - name: 'mock', - model: 'test-model', - // `~types` holds types only. Its values are never read, so they are casts. - '~types': { - providerOptions: {} as Record, - inputModalities: ['text'] as readonly ['text'], - messageMetadataByModality: { - text: undefined as unknown, - image: undefined as unknown, - audio: undefined as unknown, - video: undefined as unknown, - document: undefined as unknown, - }, - toolCapabilities: [] as ReadonlyArray, - toolCallMetadata: undefined as unknown, - systemPromptMetadata: undefined as never, - }, - chatStream: (options) => { - seen.push(options.messages) - const chunks = turns[call] ?? textTurn('') - call += 1 - return (async function* () { - yield* chunks - })() - }, - structuredOutput: async () => ({ data: {}, rawText: '{}' }), - } - return { adapter, seen } -} +import { capture, scripted, stdinFrom, textTurn, toolTurn } from './helpers' /** `/connect acme` asks for a key as a secret question, then saves it. */ const acmeKeys = definePlugin({ @@ -124,16 +40,6 @@ const acmeKeys = definePlugin({ }), }) -function capture() { - let text = '' - return { - write: (chunk: string) => (text += chunk), - get text() { - return text - }, - } -} - const olga = { id: 'olga' } const cid = { id: 'cid' } const users = new Map([olga, cid].map((user) => [user.id, user])) @@ -150,15 +56,6 @@ function accessHost() { }) } -/** A stdin that reads `lines`, as a terminal or as a pipe. */ -function stdinFrom(lines: Array, { isTTY }: { isTTY: boolean }) { - // NodeJS.ReadStream is a TTY socket. A test cannot make one without a real - // terminal, and runCli only reads the lines and `isTTY`, so this casts. - return Object.assign(Readable.from(lines), { - isTTY, - }) as unknown as NodeJS.ReadStream -} - describe('parseCliArgs', () => { it('reads flags and rejects bad values', () => { expect(parseCliArgs(['-p', 'hi', '--output', 'ndjson'])).toMatchObject({ diff --git a/packages/ai-harness-cli/tests/helpers.ts b/packages/ai-harness-cli/tests/helpers.ts new file mode 100644 index 0000000000..0cb846d311 --- /dev/null +++ b/packages/ai-harness-cli/tests/helpers.ts @@ -0,0 +1,117 @@ +import { Readable } from 'node:stream' +import { EventType } from '@tanstack/ai' +import type { AnyTextAdapter, StreamChunk } from '@tanstack/ai' + +const now = () => Date.now() + +/** The chunks of a model call that answers `text`. */ +export const textTurn = (text: string): Array => [ + { type: EventType.RUN_STARTED, runId: 'r', threadId: 't', timestamp: now() }, + { + type: EventType.TEXT_MESSAGE_START, + messageId: 'm', + role: 'assistant', + timestamp: now(), + }, + { + type: EventType.TEXT_MESSAGE_CONTENT, + messageId: 'm', + delta: text, + timestamp: now(), + }, + { type: EventType.TEXT_MESSAGE_END, messageId: 'm', timestamp: now() }, + { + type: EventType.RUN_FINISHED, + runId: 'r', + threadId: 't', + timestamp: now(), + metadata: { tanstack: { finishReason: 'stop' } }, + }, +] + +/** The chunks of a model call that calls the tool `name` with `args`. */ +export const toolTurn = (name: string, args: string): Array => [ + { type: EventType.RUN_STARTED, runId: 'r', threadId: 't', timestamp: now() }, + { + type: EventType.TOOL_CALL_START, + toolCallId: 'call_1', + toolCallName: name, + timestamp: now(), + }, + { + type: EventType.TOOL_CALL_ARGS, + toolCallId: 'call_1', + delta: args, + timestamp: now(), + }, + { type: EventType.TOOL_CALL_END, toolCallId: 'call_1', timestamp: now() }, + { + type: EventType.RUN_FINISHED, + runId: 'r', + threadId: 't', + timestamp: now(), + metadata: { tanstack: { finishReason: 'tool_calls' } }, + }, +] + +/** + * 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>) { + let call = 0 + const seen: Array> = [] + const adapter: AnyTextAdapter = { + kind: 'text', + name: 'mock', + model: 'test-model', + // `~types` holds types only. Its values are never read, so they are casts. + '~types': { + providerOptions: {} as Record, + inputModalities: ['text'] as readonly ['text'], + messageMetadataByModality: { + text: undefined as unknown, + image: undefined as unknown, + audio: undefined as unknown, + video: undefined as unknown, + document: undefined as unknown, + }, + toolCapabilities: [] as ReadonlyArray, + toolCallMetadata: undefined as unknown, + systemPromptMetadata: undefined as never, + }, + chatStream: (options) => { + seen.push(options.messages) + const chunks = turns[call] ?? textTurn('') + call += 1 + return (async function* () { + yield* chunks + })() + }, + structuredOutput: async () => ({ data: {}, rawText: '{}' }), + } + return { adapter, seen } +} + +/** An stdout or stderr that keeps what is written. */ +export function capture() { + let text = '' + return { + write: (chunk: string) => (text += chunk), + get text() { + return text + }, + } +} + +/** `stream` as the stdin of runCli, as a terminal or as a pipe. */ +export function asStdin(stream: Readable, { isTTY }: { isTTY: boolean }) { + // NodeJS.ReadStream is a TTY socket. A test cannot make one without a real + // terminal, and runCli only reads the lines and `isTTY`, so this casts. + return Object.assign(stream, { isTTY }) as unknown as NodeJS.ReadStream +} + +/** A stdin that reads `lines`, as a terminal or as a pipe. */ +export function stdinFrom(lines: Array, options: { isTTY: boolean }) { + return asStdin(Readable.from(lines), options) +} diff --git a/packages/ai-harness-cli/tests/signals.test.ts b/packages/ai-harness-cli/tests/signals.test.ts new file mode 100644 index 0000000000..53dd19c566 --- /dev/null +++ b/packages/ai-harness-cli/tests/signals.test.ts @@ -0,0 +1,299 @@ +import { EventEmitter } from 'node:events' +import { PassThrough } from 'node:stream' +import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest' +import { fakeText } from '@tanstack/ai/testing' +import { + createHarnessHost, + defineCommand, + defineHarness, + definePlugin, +} from '@tanstack/ai-harness' +import { memoryLogStore, memoryPersistence } from '@tanstack/ai-persistence' +import { z } from 'zod' +import { EXIT, runCli } from '../src' +import { asStdin, capture, stdinFrom } from './helpers' +import type { FakeResponse } from '@tanstack/ai/testing' +import type { AnyHarness } from '@tanstack/ai-harness' +import type { RunCliOptions } from '../src' + +/** Resolves when `signal` aborts. */ +function untilAborted(signal: AbortSignal | undefined) { + return new Promise((resolve) => { + if (signal?.aborted) return resolve() + signal?.addEventListener('abort', () => resolve(), { once: true }) + }) +} + +/** + * A model whose first call waits until it is aborted. With `stuck`, it + * waits for `release` instead, as a model that does not stop at once. Its + * later calls answer with `answers`. `calls` has the signal of each call. + */ +function waitingModel( + answers: Array = [], + { stuck = false } = {}, +) { + let release = () => {} + const released = new Promise((resolve) => { + release = resolve + }) + const calls: Array = [] + const model = fakeText() + model.setResponses([ + async ({ request }) => { + const signal = request.request?.signal ?? request.abortController?.signal + calls.push(signal) + await (stuck ? released : untilAborted(signal)) + // The text stops at the abort, so the call does not finish. + return { text: 'cut' } + }, + ...answers, + ]) + return { model, calls, release } +} + +/** Start runCli with signals that the test emits. */ +function startCli(harness: AnyHarness, options: RunCliOptions) { + const signals = new EventEmitter() + const stdout = capture() + const stderr = capture() + const running = runCli(harness, { stdout, stderr, signals, ...options }) + return { running, signals, stdout, stderr } +} + +/** `/connect acme` asks for a key as a secret question. */ +const acmeKeys = definePlugin({ + name: 'test/acme-keys', + setup: (ctx) => ({ + commands: { + 'connect:acme': defineCommand({ + description: 'Save an Acme API key', + run: () => + ctx.session.ask({ + message: 'Paste your Acme API key', + schema: z.string(), + secret: true, + }), + }), + }, + }), +}) + +describe('signals in line mode', () => { + it('cancels the running turn, reads the next line, and closes on a signal while nothing runs', async () => { + const { model, calls } = waitingModel([{ text: 'Second answer.' }]) + const input = new PassThrough() + const cli = startCli( + defineHarness({ name: 'test/signal-lines', adapter: model }), + { argv: [], stdin: asStdin(input, { isTTY: false }) }, + ) + input.write('hello\n') + await vi.waitFor(() => expect(calls).toHaveLength(1)) + + cli.signals.emit('SIGINT') + await vi.waitFor(() => expect(calls[0]?.aborted).toBe(true)) + input.write('again\n') + await vi.waitFor(() => expect(cli.stdout.text).toContain('Second answer.')) + cli.signals.emit('SIGINT') + + expect(await cli.running).toBe(EXIT.cancelled) + }) + + it('closes on Ctrl+C while a secret question hides typing', async () => { + const rawModes: Array = [] + const input = new PassThrough() + const stdin = Object.assign(asStdin(input, { isTTY: true }), { + setRawMode: (mode: boolean) => rawModes.push(mode), + }) + const cli = startCli( + defineHarness({ + name: 'test/signal-secret', + adapter: fakeText(), + plugins: () => [acmeKeys], + }), + { argv: [], stdin }, + ) + input.write('/connect acme\n') + await vi.waitFor(() => expect(rawModes).toEqual([true])) + + // Raw mode makes Ctrl+C a typed character. + input.write('\u0003') + + expect(await cli.running).toBe(EXIT.cancelled) + expect(rawModes).toEqual([true, false]) + }) +}) + +describe('signals in print mode', () => { + it('cancels the turn on SIGTERM and exits 143', async () => { + const { model, calls } = waitingModel() + const cli = startCli( + defineHarness({ name: 'test/signal-print', adapter: model }), + { argv: ['-p', 'hi'] }, + ) + await vi.waitFor(() => expect(calls).toHaveLength(1)) + + cli.signals.emit('SIGTERM') + + expect(await cli.running).toBe(143) + expect(calls[0]?.aborted).toBe(true) + }) + + it('closes only its session of your host on a second signal while the cancelled turn runs on', async () => { + const { model, calls, release } = waitingModel([], { stuck: true }) + const harness = defineHarness({ name: 'test/signal-stuck', adapter: model }) + const host = createHarnessHost({ persistence: memoryPersistence() }) + const close = vi.spyOn(host, 'close') + const cli = startCli(harness, { argv: ['-p', 'hi'], host }) + await vi.waitFor(() => expect(calls).toHaveLength(1)) + cli.signals.emit('SIGINT') + await vi.waitFor(() => expect(calls[0]?.aborted).toBe(true)) + + cli.signals.emit('SIGINT') + release() + + expect(await cli.running).toBe(EXIT.cancelled) + expect(close).not.toHaveBeenCalled() + expect(host.isOpen(harness, { threadId: 'main' })).toBe(false) + await host.close() + }) +}) + +describe('signals in a custom ui', () => { + it('stops a ui that is still open, and closes its session of your host', async () => { + const harness = defineHarness({ + name: 'test/signal-ui', + adapter: fakeText(), + }) + const host = createHarnessHost({ persistence: memoryPersistence() }) + let shown = false + const cli = startCli(harness, { + argv: [], + stdin: stdinFrom([], { isTTY: true }), + host, + // A ui that the user never quits. + ui: () => { + shown = true + return new Promise(() => {}) + }, + }) + await vi.waitFor(() => expect(shown).toBe(true)) + + cli.signals.emit('SIGTERM') + + expect(await cli.running).toBe(143) + expect(host.isOpen(harness, { threadId: 'main' })).toBe(false) + await host.close() + }) +}) + +describe('signals in --serve', () => { + it('stops on SIGTERM with 143, and the next host runs the cut turn again', async () => { + const persistence = { + stores: { log: memoryLogStore(), runs: memoryPersistence().stores.runs }, + } + const { model, calls } = waitingModel() + const cli = startCli( + defineHarness({ name: 'test/signal-serve', adapter: model }), + { argv: ['--serve', '--port', '0', '--token', 'secret'], persistence }, + ) + const url = await vi.waitFor(() => { + const [, served] = /at (\S+)\n/.exec(cli.stderr.text) ?? [] + expect(served).toBeDefined() + return served + }) + const receipt = await fetch(`${url}/control`, { + method: 'POST', + headers: { + authorization: 'Bearer secret', + 'content-type': 'application/json', + }, + body: JSON.stringify({ + threadId: 't', + input: { op: 'prompt', message: 'go' }, + }), + }) + const { inputId } = await receipt.json() + await vi.waitFor(() => expect(calls).toHaveLength(1)) + + cli.signals.emit('SIGTERM') + + expect(await cli.running).toBe(143) + await expect(fetch(`${url}/capabilities`)).rejects.toThrow() + const recovered = fakeText() + recovered.setResponses([{ text: 'recovered' }]) + const next = createHarnessHost({ persistence }) + const session = await next.open( + defineHarness({ name: 'test/signal-serve', adapter: recovered }), + { threadId: 't' }, + ) + expect(await session.settled(inputId)).toMatchObject({ + outcome: 'completed', + }) + await next.close() + }) +}) + +describe('signals in the stdio modes', () => { + // The MCP and ACP servers read process.stdin. The test keeps them off it, + // and keeps the events they listen for. + let listened: Array = [] + beforeEach(() => { + listened = [] + vi.spyOn(process.stdin, 'on').mockImplementation( + (event: string | symbol) => { + listened.push(event) + return process.stdin + }, + ) + }) + afterEach(() => { + vi.restoreAllMocks() + }) + + /** Start `argv` with `options`, and wait until it reads stdio. */ + async function startServing( + argv: Array, + options: RunCliOptions = {}, + ) { + const harness = defineHarness({ + name: 'test/signal-stdio', + adapter: fakeText(), + }) + const cli = startCli(harness, { argv, ...options }) + await vi.waitFor(() => expect(listened).toContain('end')) + return cli + } + + it('ends --mcp when the stdin of runCli ends', async () => { + const stdin = stdinFrom([], { isTTY: false }) + const cli = await startServing(['--mcp'], { stdin }) + + // In a real run, the MCP server reads this stdin, so it ends. + stdin.resume() + + expect(await cli.running).toBe(EXIT.ok) + }) + + it('stops --mcp on SIGINT with 130', async () => { + const cli = await startServing(['--mcp'], { + stdin: stdinFrom([], { isTTY: false }), + }) + + cli.signals.emit('SIGINT') + + expect(await cli.running).toBe(EXIT.cancelled) + }) + + it('stops --acp on SIGINT with 130, and leaves your host open', async () => { + const host = createHarnessHost({ persistence: memoryPersistence() }) + const close = vi.spyOn(host, 'close') + const cli = await startServing(['--acp'], { host }) + + cli.signals.emit('SIGINT') + + expect(await cli.running).toBe(EXIT.cancelled) + expect(close).not.toHaveBeenCalled() + await host.close() + }) +}) From 712ca71e72296c258e312027d7257bdebf7371e3 Mon Sep 17 00:00:00 2001 From: Alem Tuzlak Date: Sat, 17 Oct 2026 12:54:32 +0200 Subject: [PATCH 03/12] fix(ai-harness): one retention policy bounds every per-session collection, the feed, and the log tail --- packages/ai-harness/src/feed.ts | 65 ++++-- packages/ai-harness/src/host.ts | 47 ++++ packages/ai-harness/src/index.ts | 2 + packages/ai-harness/src/log.ts | 62 +++-- packages/ai-harness/src/session/agents.ts | 11 +- packages/ai-harness/src/session/base.ts | 247 +++++++++++++++----- packages/ai-harness/src/session/helpers.ts | 14 ++ packages/ai-harness/src/session/inputs.ts | 42 ++-- packages/ai-harness/src/session/requests.ts | 2 +- packages/ai-harness/src/session/types.ts | 9 +- packages/ai-harness/tests/feed.test.ts | 54 ++++- packages/ai-harness/tests/log.test.ts | 115 +++++++-- packages/ai-harness/tests/retention.test.ts | 168 +++++++++++++ 13 files changed, 689 insertions(+), 149 deletions(-) create mode 100644 packages/ai-harness/tests/retention.test.ts diff --git a/packages/ai-harness/src/feed.ts b/packages/ai-harness/src/feed.ts index 15e3859c60..88269af53e 100644 --- a/packages/ai-harness/src/feed.ts +++ b/packages/ai-harness/src/feed.ts @@ -15,11 +15,6 @@ import type { UsageUpdate, } from './types' -// ponytail: in-memory log capped at MAX_EVENTS per session. A reader whose -// cursor fell out of the window resumes from the oldest kept event. Durable -// cross-process delivery goes through StreamDurability in the protocol layer. -const MAX_EVENTS = 10_000 - /** * The ordered event stream of one session. Every operation publishes here. A * host without a log uses {@link SessionFeed}. A durable host uses a feed @@ -172,16 +167,38 @@ function isUsageUpdate(value: unknown): value is UsageUpdate { * The in-memory {@link EventFeed}. Cursors are opaque to callers and increase * with each event. */ +// ponytail: it keeps the newest `eventBytes` of events (the JSON length of +// each). A reader whose cursor fell out of them resumes from the oldest kept +// event. Durable cross-process delivery goes through StreamDurability in the +// protocol layer. export class SessionFeed implements EventFeed { + /** The kept events. Their cursors are consecutive numbers. */ private readonly entries: Array = [] + /** The JSON length of each event in `entries`, in the same order. */ + private readonly sizes: Array = [] + /** The sum of `sizes`. */ + private bytes = 0 + private readonly eventBytes: number private sequence = 0 private readonly waiters = new Set<() => void>() private closed = false + /** `eventBytes`: how many bytes of events to keep. */ + constructor(options: { eventBytes: number }) { + this.eventBytes = options.eventBytes + } + publish(operationId: string, event: StreamChunk): void { this.sequence += 1 + const size = jsonLength(event) this.entries.push({ cursor: String(this.sequence), operationId, event }) - if (this.entries.length > MAX_EVENTS) this.entries.shift() + this.sizes.push(size) + this.bytes += size + // The newest event stays, so a live reader still gets it. + while (this.bytes > this.eventBytes && this.entries.length > 1) { + this.entries.shift() + this.bytes -= this.sizes.shift() ?? 0 + } this.wake() } @@ -189,6 +206,17 @@ export class SessionFeed implements EventFeed { return String(this.sequence) } + /** The kept events after the cursor `after` that `filter` accepts. */ + private after( + after: number, + filter: ((entry: SessionEvent) => boolean) | undefined, + ) { + const first = this.sequence - this.entries.length + 1 + return this.entries + .slice(Math.max(0, after + 1 - first)) + .filter((entry) => (filter ? filter(entry) : true)) + } + async *read(options: { from?: Cursor signal?: AbortSignal @@ -200,11 +228,7 @@ export class SessionFeed implements EventFeed { if (!Number.isFinite(after)) after = 0 while (true) { if (options.signal?.aborted) return - const next = this.entries.filter( - (entry) => - Number(entry.cursor) > after && - (options.filter ? options.filter(entry) : true), - ) + const next = this.after(after, options.filter) for (const entry of next) { if (options.signal?.aborted) return after = Number(entry.cursor) @@ -215,12 +239,7 @@ export class SessionFeed implements EventFeed { after = Math.max(after, this.sequence) } if (this.closed || options.until?.()) { - const rest = this.entries.filter( - (entry) => - Number(entry.cursor) > after && - (options.filter ? options.filter(entry) : true), - ) - for (const entry of rest) yield entry + for (const entry of this.after(after, options.filter)) yield entry return } // Events published while the reader handled `next` woke nobody, so @@ -254,3 +273,15 @@ export class SessionFeed implements EventFeed { }) } } + +/** + * The JSON length of `event`. A value that JSON cannot encode (a `BigInt` + * from a plugin) counts as 0, so the publish does not throw. + */ +function jsonLength(event: StreamChunk) { + try { + return JSON.stringify(event).length + } catch { + return 0 + } +} diff --git a/packages/ai-harness/src/host.ts b/packages/ai-harness/src/host.ts index 62b69163e4..f1cbc7a1cb 100644 --- a/packages/ai-harness/src/host.ts +++ b/packages/ai-harness/src/host.ts @@ -145,8 +145,47 @@ export interface HarnessHostOptions

{ * Leave it out for the defaults. */ ownership?: LeaseOptions + /** How the sessions of this host live. Leave it out for the defaults. */ + lifecycle?: HostLifecycleOptions } +/** How the sessions of a host live. */ +export interface HostLifecycleOptions { + /** How much each session keeps in memory. */ + retention?: RetentionLimits +} + +/** + * How much each session keeps in memory, so a long-lived session does not + * grow. Leave a field out for its default. + */ +export interface RetentionLimits { + /** + * How many settled operations (turns, commands, resets) a session keeps, + * with their receipts and settlements. Inputs with no turn (agent runs, + * messages to them, commands, cancels) have a window of the same size. + * Past it, `session.operation(id)` is `undefined` for the oldest, and + * `settled()` of its input reads the log (a durable host). Default 500. + */ + settledOperations?: number + /** + * How many settled agent runs a session keeps. Past it, `agentRun(id)` + * and `agentRuns()` leave out the oldest, except what the log has (a + * durable host). Default 500. + */ + 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. A durable host reads older events from the + * log. Default 8 000 000. + */ + eventBytes?: number +} + +/** {@link RetentionLimits} with every default filled in. */ +export type Retention = Required + export interface OpenSessionOptions

{ /** The conversation id. Opening the same id again returns the live session. */ threadId: string @@ -936,6 +975,12 @@ export function createHarnessHost

( } const index = sessionIndex(sessionStore, emit) const coalesceMs = options.log?.coalesceMs ?? 100 + const limits = options.lifecycle?.retention + const retention = { + settledOperations: limits?.settledOperations ?? 500, + settledAgentRuns: limits?.settledAgentRuns ?? 500, + eventBytes: limits?.eventBytes ?? 8_000_000, + } /** The shared log of each log with an open session, by its store log id. */ const logs = new Map>() /** The resolved shared logs, and the reduce that folds each one, for `logState`. */ @@ -971,6 +1016,7 @@ export function createHarnessHost

( rootThread, state, coalesceMs, + eventBytes: retention.eventBytes, ...(project ? { project } : {}), ...(reduce ? { reduce } : {}), ...(metadata ? { metadata } : {}), @@ -1040,6 +1086,7 @@ export function createHarnessHost

( media, hostId, index, + retention, ...(logStore ? { log: { diff --git a/packages/ai-harness/src/index.ts b/packages/ai-harness/src/index.ts index 8f0ad2795f..5fab83b942 100644 --- a/packages/ai-harness/src/index.ts +++ b/packages/ai-harness/src/index.ts @@ -100,6 +100,7 @@ export type { HarnessHostOptions, HarnessPersistence, HostEvent, + HostLifecycleOptions, HostRevokedEvent, HostSessionDeletedEvent, HostSessionEvent, @@ -107,6 +108,7 @@ export type { HostStatusEvent, OpenSessionOptions, ResumePendingOptions, + RetentionLimits, } from './host' export { AccessDeniedError, membersOnly, membersOnlyPolicyOf } from './access' export type { diff --git a/packages/ai-harness/src/log.ts b/packages/ai-harness/src/log.ts index 830aa909dd..a402b53b1d 100644 --- a/packages/ai-harness/src/log.ts +++ b/packages/ai-harness/src/log.ts @@ -606,9 +606,6 @@ const CHECKPOINT_NAMESPACE = 'harness:log-checkpoint' // ponytail: a checkpoint every 50 appends. Tune it when a cold fold is slow. const CHECKPOINT_EVERY = 50 const READ_PAGE = 500 -// ponytail: the in-memory event tail, like SessionFeed. Older cursors page -// from the store. -const MAX_TAIL = 10_000 interface CheckpointVersions { version?: string @@ -877,15 +874,26 @@ export class SharedLog { */ private readonly staged = new Map>() /** - * The tool calls whose result landed, by `\0`. - * ponytail: one short key per call, kept until the view closes. + * The tool calls whose result landed and is not in the transcript yet, by + * `\0`. */ private readonly resulted = new Set() private timer: ReturnType | undefined private scheduled = false private chain: Promise = Promise.resolve() private failure: unknown - private readonly tail: Array<{ thread: string; entry: SessionEvent }> = [] + /** + * The newest events, like SessionFeed: 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 readonly eventBytes: number /** Every event after this position is in `tail`. */ private tailFrom: number private readonly waiters = new Set<() => void>() @@ -905,6 +913,8 @@ export class SharedLog { rootThread: string state: SharedLogState coalesceMs: number + /** How many bytes of events the in-memory tail keeps. */ + eventBytes: number project?: ProjectOptions reduce?: ReduceOptions metadata?: MetadataStore @@ -917,6 +927,7 @@ export class SharedLog { this.state = options.state this.metadata = options.metadata this.coalesceMs = options.coalesceMs + this.eventBytes = options.eventBytes this.onIdle = options.onIdle ?? (() => {}) this.fold = { rootThread: options.rootThread, @@ -1083,10 +1094,10 @@ export class SharedLog { records: ReadonlyArray, toolCallId?: string, ) { - const callKey = `${thread}\u0000${toolCallId}` // A call whose result landed (a background job) waits for the commit. - const key = - toolCallId === undefined || this.resulted.has(callKey) ? thread : callKey + const isWaiting = + toolCallId !== undefined && !this.hasResult(thread, toolCallId) + const key = isWaiting ? `${thread}\u0000${toolCallId}` : thread const staged = this.staged.get(key) ?? [] staged.push( ...records.map((record) => normalize(this.stamp(record, thread))), @@ -1094,6 +1105,15 @@ export class SharedLog { this.staged.set(key, staged) } + /** The result of the call landed: it waits for a commit, or the transcript has it. */ + private hasResult(thread: string, toolCallId: string) { + if (this.resulted.has(`${thread}\u0000${toolCallId}`)) return true + // ponytail: a scan of the transcript, only when a call stages records. + return sessionOf(this.state, thread).messages.some( + (message) => message.role === 'tool' && message.toolCallId === toolCallId, + ) + } + private commit( thread: string, options: { @@ -1270,15 +1290,27 @@ export class SharedLog { private apply(entry: LogEntry) { foldLogEntry(this.state, entry, this.fold) + const { record } = entry + const thread = threadOf(record, this.rootThread) + // A result in the transcript does not wait for a commit any more. + if (isHarnessRecord(record) && record.type === 'harness.transcript') { + for (const message of record.add) { + if (message.role !== 'tool') continue + this.resulted.delete(`${thread}\u0000${message.toolCallId}`) + } + } const event = sessionEventOf(entry) if (!event) return - this.tail.push({ - thread: threadOf(entry.record, this.rootThread), - entry: event, - }) - if (this.tail.length > MAX_TAIL) { + // The fold has JSON values only, so this cannot throw. + const size = JSON.stringify(event.event).length + 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) this.tailFrom = Number(dropped.entry.cursor) + if (!dropped) break + this.tailBytes -= dropped.size + this.tailFrom = Number(dropped.entry.cursor) } } diff --git a/packages/ai-harness/src/session/agents.ts b/packages/ai-harness/src/session/agents.ts index e3d751090f..1632323885 100644 --- a/packages/ai-harness/src/session/agents.ts +++ b/packages/ai-harness/src/session/agents.ts @@ -197,8 +197,9 @@ export abstract class SessionAgents< /** * The agent run `operationId`: a run of this session, or on a durable * host a run in the log that ended, so a follow-up continues it after a - * restart. `undefined` for an unknown run, and for one that another host - * runs. + * restart. `undefined` for an unknown run, for one that another host + * runs, and on a host without a log for a run older than the window of + * settled agent runs (`lifecycle.retention.settledAgentRuns`). */ agentRun(operationId: string): AgentRun | undefined { const chain = this.chainOf.get(operationId) ?? this.loggedChain(operationId) @@ -208,8 +209,10 @@ export abstract class SessionAgents< } /** - * The agent runs of this session, oldest first. On a durable host it also - * lists the runs in the log, so it works after a restart. + * The agent runs of this session, oldest first: the running ones, and the + * settled ones in the window of `lifecycle.retention.settledAgentRuns`. On + * a durable host it also lists the runs in the log, so it works after a + * restart. */ agentRuns(): Array { const live = [...this.agentChains.values()].flatMap((chain) => diff --git a/packages/ai-harness/src/session/base.ts b/packages/ai-harness/src/session/base.ts index c8e408f99e..12e52eab9f 100644 --- a/packages/ai-harness/src/session/base.ts +++ b/packages/ai-harness/src/session/base.ts @@ -27,6 +27,7 @@ import { createMessageId, customEvent, cacheObject, + dropOldest, scopeOf, storedPrincipal, notOpenMessages, @@ -104,15 +105,121 @@ import type { import type { HarnessSession } from './harness-session' /** - * How many settled operations a session keeps, with their receipts and - * settlements. Older ones go, so a long-lived session does not grow. + * 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 + * operation. An input goes as one piece, so no map keeps a part of it. + * + * - A chat turn input goes with its operation, when that operation falls out + * of the window of settled operations. + * - Every other input (an agent run, a message to one, a command, a cancel, a + * setting) joins its own window of the same size when it ends. So these + * inputs never push a turn out. */ -// ponytail: a fixed window. After it, `operation(id)` is undefined, and a -// duplicate input or `settled()` falls back to the log (a durable host) or -// the inbox. The duplicate check of an input with no operation of its own (a -// command, a cancel) keeps its id forever: one string each. Make the window -// an option if a host needs more. -export const KEPT_SETTLED = 500 +// ponytail: after its window, `settled()` and a duplicate input fall back to +// the log (a durable host) or the inbox. An input whose work stops on a +// failed store write, and never ends, stays until the session closes. +export class InputRecords { + /** + * The input and its sender as JSON, by input id: the duplicate check. + * `recover()` skips these inputs. + */ + readonly admitted = new Map() + /** The inputs that joined each turn or agent run, by operation id. */ + readonly turnJoins = new Map>() + /** The receipt of each chat input this session answered, by input id. */ + private readonly receipts = new Map() + /** How the inputs of this session ended, for `settled()`, by input id. */ + private readonly settlements = new Map() + private readonly turns = new Map>() + private readonly inputOf = new Map() + /** The turn operation of each chat input, by input id. See `bind`. */ + readonly turnOperations: ReadonlyMap> = + this.turns + /** The input of each turn operation, by operation id. See `bind`. */ + readonly operationInputs: ReadonlyMap = this.inputOf + /** The inputs that ended and have no turn operation, oldest first. */ + private readonly ended = new Set() + private readonly limit: number + + /** `limit`: how many ended inputs with no turn operation to keep. */ + constructor(limit: number) { + this.limit = limit + } + + /** Link a turn operation and its input. The input then goes with it. */ + bind(inputId: string, operation: OperationImpl) { + this.turns.set(inputId, operation) + this.inputOf.set(operation.id, inputId) + this.ended.delete(inputId) + } + + receipt(inputId: string) { + return this.receipts.get(inputId) + } + + settlement(inputId: string) { + return this.settlements.get(inputId) + } + + /** Keep the receipt of a chat input, for a duplicate. A refusal ends it. */ + keep(receipt: Receipt) { + this.receipts.set(receipt.inputId, receipt) + if (receipt.status === 'rejected') this.end(receipt.inputId) + return receipt + } + + /** + * The session refused the input. It leaves the duplicate check: the store + * has it, so a retry with its id is still a duplicate. + */ + reject(receipt: Receipt) { + this.admitted.delete(receipt.inputId) + return this.keep(receipt) + } + + /** The input ended with `settlement`. */ + settle(settlement: InputSettlement) { + const { inputId } = settlement + // A turn input leaves the duplicate check with its operation: a client + // may send its id again until then. Other inputs leave now. + if (!this.turns.has(inputId)) this.admitted.delete(inputId) + this.settlements.set(inputId, settlement) + this.end(inputId) + } + + /** + * The input ended. One with no turn operation goes to the end of the + * window of ended inputs, and the oldest past the window are forgotten. + */ + end(inputId: string) { + if (this.turns.has(inputId)) return + // A command ends when it is stored, and again when it settles. + this.ended.delete(inputId) + this.ended.add(inputId) + for (const old of dropOldest(this.ended, this.limit)) this.forget(old) + } + + /** Forget what is kept for the settled operation `operationId`. */ + forgetOperation(operationId: string) { + this.turnJoins.delete(operationId) + const inputId = this.inputOf.get(operationId) + this.inputOf.delete(operationId) + // A newer operation of the input (a steer that runs as its own turn) + // keeps it. + if (inputId === undefined || this.turns.get(inputId)?.id !== operationId) { + return + } + this.turns.delete(inputId) + this.forget(inputId) + } + + private forget(inputId: string) { + this.admitted.delete(inputId) + this.receipts.delete(inputId) + this.settlements.delete(inputId) + this.ended.delete(inputId) + } +} /** * The fields of a {@link HarnessSession}, and the helpers every layer uses. The @@ -180,30 +287,39 @@ export abstract class SessionBase { /** Why the session log stopped taking writes. */ protected logFailure: Error | undefined protected readonly log: SessionDependencies['log'] - /** - * Inputs this session stored or recovered, by id, as JSON: the duplicate - * check. `recover()` skips them. - */ - protected readonly admitted = new Map() - /** The receipt of each chat input this session answered, by input id. */ - protected readonly receipts = new Map() - /** The turn operation of each chat input, by input id. */ - protected readonly turnOperations = new Map< - string, - OperationImpl - >() - /** The input of each turn operation, by operation id. */ - protected readonly operationInputs = new Map() - /** The inputs that joined each turn, by operation id. */ - protected readonly turnJoins = new Map>() - /** How the chat inputs of this session ended, for `settled()`. */ - protected readonly settlements = new Map() + /** How much this session keeps in memory. */ + protected readonly retention: SessionDependencies['retention'] + /** What this session keeps for its inputs. */ + protected readonly inputRecords: InputRecords protected readonly agentRegistry = new AgentRegistry() protected readonly operations = new Map>() - /** The ids of the settled operations in `operations`, oldest first. */ - protected readonly settledIds: Array = [] - /** The ids of forgotten operations, so a client `runId` stays unique. */ + /** The settled operations in `operations` that are not agent runs, oldest first. */ + protected readonly settledIds = new Set() + /** The settled agent runs in `operations`, oldest first. */ + protected readonly settledRuns = new Set() + /** + * The ids of forgotten operations, oldest first, so a client `runId` does + * not take one. + */ + // ponytail: as many ids as both windows keep. Past that, only + // `stores.runs` refuses a reused run id. protected readonly retiredIds = new Set() + /** See {@link InputRecords.admitted}. */ + protected get admitted() { + return this.inputRecords.admitted + } + /** See {@link InputRecords.turnOperations}. */ + protected get turnOperations() { + return this.inputRecords.turnOperations + } + /** See {@link InputRecords.operationInputs}. */ + protected get operationInputs() { + return this.inputRecords.operationInputs + } + /** See {@link InputRecords.turnJoins}. */ + protected get turnJoins() { + return this.inputRecords.turnJoins + } /** The agent chains of this session, by the input id of the first run. */ protected readonly agentChains = new Map() /** The chain of each agent run operation, by operation id. */ @@ -368,8 +484,14 @@ export abstract class SessionBase { options: this.harness.media, }) this.principal = deps.principal + this.retention = deps.retention + this.inputRecords = new InputRecords(deps.retention.settledOperations) this.notifyHost = deps.onEvent - this.feed = tapFeed(new SessionFeed(), deps.onEvent, this.eventListeners) + this.feed = tapFeed( + new SessionFeed({ eventBytes: deps.retention.eventBytes }), + deps.onEvent, + this.eventListeners, + ) this.onClose = deps.onClose this.hostId = deps.hostId this.index = deps.index @@ -434,49 +556,58 @@ export abstract class SessionBase { } /** - * An operation of this session by id: a running one, an agent run, or one - * of the last 500 other operations that settled. + * An operation of this session by id: a running one, or one of the last + * operations that settled. The host option `lifecycle.retention` sets how + * many settled operations and agent runs a session keeps (default 500 + * each). */ operation(id: string): Operation | undefined { return this.operations.get(id) } - /** Keep `operation` in `operations`. Once settled, it joins the window. */ - protected track(operation: OperationImpl): void { + /** + * Keep `operation` in `operations`. Once settled, it joins its window: the + * agent runs, or the other operations. + */ + protected track(operation: OperationImpl) { this.operations.set(operation.id, operation) - // ponytail: agent runs stay, so `agentRuns()` lists every run and a - // message to an ended run starts its follow-up. Memory grows with the - // number of agent runs. Bound them too if a session runs thousands. - if (operation.kind === 'agent') return + const isRun = operation.kind === 'agent' + const window = isRun ? this.settledRuns : this.settledIds + const limit = isRun + ? this.retention.settledAgentRuns + : this.retention.settledOperations const settled = () => { - this.settledIds.push(operation.id) - const old = - this.settledIds.length > KEPT_SETTLED && this.settledIds.shift() - if (old) this.forget(old) + window.add(operation.id) + for (const old of dropOldest(window, limit)) this.forget(old) } void operation.then(settled, settled) } /** Drop what this session keeps for the settled operation `id`. */ - protected forget(id: string): void { + protected forget(id: string) { + const chain = this.chainOf.get(id) this.operations.delete(id) - // ponytail: only the id stays, so a reused `runId` is still refused. + this.chainOf.delete(id) this.retiredIds.add(id) - this.turnJoins.delete(id) - const inputId = this.operationInputs.get(id) - this.operationInputs.delete(id) - if (inputId !== undefined && this.turnOperations.get(inputId)?.id === id) { - this.turnOperations.delete(inputId) - this.admitted.delete(inputId) - this.receipts.delete(inputId) - this.settlements.delete(inputId) + const { settledOperations, settledAgentRuns } = this.retention + dropOldest(this.retiredIds, settledOperations + settledAgentRuns) + this.inputRecords.forgetOperation(id) + if (chain) this.forgetRun(chain, id) + } + + /** + * Drop the run `id` from `chain`. A chain with no run left and no + * follow-up that waits goes too. + */ + private forgetRun(chain: AgentChain, id: string) { + const at = chain.runs.findIndex((run) => run.operation.id === id) + if (at >= 0) chain.runs.splice(at, 1) + if (chain.runs.length > 0 || chain.followUps.length > 0) return + if (this.agentChains.get(chain.runInputId) === chain) { + this.agentChains.delete(chain.runInputId) } - // Steers and agent inputs have no operation of their own here. - for (const map of [this.receipts, this.settlements]) { - for (const key of map.keys()) { - if (map.size <= KEPT_SETTLED) break - map.delete(key) - } + for (const parent of this.agentChains.values()) { + parent.children.delete(chain) } } diff --git a/packages/ai-harness/src/session/helpers.ts b/packages/ai-harness/src/session/helpers.ts index 8ba1dd2cf9..8dd38543ab 100644 --- a/packages/ai-harness/src/session/helpers.ts +++ b/packages/ai-harness/src/session/helpers.ts @@ -573,6 +573,20 @@ export const WORK_OPS = new Set([ /** The inputs that change a waiting input. */ export const CONTROL_OPS = new Set(['cancelInput', 'setDelivery']) +/** + * Delete the oldest items of `set` (the first added) until it has at most + * `limit` items. Returns the deleted items, oldest first. + */ +export function dropOldest(set: Set, limit: number) { + const dropped: Array = [] + for (const item of set) { + if (set.size <= limit) break + set.delete(item) + dropped.push(item) + } + return dropped +} + /** Why a turn stops when its input passes its time limit. */ export const TIMEOUT_REASON = 'harness:input-timeout' diff --git a/packages/ai-harness/src/session/inputs.ts b/packages/ai-harness/src/session/inputs.ts index 35606a3088..d2c4608f66 100644 --- a/packages/ai-harness/src/session/inputs.ts +++ b/packages/ai-harness/src/session/inputs.ts @@ -11,8 +11,9 @@ import { WORK_OPS, lastAssistantText, checkHostRecords, + dropOldest, } from './helpers' -import { KEPT_SETTLED, SessionBase } from './base' +import { SessionBase } from './base' import type { AccessRun } from '../access' import type { OperationImpl } from '../operation' import type { LogRecord } from '@tanstack/ai-persistence' @@ -426,6 +427,8 @@ export abstract class SessionInputs< * the other. */ protected async applyControl(inputId: string, change: HarnessRecord) { + // Also an input that recovery took: it has no other end. + this.inputRecords.end(inputId) if (!this.writer) { await this.inbox.markApplied({ threadId: this.storeThreadId, @@ -509,8 +512,11 @@ export abstract class SessionInputs< }) // The inbox had the id before this session: a restart. if (stored.createdAt !== at || stored.status !== 'pending') { - if (inputKey(stored.input, stored.principal) === payload) + if (inputKey(stored.input, stored.principal) === payload) { + // Nothing runs again, so the input ends here. + this.inputRecords.end(inputId) return 'duplicate' + } this.admitted.delete(inputId) return 'conflict' } @@ -520,6 +526,8 @@ export abstract class SessionInputs< this.admitted.delete(inputId) throw error } + // An input with no work to run (a cancel, a setting) ends once stored. + if (!WORK_OPS.has(input.op)) this.inputRecords.end(inputId) this.feed.publish( 'session', customEvent(HARNESS_EVENTS.inputAccepted, { inputId, op: input.op }), @@ -544,10 +552,7 @@ export abstract class SessionInputs< // window of settled operations. if (CHAT_OPS.has(input.op)) { this.forwarded.add(inputId) - for (const old of this.forwarded) { - if (this.forwarded.size <= KEPT_SETTLED) break - this.forwarded.delete(old) - } + dropOldest(this.forwarded, this.retention.settledOperations) } await inbox.append({ inputId, @@ -628,8 +633,7 @@ export abstract class SessionInputs< .catch((error: unknown) => this.warnSession('harness:inbox', error)) } const receipt: Receipt = { inputId, status: 'rejected', reason } - this.admitted.delete(inputId) - this.receipts.set(inputId, receipt) + this.inputRecords.reject(receipt) this.feed.publish( 'session', customEvent(HARNESS_EVENTS.inputRejected, { inputId, reason }), @@ -653,11 +657,7 @@ export abstract class SessionInputs< ) } for (const settlement of settlements) { - // A turn input leaves the set with its operation (`forget`): a client - // may send its id again until then. Other inputs leave now. - if (!this.turnOperations.has(settlement.inputId)) - this.admitted.delete(settlement.inputId) - this.settlements.set(settlement.inputId, settlement) + this.inputRecords.settle(settlement) this.feed.publish( settlement.operationId ?? 'session', customEvent(HARNESS_EVENTS.inputSettled, { ...settlement }), @@ -676,9 +676,9 @@ export abstract class SessionInputs< status: 'rejected' as const, reason: rejectedReason(logged), } - : this.receipts.get(inputId) + : this.inputRecords.receipt(inputId) if (refused?.status === 'rejected') throw new InputRejectedError(refused) - return this.settlements.get(inputId) + return this.inputRecords.settlement(inputId) } protected isChatInput(inputId: string) { @@ -694,7 +694,7 @@ export abstract class SessionInputs< ) return ( this.turnOperations.has(inputId) || - this.receipts.has(inputId) || + this.inputRecords.receipt(inputId) !== undefined || this.forwarded.has(inputId) ) } @@ -704,14 +704,12 @@ export abstract class SessionInputs< inputId: string, operation: OperationImpl, ) { - this.turnOperations.set(inputId, operation) - this.operationInputs.set(operation.id, inputId) + this.inputRecords.bind(inputId, operation) } /** Keep a chat input's receipt, for a duplicate of the input. */ - protected keep(receipt: Receipt): Receipt { - this.receipts.set(receipt.inputId, receipt) - return receipt + protected keep(receipt: Receipt) { + return this.inputRecords.keep(receipt) } /** Resolve a turn's receipt, and keep it for a duplicate. */ @@ -736,7 +734,7 @@ export abstract class SessionInputs< return { inputId, status: 'rejected', reason: admission } } if (admission === 'forwarded') return { inputId, status: 'queued' } - const kept = this.receipts.get(inputId) + const kept = this.inputRecords.receipt(inputId) if (kept) return kept const operationId = this.writer?.state.inputs.get(inputId)?.operationId return { diff --git a/packages/ai-harness/src/session/requests.ts b/packages/ai-harness/src/session/requests.ts index de9d6aeb34..1bd80330af 100644 --- a/packages/ai-harness/src/session/requests.ts +++ b/packages/ai-harness/src/session/requests.ts @@ -1078,7 +1078,6 @@ export abstract class SessionRequests< return } operation.resolveReceipt(receipt) - this.followers.add(operation) if (!this.writer) { // Without a log, the result of an input from before a restart, or of // one the session no longer keeps, is gone. @@ -1090,6 +1089,7 @@ export abstract class SessionRequests< ) return } + this.followers.add(operation) void this.settled(inputId) .then( (settlement) => this.settleOperation(operation, settlement), diff --git a/packages/ai-harness/src/session/types.ts b/packages/ai-harness/src/session/types.ts index 6234341112..3fa0986a67 100644 --- a/packages/ai-harness/src/session/types.ts +++ b/packages/ai-harness/src/session/types.ts @@ -24,7 +24,7 @@ import type { ConfigOption } from '../config' import type { SessionUsage } from '../usage' import type { AgentInputOf, AgentResultOf, AnyAgent } from '../agents' import type { AnyHarness, HarnessAgentsOf } from '../define' -import type { HarnessPersistence, SessionIndexWriter } from '../host' +import type { HarnessPersistence, Retention, SessionIndexWriter } from '../host' import type { PluginLifetime } from '../plugins' import type { AgentRun, @@ -221,6 +221,8 @@ export interface SessionDependencies { hostId: string /** Writes the session index entry of the thread. */ index: SessionIndexWriter + /** How much the session keeps in memory. */ + retention: Retention /** The session log of a durable host. */ log?: { store: LogStore @@ -369,7 +371,10 @@ export interface AgentChain { budget: SubagentBudget /** The newest run that started. */ current: AgentRunEntry - /** Every run that started, oldest first. */ + /** + * Every run that started and that the session still keeps, oldest first. + * A run leaves it when it falls out of the window of settled agent runs. + */ runs: Array /** Steers that wait for the current run's next model call. */ steers: Array diff --git a/packages/ai-harness/tests/feed.test.ts b/packages/ai-harness/tests/feed.test.ts index db5cc8febb..c5b2dc5b37 100644 --- a/packages/ai-harness/tests/feed.test.ts +++ b/packages/ai-harness/tests/feed.test.ts @@ -2,17 +2,31 @@ import { describe, expect, it } from 'vitest' import { EventType } from '@tanstack/ai' import { SessionFeed } from '../src/feed' import type { StreamChunk } from '@tanstack/ai' +import type { SessionEvent } from '../src/types' +/** Room for every event of a test. */ +const EVENT_BYTES = 1_000_000 + +/** A custom event. Names of one length give events of one JSON length. */ const custom = (name: string): StreamChunk => ({ type: EventType.CUSTOM, name, value: {}, - timestamp: Date.now(), + timestamp: 0, }) +/** The names of the custom events that `events` yields, until it ends. */ +async function namesOf(events: AsyncIterable) { + const names: Array = [] + for await (const entry of events) { + names.push(entry.event.type === EventType.CUSTOM ? entry.event.name : '') + } + return names +} + describe('SessionFeed', () => { it('delivers events published while a reader handles earlier ones', async () => { - const feed = new SessionFeed() + const feed = new SessionFeed({ eventBytes: EVENT_BYTES }) const controller = new AbortController() const seen: Array = [] const reading = (async () => { @@ -34,24 +48,42 @@ describe('SessionFeed', () => { }) it('skips filtered events and replays from a cursor', async () => { - const feed = new SessionFeed() + const feed = new SessionFeed({ eventBytes: EVENT_BYTES }) feed.publish('a', custom('one')) feed.publish('b', custom('two')) feed.publish('a', custom('three')) feed.close() - const names: Array = [] - for await (const entry of feed.read({ - from: '1', - filter: (item) => item.operationId === 'a', - })) { - names.push(entry.event.type === EventType.CUSTOM ? entry.event.name : '') - } + const names = await namesOf( + feed.read({ from: '1', filter: (item) => item.operationId === 'a' }), + ) expect(names).toEqual(['three']) expect(feed.head()).toBe('3') }) + it('keeps the newest events that fit in eventBytes, and starts an older cursor at the oldest kept', async () => { + const size = JSON.stringify(custom('e1')).length + const feed = new SessionFeed({ 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: '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 }) + feed.publish('op', custom('e1')) + feed.publish('op', custom('e2')) + feed.close() + + expect(await namesOf(feed.read({ from: '0' }))).toEqual(['e2']) + }) + it('ends a read when wake() finds until() true, with no new event', async () => { - const feed = new SessionFeed() + const feed = new SessionFeed({ eventBytes: EVENT_BYTES }) let done = false const reading = (async () => { let count = 0 diff --git a/packages/ai-harness/tests/log.test.ts b/packages/ai-harness/tests/log.test.ts index 5da87c7804..c711463fe9 100644 --- a/packages/ai-harness/tests/log.test.ts +++ b/packages/ai-harness/tests/log.test.ts @@ -30,11 +30,14 @@ import type { SessionEvent } from '../src/types' const THREAD = 't1' /** The id of the log of `THREAD` in the store. */ const LOG = logIdOf({ logId: THREAD }) +/** The tail of a test log: room for every event of a test. */ +const EVENT_BYTES = 1_000_000 function newWriter( store: LogStore, options: { coalesceMs?: number + eventBytes?: number project?: ProjectOptions metadata?: MetadataStore onFailure?: (error: unknown) => void @@ -47,6 +50,7 @@ function newWriter( rootThread: THREAD, state: options.state ?? emptySharedLogState(), coalesceMs: options.coalesceMs ?? 0, + eventBytes: options.eventBytes ?? EVENT_BYTES, ...(options.project ? { project: options.project } : {}), ...(options.metadata ? { metadata: options.metadata } : {}), }) @@ -276,6 +280,23 @@ describe('tool call records', () => { toolCallId, message: { role: 'tool', toolCallId, content: 'ok' }, }) + /** A turn that called `pay` as `call-1`, with the call's result. */ + const paid: Array = [ + user('u1', 'pay'), + { + id: 'a1', + role: 'assistant', + content: '', + toolCalls: [ + { + id: 'call-1', + type: 'function', + function: { name: 'pay', arguments: '{}' }, + }, + ], + }, + { role: 'tool', toolCallId: 'call-1', content: 'ok' }, + ] it('appends the staged records of a call with its result, not with another commit', async () => { const { store, batches } = recordingStore() @@ -293,18 +314,6 @@ describe('tool call records', () => { it('forgets the result, start, and steps of a call once its result is in the transcript', async () => { const writer = newWriter(memoryLogStore()) - const call: ModelMessage = { - id: 'a1', - role: 'assistant', - content: '', - toolCalls: [ - { - id: 'call-1', - type: 'function', - function: { name: 'pay', arguments: '{}' }, - }, - ], - } await writer.append([ { type: 'harness.tool.started', @@ -322,18 +331,80 @@ describe('tool call records', () => { ]) expect(writer.state.toolResults.size).toBe(1) - await writer.commit({ - messages: [ - user('u1', 'pay'), - call, - { role: 'tool', toolCallId: 'call-1', content: 'ok' }, - ], - }) + await writer.commit({ messages: paid }) expect(writer.state.toolResults.size).toBe(0) expect(writer.state.started.size).toBe(0) expect(writer.state.steps.size).toBe(0) }) + + it('appends what a call stages after its result is in the transcript with the next commit, and keeps no key for it', async () => { + const { store, batches } = recordingStore() + const log = new SharedLog({ + store, + logId: LOG, + rootThread: THREAD, + state: emptySharedLogState(), + coalesceMs: 0, + eventBytes: EVENT_BYTES, + }) + const writer = log.view(THREAD, () => {}) + await writer.append([result('call-1')]) + await writer.commit({ messages: paid }) + + // A background job of the call stages a record after that. + writer.stage([{ type: 'app.late' }], 'call-1') + await writer.commit({ messages: [...paid, user('u2', 'next')] }) + + expect(batches).toEqual([ + ['harness.tool.result'], + ['harness.transcript'], + ['harness.transcript', 'app.late'], + ]) + // The set of calls that wait for a commit has no public read. + const resulted: unknown = Reflect.get(log, 'resulted') + expect(resulted).toEqual(new Set()) + }) +}) + +describe('the event tail', () => { + const custom = (name: string): StreamChunk => ({ + type: EventType.CUSTOM, + name, + value: {}, + timestamp: 0, + }) + + it('keeps eventBytes of events in memory, and reads older ones from the store', async () => { + const inner = memoryLogStore() + let reads = 0 + const store = wrapLogStore(inner, { + read: (options) => { + reads += 1 + return inner.read(options) + }, + }) + const size = JSON.stringify(custom('e1')).length + const writer = newWriter(store, { eventBytes: size * 2 }) + for (const name of ['e1', 'e2', 'e3', 'e4']) { + writer.publish('op-1', custom(name)) + } + await writer.flush() + reads = 0 + const namesFrom = async (from: string) => { + const names: Array = [] + for await (const { event } of writer.read({ from, until: () => true })) { + names.push(event.type === EventType.CUSTOM ? event.name : '') + } + return names + } + + // Events 3 and 4 are in memory. Events 1 and 2 come from the store. + expect(await namesFrom('0')).toEqual(['e1', 'e2', 'e3', 'e4']) + expect(reads).toBe(1) + expect(await namesFrom('2')).toEqual(['e3', 'e4']) + expect(reads).toBe(1) + }) }) describe('input fold', () => { @@ -388,6 +459,7 @@ describe('a closed view', () => { rootThread: THREAD, state: emptySharedLogState(), coalesceMs: 0, + eventBytes: EVENT_BYTES, }) const closed = log.view('a', () => {}) // Another session keeps the log open. @@ -884,6 +956,7 @@ describe('fold checkpoints', () => { rootThread: THREAD, state: emptySharedLogState(reduce), coalesceMs: 0, + eventBytes: EVENT_BYTES, reduce, metadata, }) @@ -1107,6 +1180,7 @@ describe('a log shared by sessions', () => { rootThread: 'log-1', state: emptySharedLogState(), coalesceMs: 0, + eventBytes: EVENT_BYTES, }) return { log, @@ -1183,6 +1257,7 @@ describe('a log shared by sessions', () => { rootThread: 'log-1', state: emptySharedLogState(reduce), coalesceMs: 0, + eventBytes: EVENT_BYTES, reduce, }) const root = log.view('log-1', () => {}) @@ -1248,6 +1323,7 @@ describe('a log shared by sessions', () => { rootThread: 'log-1', state: emptySharedLogState(), coalesceMs: 0, + eventBytes: EVENT_BYTES, onIdle, }) const rootFailure = vi.fn() @@ -1312,6 +1388,7 @@ describe('a log shared by sessions', () => { rootThread: THREAD, state: emptySharedLogState(), coalesceMs: 0, + eventBytes: EVENT_BYTES, }) const view = log.view(THREAD, () => {}) void view.append([{ type: 'app.last' }]) diff --git a/packages/ai-harness/tests/retention.test.ts b/packages/ai-harness/tests/retention.test.ts new file mode 100644 index 0000000000..9792913706 --- /dev/null +++ b/packages/ai-harness/tests/retention.test.ts @@ -0,0 +1,168 @@ +import { describe, expect, it } from 'vitest' +import { defineAgent } from '@tanstack/ai' +import { memoryPersistence } from '@tanstack/ai-persistence' +import { createHarnessHost, defineHarness } from '../src' +import { mockAdapter, text } from './helpers' +import type { HarnessPersistence, RetentionLimits, SessionEvent } from '../src' + +/** + * A session on a host without a log, with `retention`. The model answers + * `done` to every call, and the agent `writer` runs one model call. + */ +async function openWith( + retention: RetentionLimits, + persistence: HarnessPersistence = memoryPersistence(), +) { + const host = createHarnessHost({ persistence, lifecycle: { retention } }) + const { adapter } = mockAdapter(() => text('done')) + const writer = defineAgent({ + name: 'writer', + description: 'Writes', + run: (ctx) => ctx.chat({ adapter }), + }) + const harness = defineHarness({ + name: 'test/retention', + adapter, + agents: [writer], + }) + const session = await host.open(harness, { threadId: 't1' }) + return { host, session } +} + +/** + * The size of the private set `name` of `owner`. A leak in these sets has no + * public signal, so the test reads the set by name. + */ +function privateSetSize(owner: object, name: string) { + const value: unknown = Reflect.get(owner, name) + return value instanceof Set ? value.size : undefined +} + +async function collect(events: AsyncIterable) { + const kept: Array = [] + for await (const entry of events) kept.push(entry) + return kept +} + +describe('retention', () => { + it('keeps the settlement of a turn in the window while many agent runs settle', async () => { + const { host, session } = await openWith({ + settledOperations: 5, + settledAgentRuns: 2, + }) + for (const inputId of ['t-0', 't-1', 't-2']) { + await session.send('Hi', { inputId }) + } + for (let index = 0; index < 12; index += 1) { + await session.agents.writer.run() + } + for (const inputId of ['t-3', 't-4']) { + await session.send('Hi', { inputId }) + } + + await expect(session.settled('t-0')).resolves.toMatchObject({ + inputId: 't-0', + outcome: 'completed', + }) + // A sixth turn pushes the oldest turn out of the window. + await session.send('Hi', { inputId: 't-5' }) + await expect(session.settled('t-0')).rejects.toThrow('knows no chat input') + await host.close() + }) + + it('keeps the receipt of a turn while many inputs with no turn end', async () => { + const { host, session } = await openWith({ settledOperations: 3 }) + const first = session.send('Hi', { inputId: 't-0' }) + await first + for (let index = 0; index < 10; index += 1) { + // Nothing runs, so each cancel is refused. + await session.cancel() + await session.configure({ settings: { instructions: `Rule ${index}.` } }) + } + + const again = session.send('Hi', { inputId: 't-0' }) + + expect(again.id).toBe(first.id) + expect(await again.receipt).toEqual(await first.receipt) + await host.close() + }) + + it('forgets agent runs past settledAgentRuns', async () => { + const { host, session } = await openWith({ settledAgentRuns: 3 }) + const ids: Array = [] + for (let index = 0; index < 10; index += 1) { + const run = session.agents.writer.run() + ids.push(run.id) + await run + } + + const kept = ids.slice(7) + expect(ids.filter((id) => session.operation(id) !== undefined)).toEqual( + kept, + ) + expect(session.agentRuns().map((run) => run.operationId)).toEqual(kept) + expect(session.agentRun(ids[0] ?? '')).toBe(undefined) + await host.close() + }) + + it('keeps at most eventBytes of events for events()', async () => { + // Every event is longer than 1 byte, so only the newest one stays. + const { host, session } = await openWith({ eventBytes: 1 }) + await session.send('Hi') + await session.close() + + expect(await collect(session.events({ from: '0' }))).toHaveLength(1) + await host.close() + }) + + it('keeps the ids of forgotten operations only as long as both windows', async () => { + // Without `stores.runs`, only the session refuses a reused run id. The + // interrupts store needs the runs store, so it goes too. + const persistence = { + stores: { + ...memoryPersistence().stores, + runs: undefined, + interrupts: undefined, + }, + } + const { host, session } = await openWith( + { settledOperations: 2, settledAgentRuns: 1 }, + persistence, + ) + await session.send('Hi', { runId: 'run-0' }) + for (let index = 0; index < 2; index += 1) await session.send('Hi') + expect(session.operation('run-0')).toBe(undefined) + + const reused = session.send('Hi', { runId: 'run-0' }) + expect(await reused.receipt).toMatchObject({ + status: 'rejected', + reason: 'conflict', + }) + await reused.then( + () => undefined, + () => undefined, + ) + // Both windows keep 3 operations. Then the id is free again. + for (let index = 0; index < 4; index += 1) await session.send('Hi') + + expect(privateSetSize(session, 'retiredIds')).toBe(3) + await expect(session.send('Hi', { runId: 'run-0' })).resolves.toEqual({ + text: 'done', + media: [], + }) + await host.close() + }) + + it('keeps no follower of a duplicate input on a host without a log', async () => { + const { host, session } = await openWith({ settledOperations: 1 }) + await session.send('Hi', { inputId: 't-0' }) + // The next turn pushes the first one out of the window. + await session.send('Hi', { inputId: 't-1' }) + + await expect(session.send('Hi', { inputId: 't-0' })).rejects.toThrow( + 'its result is gone', + ) + expect(privateSetSize(session, 'followers')).toBe(0) + await host.close() + }) +}) From 277eeb9cb9cd07e172e313f4c53d7e4a8b460b8c Mon Sep 17 00:00:00 2001 From: Alem Tuzlak Date: Sat, 17 Oct 2026 13:37:02 +0200 Subject: [PATCH 04/12] feat(ai-harness, ai-mcp, ai-acp, ai-harness-cli): reference-counted session handles with idle close and maxOpen --- packages/ai-acp/src/agent/index.ts | 216 +++--- packages/ai-acp/tests/agent-handlers.test.ts | 37 +- packages/ai-harness-cli/src/index.ts | 18 +- packages/ai-harness-cli/tests/signals.test.ts | 20 +- packages/ai-harness/src/harness-text.ts | 15 +- packages/ai-harness/src/host.ts | 615 +++++++++++++----- packages/ai-harness/src/http.ts | 274 +++++--- packages/ai-harness/src/session/base.ts | 16 +- packages/ai-harness/src/session/handle.ts | 432 ++++++++++-- .../ai-harness/src/session/harness-session.ts | 48 +- packages/ai-harness/src/session/requests.ts | 16 +- packages/ai-harness/src/session/types.ts | 5 + packages/ai-harness/src/worker.ts | 2 + packages/ai-harness/tests/agent-runs.test.ts | 3 +- packages/ai-harness/tests/auth.test.ts | 3 +- packages/ai-harness/tests/fork.test.ts | 15 +- packages/ai-harness/tests/forwarding.test.ts | 12 +- packages/ai-harness/tests/handles.test.ts | 275 ++++++++ .../ai-harness/tests/harness-text.test.ts | 16 +- packages/ai-harness/tests/oauth-edges.test.ts | 3 +- packages/ai-harness/tests/reload.test.ts | 2 +- packages/ai-harness/tests/retention.test.ts | 7 +- .../ai-harness/tests/session-plugins.test.ts | 6 +- packages/ai-harness/tests/session.test.ts | 27 +- .../ai-harness/tests/sessions-index.test.ts | 18 +- packages/ai-harness/tests/shared-log.test.ts | 32 +- packages/ai-mcp/src/harness.ts | 25 +- packages/ai-mcp/tests/harness.test.ts | 27 + 28 files changed, 1673 insertions(+), 512 deletions(-) create mode 100644 packages/ai-harness/tests/handles.test.ts diff --git a/packages/ai-acp/src/agent/index.ts b/packages/ai-acp/src/agent/index.ts index 41a534de02..982800fca8 100644 --- a/packages/ai-acp/src/agent/index.ts +++ b/packages/ai-acp/src/agent/index.ts @@ -249,6 +249,10 @@ const APPROVAL_OPTIONS = [ * 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. + * * ACP v2 is a draft, so this API is experimental too. */ export function createAcpAgent

( @@ -260,15 +264,38 @@ export function createAcpAgent

( 'createAcpAgent needs a principal: the host has access. Pass the local user as principal.', ) } - /** The ids of the sessions that this agent opened. */ - const opened = new Set() + /** + * 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. + */ + const handles = new Map>() /** The last pump of each session. The pumps of one session run in order. */ const tails = new Map>() - // The host gives back a handle to the live session of a thread, or opens - // it again after it closed, so the agent keeps no session objects. - const live = (sessionId: string) => - host.open(harness, { threadId: sessionId, principal }) + /** The handle of a session: the one that the agent holds, or a new one. */ + const live = (sessionId: string) => { + const held = handles.get(sessionId) + if (held) return held + 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) + }) + return opening + } + + /** Release the handle of a session. The session closes on idle. */ + const release = async (sessionId: string) => { + const held = handles.get(sessionId) + if (!held) return + handles.delete(sessionId) + await held.then( + (session) => session.close(), + () => {}, + ) + } /** * The session of an id that this agent or the session index knows. A @@ -278,7 +305,7 @@ export function createAcpAgent

( */ const known = async (sessionId: string) => { try { - if (!opened.has(sessionId)) { + if (!handles.has(sessionId)) { const entry = await host.sessions.get({ threadId: sessionId, principal, @@ -287,9 +314,7 @@ export function createAcpAgent

( entry?.harness === harness.name && entry.parentThreadId === undefined if (!isOwnSession) throw RequestError.resourceNotFound(sessionId) } - const session = await live(sessionId) - opened.add(sessionId) - return session + return await live(sessionId) } catch (error) { if (error instanceof AccessDeniedError) { throw RequestError.resourceNotFound(sessionId) @@ -463,94 +488,97 @@ export function createAcpAgent

( ) } - return agent() - .onRequest('initialize', () => ({ - protocolVersion: PROTOCOL_VERSION, - info: { name: harness.name, version: options.version ?? '0.0.0' }, - // Text and resource links are the baseline. Images, audio, and embedded - // resources go to the media store of the session. - capabilities: { - session: { prompt: { image: {}, audio: {}, embeddedContext: {} } }, - }, - })) - .onRequest('session/new', async () => { - const sessionId = `acp-${Date.now().toString(36)}-${Math.random().toString(36).slice(2, 9)}` - opened.add(sessionId) - await live(sessionId) - return { sessionId } - }) - .onRequest('session/list', async ({ params }) => { - // ponytail: the harness keeps no cwd per thread, so every session - // reports the cwd of the agent. Store one per thread when one agent - // serves more than one project. - const cwd = options.cwd ?? process.cwd() - if (params.cwd && params.cwd !== cwd) return { sessions: [] } - const page = await host.sessions.list({ - principal, - harness: harness.name, - limit: LIST_PAGE_SIZE, - ...(params.cursor ? { cursor: params.cursor } : {}), + return ( + agent() + .onConnect((connection) => { + // The handles of the sessions go with the connection. + void connection.closed.then(() => + Promise.all([...handles.keys()].map(release)), + ) }) - return { - sessions: page.entries.map((entry) => ({ - sessionId: entry.threadId, - cwd, - title: entry.title ?? null, - updatedAt: new Date(entry.updatedAt).toISOString(), - })), - nextCursor: page.truncated ? page.cursor : null, - } - }) - .onRequest('session/resume', async ({ params }) => { - await known(params.sessionId) - return {} - }) - .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 + .onRequest('initialize', () => ({ + protocolVersion: PROTOCOL_VERSION, + info: { name: harness.name, version: options.version ?? '0.0.0' }, + // Text and resource links are the baseline. Images, audio, and embedded + // resources go to the media store of the session. + capabilities: { + session: { prompt: { image: {}, audio: {}, embeddedContext: {} } }, }, - ) - const messageId = `user-${crypto.randomUUID()}` - await notify(client, sessionId, { - sessionUpdate: 'user_message', - messageId, - content: params.prompt, + })) + .onRequest('session/new', async () => { + const sessionId = `acp-${Date.now().toString(36)}-${Math.random().toString(36).slice(2, 9)}` + await live(sessionId) + return { sessionId } }) - // `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), - ) - tails.set(sessionId, tail) - void tail.then(() => { - if (tails.get(sessionId) === tail) tails.delete(sessionId) + .onRequest('session/list', async ({ params }) => { + // ponytail: the harness keeps no cwd per thread, so every session + // reports the cwd of the agent. Store one per thread when one agent + // serves more than one project. + const cwd = options.cwd ?? process.cwd() + if (params.cwd && params.cwd !== cwd) return { sessions: [] } + const page = await host.sessions.list({ + principal, + harness: harness.name, + limit: LIST_PAGE_SIZE, + ...(params.cursor ? { cursor: params.cursor } : {}), + }) + return { + sessions: page.entries.map((entry) => ({ + sessionId: entry.threadId, + cwd, + title: entry.title ?? null, + updatedAt: new Date(entry.updatedAt).toISOString(), + })), + nextCursor: page.truncated ? page.cursor : null, + } }) - return { messageId } - }) - .onRequest('session/close', async ({ params }) => { - if (!opened.delete(params.sessionId)) return - // ponytail: a handle has no close of its own yet, so this closes the - // live session of the thread on this host, for every handle. - await (await live(params.sessionId)).close() - }) - .onNotification('session/cancel', async ({ params }) => { - // While the client asks for an approval, no turn runs: the client - // answers that request with `cancelled`, and the turn stops there. - if (!opened.has(params.sessionId)) return - await (await live(params.sessionId)).cancel() - }) + .onRequest('session/resume', async ({ params }) => { + await known(params.sessionId) + return {} + }) + .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), + ) + tails.set(sessionId, tail) + void tail.then(() => { + if (tails.get(sessionId) === tail) tails.delete(sessionId) + }) + return { messageId } + }) + // Only the handle of this agent goes. Other handles keep the session. + .onRequest('session/close', ({ params }) => release(params.sessionId)) + .onNotification('session/cancel', async ({ params }) => { + // While the client asks for an approval, no turn runs: the client + // answers that request with `cancelled`, and the turn stops there. + if (!handles.has(params.sessionId)) return + await (await live(params.sessionId)).cancel() + }) + ) } /** diff --git a/packages/ai-acp/tests/agent-handlers.test.ts b/packages/ai-acp/tests/agent-handlers.test.ts index 2f34acc0fc..6755b79fb1 100644 --- a/packages/ai-acp/tests/agent-handlers.test.ts +++ b/packages/ai-acp/tests/agent-handlers.test.ts @@ -9,7 +9,7 @@ import { } from '@tanstack/ai-harness' import { createAcpAgent, serveAcp, toSessionUpdate } from '../src/agent' import type { AnyTextAdapter, ModelMessage, StreamChunk } from '@tanstack/ai' -import type { HarnessAccess } from '@tanstack/ai-harness' +import type { HarnessAccess, HostLifecycleOptions } from '@tanstack/ai-harness' import type { ClientApp, RequestPermissionOutcome, @@ -136,7 +136,10 @@ function model(): AnyTextAdapter { } /** With `access`, the host checks each action of each principal. */ -function setup({ access }: { access?: HarnessAccess } = {}) { +function setup({ + access, + lifecycle, +}: { access?: HarnessAccess; lifecycle?: HostLifecycleOptions } = {}) { const remove = vi.fn(async () => ({ removed: true })) const harness = defineHarness({ name: 'test/acp-handlers', @@ -161,7 +164,7 @@ function setup({ access }: { access?: HarnessAccess } = {}) { ], }) const persistence = memoryPersistence() - const host = createHarnessHost({ persistence, access }) + const host = createHarnessHost({ persistence, access, lifecycle }) const updates: Array<{ sessionId: string; update: SessionUpdate }> = [] const permissions: Array<{ title: string; toolCallId: string }> = [] // The answer of the editor to each permission request. Default: Reject. @@ -335,6 +338,34 @@ describe('ACP agent handlers', () => { close = undefined }) + it('releases its handle at session/close, and every handle when the connection closes', async () => { + // A session that no handle holds closes at once when it is idle. + const { harness, host, editor } = 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', { + cwd: '/tmp', + mcpServers: [], + }) + kept = (await ctx.request('session/new', { cwd: '/tmp', mcpServers: [] })) + .sessionId + expect(isOpen(closed.sessionId)).toBe(true) + + await ctx.request('session/close', { sessionId: closed.sessionId }) + + expect(isOpen(closed.sessionId)).toBe(false) + expect(isOpen(kept)).toBe(true) + }) + // 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 () => { const { harness, host, editor, permissions, remove, idle, textOf } = setup() close = () => host.close() diff --git a/packages/ai-harness-cli/src/index.ts b/packages/ai-harness-cli/src/index.ts index ba91ba39d0..a3f9401d7d 100644 --- a/packages/ai-harness-cli/src/index.ts +++ b/packages/ai-harness-cli/src/index.ts @@ -80,7 +80,8 @@ export interface RunCliOptions

{ * signal while no turn runs, closes what the CLI opened, and `runCli` * resolves to 130 (SIGINT) or 143 (SIGTERM). A host that the CLI made * closes with `recoverable: true`, so a durable host runs the cut work - * again. Your own `host` stays open: the CLI closes only its session. + * again. Your own `host` stays open: the CLI releases only its handle of + * the session. */ signals?: SignalSource } @@ -194,6 +195,8 @@ export async function runCli

( // have no turn of their own, so a signal stops them at once. let cancel = () => false const stop = stopOnSignals(options.signals ?? process, () => cancel()) + /** The handle of the session modes. The CLI releases it at the end. */ + let session: HarnessSession | undefined try { if (args.acp) { const acp = await loadOptional( @@ -301,7 +304,7 @@ export async function runCli

( return code } - const session = await host.open(harness, { + session = await host.open(harness, { threadId: args.thread, ...(principal ? { principal } : {}), }) @@ -329,9 +332,13 @@ export async function runCli

( // A signal that cancelled the turn of --print gives its exit code. return code === EXIT.cancelled ? (stop.exitCode() ?? code) : code } - // The mode ends when its session closes. Work that the close cuts can - // fail, and the signal gives the exit code. - await session.close({ recoverable: true }) + // A host that the CLI made closes now, so the mode ends with its + // session, and its running work stops as on a crash: a durable host runs + // it again. On your own host, the CLI releases its handle, and the mode + // ends with the cancelled turn. Work that the close cuts can fail, and + // the signal gives the exit code. + if (options.host) await session.close() + else await host.close({ recoverable: true }) await running.catch(() => {}) return await stop.stopped } catch (error) { @@ -339,6 +346,7 @@ export async function runCli

( return EXIT.failed } finally { stop.dispose() + await session?.close() // A host the caller gave stays open: the caller owns it. After a // signal, its running work stops as on a crash, so a durable host runs // it again. diff --git a/packages/ai-harness-cli/tests/signals.test.ts b/packages/ai-harness-cli/tests/signals.test.ts index 53dd19c566..b6c79d7621 100644 --- a/packages/ai-harness-cli/tests/signals.test.ts +++ b/packages/ai-harness-cli/tests/signals.test.ts @@ -139,10 +139,14 @@ describe('signals in print mode', () => { expect(calls[0]?.aborted).toBe(true) }) - it('closes only its session of your host on a second signal while the cancelled turn runs on', async () => { + it('releases only its handle of your host on a second signal while the cancelled turn runs on', async () => { const { model, calls, release } = waitingModel([], { stuck: true }) const harness = defineHarness({ name: 'test/signal-stuck', adapter: model }) - const host = createHarnessHost({ persistence: memoryPersistence() }) + // A session that no handle holds closes at once when it is idle. + const host = createHarnessHost({ + persistence: memoryPersistence(), + lifecycle: { idleCloseMs: 0 }, + }) const close = vi.spyOn(host, 'close') const cli = startCli(harness, { argv: ['-p', 'hi'], host }) await vi.waitFor(() => expect(calls).toHaveLength(1)) @@ -154,18 +158,24 @@ describe('signals in print mode', () => { expect(await cli.running).toBe(EXIT.cancelled) expect(close).not.toHaveBeenCalled() - expect(host.isOpen(harness, { threadId: 'main' })).toBe(false) + await vi.waitFor(() => + expect(host.isOpen(harness, { threadId: 'main' })).toBe(false), + ) await host.close() }) }) describe('signals in a custom ui', () => { - it('stops a ui that is still open, and closes its session of your host', async () => { + it('stops a ui that is still open, and releases its handle of your host', async () => { const harness = defineHarness({ name: 'test/signal-ui', adapter: fakeText(), }) - const host = createHarnessHost({ persistence: memoryPersistence() }) + // A session that no handle holds closes at once when it is idle. + const host = createHarnessHost({ + persistence: memoryPersistence(), + lifecycle: { idleCloseMs: 0 }, + }) let shown = false const cli = startCli(harness, { argv: [], diff --git a/packages/ai-harness/src/harness-text.ts b/packages/ai-harness/src/harness-text.ts index 7b89194a9e..bea50b0847 100644 --- a/packages/ai-harness/src/harness-text.ts +++ b/packages/ai-harness/src/harness-text.ts @@ -290,10 +290,9 @@ export function harnessText( // No thread id: a new inner session per call, as codexText does. const threadId = `${chatOptions.threadId ?? runId}:${harness.name}` const signal = signalOf(chatOptions.request) - // A session this call opens closes with it. One the caller holds open - // stays, as in host.fork. + // The call releases its handle at the end. The session then closes + // on idle, unless the caller holds it too. const user = principal ? { principal } : {} - const wasOpen = host.isOpen(harness, { threadId, ...user }) const session = await host.open(harness, { threadId, ...user }) // The outer call's tools join the inner turn. Its system prompts do // not: the harness owns its prompt stack. @@ -351,15 +350,7 @@ export function harnessText( } finally { // A long-lived signal must not keep one listener per call. signal?.removeEventListener('abort', cancel) - if (!wasOpen) - await session - .close() - .catch((error: unknown) => - console.error( - '[@tanstack/ai-harness] The session did not close.', - error, - ), - ) + await session.close() } })(), structuredOutput: () => diff --git a/packages/ai-harness/src/host.ts b/packages/ai-harness/src/host.ts index f1cbc7a1cb..109df9152f 100644 --- a/packages/ai-harness/src/host.ts +++ b/packages/ai-harness/src/host.ts @@ -151,6 +151,18 @@ export interface HarnessHostOptions

{ /** How the sessions of a host live. */ export interface HostLifecycleOptions { + /** + * How long a session that no handle holds stays open after it goes idle, + * in milliseconds. Then the host closes it, and the next call on one of + * its handles opens it again. `0` closes it at once. Default 60 000. + */ + idleCloseMs?: number + /** + * How many sessions the host keeps open. An open past it first closes the + * least recently used sessions that no handle holds and that are idle. + * When there are none, the open goes on, with one warning. Default 500. + */ + maxOpen?: number /** How much each session keeps in memory. */ retention?: RetentionLimits } @@ -452,8 +464,9 @@ export interface ResumePendingOptions { harnesses: ReadonlyArray /** * What happens to a session that the sweep opened. `'whenIdle'` - * (default): it closes when its pending work ends, unless the app opened - * the thread too. `'never'`: it stays open, like any opened session. + * (default): it closes as soon as its pending work ends, unless the app + * opened the thread too. `'never'`: it closes like any session that no + * handle holds, after `lifecycle.idleCloseMs`. */ close?: 'whenIdle' | 'never' /** @@ -477,6 +490,11 @@ export interface HarnessHost

{ * `parentThreadId`) does not open: open its root thread, and read the * child with `session.agentTranscript`. * + * Release the handle with `session.close()`, or with `await using`. The + * session closes when no handle holds it and it was idle for + * `lifecycle.idleCloseMs`. A call on a handle after that opens the session + * again. Throws after `host.close()`. + * * With access, `principal` is required. A thread that exists needs the * `read` action, else the open throws `AccessDeniedError`. A new thread * gets the principal as its owner. Then each action of the handle asks @@ -493,7 +511,8 @@ export interface HarnessHost

{ ) => Promise> /** * Is a session of `harness` open on `threadId` on this host, in the tenant - * of `principal`? + * of `principal`? A session that closed on idle is not open, also when a + * handle of it still exists. */ isOpen: ( harness: AnyHarness, @@ -557,9 +576,12 @@ export interface HarnessHost

{ */ reload: (harness: AnyHarness) => Promise /** - * Close every live session. With `recoverable`, their running work stops - * with no settlement, and the next host that opens a thread runs it again. - * See `session.close`. + * Close every live session, for every handle, and refuse later opens. + * After it, a call on a handle throws, except a read that gives a value at + * once (`snapshot()`): it reads the last state. With `recoverable`, the + * running work stops with no settlement, and the next host that opens a + * thread runs it again: use it when a durable host shuts down, for + * example for a deploy. Without it, the running work is cancelled. * * @example * ```ts @@ -802,6 +824,43 @@ export function sessionOfTurn

( return turnFinders.get(host)?.(lookup) ?? Promise.resolve(undefined) } +/** What a host keeps for each session key: its handles and its idle timer. */ +interface Hold { + /** The handles that hold the session now. */ + handles: number + /** The idle close that waits. */ + timer: ReturnType | undefined + /** A dropped handle of this session warned already. */ + warned: boolean +} + +/** What a host counts, for a leak check. */ +interface HostCounts { + /** Open sessions, and sessions that open. */ + sessions: number + /** Handles that hold a session. */ + handles: number + /** Session keys with a hold record. */ + holds: number + /** Statuses that `host.events()` starts with. */ + statuses: number + /** Sessions that close now. */ + closing: number +} + +/** How each host counts. Not public API. */ +const counters = new WeakMap HostCounts>() + +/** + * @internal What `host` keeps now: sessions, handles, hold records, + * statuses, and closes. Tests read it to find leaks. + */ +export function hostCountsOf

(host: HarnessHost

) { + const count = counters.get(host) + if (!count) throw new Error('Not a host of createHarnessHost.') + return count() +} + /** The log and run stores of each durable host. Not public API. */ const durableStores = new WeakMap() @@ -1034,15 +1093,192 @@ export function createHarnessHost

( return loaded } const sessions = new Map>() + /** The sessions of `sessions` that finished opening, by key. */ + const live = new Map() /** * The keys of sessions that only a `resumePending` sweep opened. With - * `close: 'whenIdle'`, the sweep closes only these. An open by the app - * takes a key out. + * `close: 'whenIdle'`, such a session closes at once when it goes idle. An + * open by the app takes a key out. */ const sweptOnly = new Set() - /** The closes of swept sessions that left the cache, by key. */ + /** The closes of sessions that left the cache, by key. */ const closing = new Map>() const hostId = `host-${Date.now().toString(36)}-${Math.random().toString(36).slice(2, 9)}` + const idleCloseMs = options.lifecycle?.idleCloseMs ?? 60_000 + const maxOpen = options.lifecycle?.maxOpen ?? 500 + /** `host.close()` ran: no session opens any more. */ + let isClosed = false + /** + * The handles and the idle timer of each session key, the least recently + * used key first. Each open session has one. + */ + const holds = new Map() + /** The open went past `maxOpen` with no session to close, and warned. */ + let warnedFull = false + + const hostClosed = () => + new Error('The host is closed. Create a new host to open sessions.') + + /** The hold of `key`, now the most recently used. */ + const touch = (key: string) => { + const held = holds.get(key) ?? { + handles: 0, + timer: undefined, + warned: false, + } + holds.delete(key) + holds.set(key, held) + return held + } + + /** Stop the idle close of `key` that waits. */ + const cancelIdle = (key: string) => { + const held = holds.get(key) + clearTimeout(held?.timer) + if (held) held.timer = undefined + } + + /** Drop the hold record of `key` when no handle and no session uses it. */ + const dropHold = (key: string) => { + const held = holds.get(key) + if (!held || held.handles > 0 || sessions.has(key)) return + clearTimeout(held.timer) + holds.delete(key) + } + + /** Keep the close of `key`, so the next open of the key waits for it. */ + const trackClose = (key: string, closed: Promise) => { + const settled = closed.then( + () => {}, + () => {}, + ) + closing.set(key, settled) + void settled.then(() => { + if (closing.get(key) === settled) closing.delete(key) + }) + } + + /** Take `session` out of the open sessions of `key`, when it is the one. */ + const forget = (key: string, session: SessionImpl) => { + if (live.get(key) !== session) return + live.delete(key) + sessions.delete(key) + statuses.delete(key) + sweptOnly.delete(key) + cancelIdle(key) + dropHold(key) + } + + /** + * Close the session of `key` for every handle. From here on, an open of + * the key gets a new session, which waits for this close. A handle opens + * the new one on its next call. + */ + const closeShared = ( + key: string, + session: SessionImpl, + closeOptions?: { recoverable?: boolean }, + ) => { + forget(key, session) + const closed = session.close(closeOptions) + trackClose(key, closed) + return closed + } + + /** `closeShared` of an idle session: a failed close only logs. */ + const closeIdle = (key: string, session: SessionImpl) => + // ponytail: an idle session has no work to lose. + closeShared(key, session).catch((error: unknown) => + console.error( + '[@tanstack/ai-harness] An idle session did not close.', + error, + ), + ) + + /** No handle holds the open session `session` of `key`, and it is idle. */ + const isUnused = (key: string, session: SessionImpl) => + live.get(key) === session && + (holds.get(key)?.handles ?? 0) === 0 && + !session.isClosing() && + session.isIdle() + + /** + * Start the idle close of the session of `key` when no handle holds it and + * it is idle. With a delay of 0, it closes at once. Resolves after a close + * that it starts. + */ + const armIdle = (key: string) => { + const session = live.get(key) + const held = holds.get(key) + if (!session || !held) { + dropHold(key) + return Promise.resolve() + } + if (held.timer !== undefined || !isUnused(key, session)) { + return Promise.resolve() + } + const delay = sweptOnly.has(key) ? 0 : idleCloseMs + if (delay <= 0) return closeIdle(key, session) + const timer = setTimeout(() => { + held.timer = undefined + // A session that is busy now arms again when it goes idle. + if (isUnused(key, session)) void closeIdle(key, session) + }, delay) + // The timer must not keep a process alive. + if (typeof timer === 'object' && 'unref' in timer) timer.unref() + held.timer = timer + return Promise.resolve() + } + + /** One more handle on the session of `key`: it does not close on idle. */ + const hold = (key: string) => { + touch(key).handles += 1 + cancelIdle(key) + } + + /** One handle less on the session of `key`. Resolves after its close. */ + const release = (key: string) => { + const held = holds.get(key) + if (!held) return Promise.resolve() + held.handles = Math.max(0, held.handles - 1) + touch(key) + return armIdle(key) + } + + /** The first dropped handle of a session warns. */ + const warnDropped = (key: string, threadId: string) => { + const held = holds.get(key) + if (!held || held.warned) return + held.warned = true + console.warn( + `[@tanstack/ai-harness] A handle of thread ${JSON.stringify(threadId)} was dropped without close(), so the host released it. Call session.close(), or open the session with await using.`, + ) + } + + /** + * Room for one more session: close the least recently used sessions that + * no handle holds and that are idle. When none is, warn once. + */ + const makeRoom = () => { + let excess = sessions.size + 1 - maxOpen + if (excess <= 0) { + warnedFull = false + return + } + const keys = [...holds.keys()] + for (const key of keys) { + if (excess <= 0) break + const session = live.get(key) + if (!session || !isUnused(key, session)) continue + void closeIdle(key, session) + excess -= 1 + } + if (excess <= 0 || warnedFull) return + warnedFull = true + console.warn( + `[@tanstack/ai-harness] The host has ${sessions.size} open sessions, more than lifecycle.maxOpen (${maxOpen}), and each one is busy or held by a handle. Release handles with session.close().`, + ) + } /** Throw when the index entry of the thread names a parent thread. */ const refuseChild = async (key: SessionIndexKey) => { @@ -1054,22 +1290,33 @@ export function createHarnessHost

( } /** - * `host.open`, in the tenant `tenantId`. It gives the session class, so - * the host can call the members that the `HarnessSession` type does not - * have. An internal open can pass a tenant with no principal, so it does - * not claim an entry with no owner. + * `host.open`, in the tenant `tenantId`, with no checks and no handle. It + * gives the session class, so the host can call the members that the + * `HarnessSession` type does not have. An internal open can pass a tenant + * with no principal, so it does not claim an entry with no owner. Hold the + * key first (`hold`), or the session can close on idle at once. */ const openSession = ( harness: THarness, { threadId, principal, logId: openLogId, promptCache }: OpenSessionOptions, tenantId: string | undefined, ) => { + if (isClosed) return Promise.reject(hostClosed()) const key = sessionKey(harness.name, tenantId, threadId) const of = { tenantId, threadId } + touch(key) let session = sessions.get(key) + // A session that closes by itself (its log failed) takes no new work. + const stopping = live.get(key) + if (session && stopping?.isClosing()) { + forget(key, stopping) + trackClose(key, stopping.close()) + session = undefined + } // The app opened a thread that a sweep opened: it stays open. if (session) sweptOnly.delete(key) if (!session) { + makeRoom() const logId = openLogId ?? threadId // The one place where the log id of the app becomes the id in the // store: the tenant is part of it. @@ -1112,37 +1359,50 @@ export function createHarnessHost

( ...(promptCache ? { promptCache } : {}), ...(access ? { access: { envKeys: access.envKeys === true } } : {}), onEvent: (event) => { + // A session that left the cache has no status any more. + if (sessions.get(key) !== session) return const status = statusAfter(created, event) + if (status === 'running') cancelIdle(key) if (status) setStatus(key, of, status) }, - onClose: () => { - // A sweep may have put a new session under this key already. It - // opens only after this one closed, so its status is not set yet. - if (sessions.get(key) === session) sessions.delete(key) - statuses.delete(key) - }, + // A new session under this key opens only after this one closed. + onClose: () => forget(key, created), + // After the call: the session may close, but not inside its own code. + onIdle: () => + queueMicrotask(() => { + if (live.get(key) === created) void armIdle(key) + }), }) - // A swept session under this key may still close: it lets go of - // the thread first. + // A session under this key may still close: it lets go of the thread + // first. const closed = closing.get(key) ?? Promise.resolve() session = closed .then(() => refuseChild(of)) .then(() => created.open()) .then( () => { + if (sessions.get(key) !== session) return created + live.set(key, created) // A recovered turn can run already. Else the new session is // idle. if (!statuses.has(key)) setStatus(key, of, 'idle') + // A session that no handle holds (a sweep) closes on idle too. + void armIdle(key) return created }, (error: unknown) => { - sessions.delete(key) - statuses.delete(key) + if (sessions.get(key) === session) { + sessions.delete(key) + statuses.delete(key) + } + dropHold(key) throw error }, ) sessions.set(key, session) } + // The key names the harness, so the session is a session of `harness`. + // A map of sessions cannot carry a harness type per key. return session as Promise> } @@ -1300,18 +1560,84 @@ export function createHarnessHost

( } } - /** The handle of `session` for `principal`. With access, it checks. */ - const handleOf = ( - session: SessionImpl, - principal: P | undefined, - ) => - sessionHandle({ + /** + * `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. + */ + const openChecked = async ( + harness: THarness, + openOptions: OpenSessionOptions

, + ) => { + if (!access) { + return openSession(harness, openOptions, openOptions.principal?.tenantId) + } + const principal = need(openOptions.principal, 'host.open') + const { threadId } = openOptions + const ids = { tenantId: principal.tenantId, threadId } + const before = await index.get(ids) + // The open refuses a child thread, and says why. + const isRoot = before !== undefined && before.parentThreadId === undefined + if (isRoot) { + await allow({ principal, action: 'read', threadId, entry: before }) + } + const session = await openSession(harness, openOptions, principal.tenantId) + if (!before) { + // Two users can create one thread at once. Only the user whose open + // wrote the entry needs no check. + const created = await index.get(ids) + if (created?.ownerId !== principal.id) { + await allow({ principal, action: 'read', threadId, entry: created }) + } + } + return session + } + + /** + * The handle of the session of `key` for the principal of `openOptions`. + * It holds the key already (`hold`), and releases it on close. With access, + * it checks. A call on it after the session closed opens it again with + * `openOptions`. + */ + const handleOf = ({ + harness, + session, + openOptions, + key, + }: { + harness: THarness + session: SessionImpl + openOptions: OpenSessionOptions

+ key: string + }) => { + const { principal } = openOptions + const { threadId } = session + return sessionHandle({ session, principal, - ...(access && principal - ? { guard: guardOf(principal, session.threadId) } - : {}), + ...(access && principal ? { guard: guardOf(principal, threadId) } : {}), + lease: { + current: () => { + // After `host.close()`, a call refuses at once. + const open = isClosed ? undefined : live.get(key) + // The key names the harness, so the session is a session of + // `harness`, as in `openSession`. + return open && !open.isClosing() + ? (open as SessionImpl) + : undefined + }, + reopen: () => { + if (isClosed) throw hostClosed() + return openChecked(harness, openOptions) + }, + hold: () => hold(key), + release: (releaseOptions) => { + if (releaseOptions?.dropped) warnDropped(key, threadId) + return release(key) + }, + }, }) + } /** The open sessions of a thread of a tenant on this host, of any harness. */ const openSessionsOf = async ( @@ -1319,17 +1645,20 @@ export function createHarnessHost

( threadId: string, ) => { const suffix = `\u0000${logIdOf({ tenantId, logId: threadId })}` - const live = await Promise.allSettled( + const opened = await Promise.allSettled( [...sessions].flatMap(([key, session]) => key.endsWith(suffix) ? [session] : [], ), ) - return live.flatMap((entry) => + return opened.flatMap((entry) => entry.status === 'fulfilled' ? [entry.value] : [], ) } - /** `host.fork` after its checks: the session class of the new thread. */ + /** + * `host.fork` after its checks: the session class of the new thread, and + * its key, which the fork holds. Give the hold to a handle, or release it. + */ const forkInto = async ( harness: THarness, { threadId, newThreadId, at, principal, logId }: ForkSessionOptions

, @@ -1338,16 +1667,17 @@ export function createHarnessHost

( throw new Error('A fork needs a new thread: newThreadId is threadId.') } const tenantId = principal?.tenantId - // A source that only this fork opens closes again at the end. - const wasOpen = sessions.has(sessionKey(harness.name, tenantId, threadId)) - // In the tenant of the fork, with no principal: the fork does not - // become the owner of a source with no owner. - const source = await openSession( - harness, - { threadId, ...(logId ? { logId } : {}) }, - tenantId, - ) + // The fork holds the source while it reads it. + const sourceKey = sessionKey(harness.name, tenantId, threadId) + hold(sourceKey) try { + // In the tenant of the fork, with no principal: the fork does not + // become the owner of a source with no owner. + const source = await openSession( + harness, + { threadId, ...(logId ? { logId } : {}) }, + tenantId, + ) const transcript = await source.transcript() const end = at === undefined @@ -1378,61 +1708,40 @@ export function createHarnessHost

( `Thread ${JSON.stringify(newThreadId)} is not empty. Fork into a new thread.`, ) } - const fork = await openSession( - harness, - { threadId: newThreadId, ...(principal ? { principal } : {}) }, - tenantId, - ) - await fork.adoptFork(source, transcript.slice(0, end)) - return fork + const key = sessionKey(harness.name, tenantId, newThreadId) + hold(key) + try { + const fork = await openSession( + harness, + { threadId: newThreadId, ...(principal ? { principal } : {}) }, + tenantId, + ) + await fork.adoptFork(source, transcript.slice(0, end)) + return { fork, key } + } catch (error) { + await release(key) + throw error + } } finally { - // A failed close must not hide the fork's own error. - if (!wasOpen) - await source - .close() - .catch((error: unknown) => - console.error( - '[@tanstack/ai-harness] The source session did not close.', - error, - ), - ) + // A source that only this fork opened closes on idle. + await release(sourceKey) } } const host: HarnessHost

= { async open(harness, openOptions) { - if (!access) { - const { principal } = openOptions - const session = await openSession( - harness, - openOptions, - principal?.tenantId, - ) - return handleOf(session, principal) - } - const principal = need(openOptions.principal, 'host.open') - const { threadId } = openOptions - const key = { tenantId: principal.tenantId, threadId } - const before = await index.get(key) - // The open refuses a child thread, and says why. - const isRoot = before !== undefined && before.parentThreadId === undefined - if (isRoot) { - await allow({ principal, action: 'read', threadId, entry: before }) - } - const session = await openSession( - harness, - openOptions, - principal.tenantId, - ) - if (!before) { - // Two users can create one thread at once. Only the user whose open - // wrote the entry needs no check. - const created = await index.get(key) - if (created?.ownerId !== principal.id) { - await allow({ principal, action: 'read', threadId, entry: created }) - } + if (isClosed) throw hostClosed() + const tenantId = openOptions.principal?.tenantId + const key = sessionKey(harness.name, tenantId, openOptions.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 }) + } catch (error) { + await release(key) + throw error } - return handleOf(session, principal) }, isOpen(harness, { threadId, principal }) { if (access) need(principal, 'host.isOpen') @@ -1442,10 +1751,13 @@ export function createHarnessHost

( }, async fork(harness, forkOptions) { await checkFork('host.fork', forkOptions) - return handleOf( - await forkInto(harness, forkOptions), - forkOptions.principal, - ) + const { fork, key } = await forkInto(harness, forkOptions) + const { principal } = forkOptions + const openOptions = { + threadId: fork.threadId, + ...(principal ? { principal } : {}), + } + return handleOf({ harness, session: fork, openOptions, key }) }, revoke: ({ threadId, principal }) => revokeUser({ @@ -1497,13 +1809,13 @@ export function createHarnessHost

( // A claim is on the store key of the thread: its tenant and its id. const { tenantId, logId: threadId } = parseLogId(claimed) const key = sessionKey(name, tenantId, threadId) - // Only a session that this sweep opens closes when idle. A session - // that the app has open stays open. + // Only a session that this sweep opens closes at once when idle. A + // session that the app has open stays open. No handle holds the + // session, so it closes on idle in any case. const isOwn = close === 'whenIdle' && !sessions.has(key) if (isOwn) sweptOnly.add(key) - let session: SessionImpl try { - session = await openSession(harness, { threadId }, tenantId) + await openSession(harness, { threadId }, tenantId) } catch (error) { // One thread that cannot open must not stop the sweep. The claim // goes back, so the next sweep tries again. @@ -1514,23 +1826,6 @@ export function createHarnessHost

( failures.push(error) continue } - if (isOwn) { - const stop = session.onIdle(() => { - stop() - // The app opened it since: it stays open. - if (!sweptOnly.delete(key)) return - // From here on, an open gets a new session, so a new prompt - // never reaches this one while it closes. - sessions.delete(key) - // ponytail: a failed close is dropped, as an idle session has no - // work to lose. - const closed = session.close().catch(() => {}) - closing.set(key, closed) - void closed.then(() => { - if (closing.get(key) === closed) closing.delete(key) - }) - }) - } opened.push({ threadId, harness: name }) } if (failures.length > 0) { @@ -1546,9 +1841,9 @@ export function createHarnessHost

( return opened }, async recover(threadId) { - const live = await Promise.allSettled(sessions.values()) + const opened = await Promise.allSettled(sessions.values()) await Promise.all( - live.flatMap((entry) => + opened.flatMap((entry) => entry.status === 'fulfilled' && (threadId === undefined || entry.value.threadId === threadId) ? [entry.value.recover()] @@ -1558,13 +1853,13 @@ export function createHarnessHost

( }, async reload(harness) { const prefix = `${harness.name}\u0000` - const live = await Promise.allSettled( + const opened = await Promise.allSettled( [...sessions].flatMap(([key, session]) => key.startsWith(prefix) ? [session] : [], ), ) const results = await Promise.allSettled( - live.flatMap((entry) => + opened.flatMap((entry) => entry.status === 'fulfilled' ? [entry.value.reload()] : [], ), ) @@ -1573,10 +1868,13 @@ export function createHarnessHost

( } }, async close(closeOptions) { + // First: no open starts while the sessions close. + isClosed = true + for (const key of holds.keys()) cancelIdle(key) await Promise.allSettled(closing.values()) - const live = await Promise.allSettled(sessions.values()) + const opened = await Promise.allSettled(sessions.values()) await Promise.all( - live + opened .filter((entry) => entry.status === 'fulfilled') .map((entry) => entry.value.close(closeOptions)), ) @@ -1671,10 +1969,10 @@ export function createHarnessHost

( // thread id cannot read it. ponytail: the log or the messages only; // PR 4 adds a purge of blobs, metadata, and children. if (entry?.harness !== undefined) { - const open = sessions.get( - sessionKey(entry.harness, tenantId, threadId), - ) - if (open) await (await open).close() + const key = sessionKey(entry.harness, tenantId, threadId) + const opening = sessions.get(key) + // For every handle. A handle that calls again opens a new thread. + if (opening) await closeShared(key, await opening) } const storeId = logIdOf({ tenantId, logId: threadId }) const { log, messages } = persistence.stores @@ -1694,30 +1992,41 @@ export function createHarnessHost

( // `host.fork` takes the last message to copy, so `before` becomes // the message just before it, or `null` before the first one. let last: string | null = 'through' in at ? at.through : null - if ('before' in at) { - // With no principal, as `host.fork` opens the source. - const opened = await openSession(harness, { threadId }, tenantId) - const transcript = await opened.transcript() - const position = transcript.findIndex(({ id }) => id === at.before) - if (position === -1) { - throw new Error( - `The transcript has no message with id ${at.before}.`, - ) - } - if (position > 0) { - const previous = transcript[position - 1]?.id - if (!previous) { - throw new Error(`The message before ${at.before} has no id.`) + // The source stays open from the read of `before` to the fork. + const sourceKey = sessionKey(harness.name, tenantId, threadId) + hold(sourceKey) + let forked: SessionImpl + try { + if ('before' in at) { + // With no principal, as `host.fork` opens the source. + const opened = await openSession(harness, { threadId }, tenantId) + const transcript = await opened.transcript() + const position = transcript.findIndex(({ id }) => id === at.before) + if (position === -1) { + throw new Error( + `The transcript has no message with id ${at.before}.`, + ) + } + if (position > 0) { + const previous = transcript[position - 1]?.id + if (!previous) { + throw new Error(`The message before ${at.before} has no id.`) + } + last = previous } - last = previous } + const { fork, key } = await forkInto(harness, { + threadId, + newThreadId: `thread-${crypto.randomUUID()}`, + at: last, + ...(principal ? { principal } : {}), + }) + forked = fork + // No handle takes the fork: it closes on idle. + await release(key) + } finally { + await release(sourceKey) } - const forked = await forkInto(harness, { - threadId, - newThreadId: `thread-${crypto.randomUUID()}`, - at: last, - ...(principal ? { principal } : {}), - }) const mode = members ?? (access ? (policy?.fork.members ?? 'self') : undefined) const { ownerId, members: people } = forkedMembers( @@ -1834,6 +2143,16 @@ export function createHarnessHost

( durableStores.set(host, { log: logStore, ...(runs ? { runs } : {}) }) } if (access) accesses.set(host, access) + counters.set(host, () => ({ + sessions: sessions.size, + handles: [...holds.values()].reduce( + (total, held) => total + held.handles, + 0, + ), + holds: holds.size, + statuses: statuses.size, + closing: closing.size, + })) turnFinders.set(host, async ({ harness, tenantId, operationId }) => { // The key of a session with an empty thread id is the start of the keys // of the harness and tenant: the encoded tenant has no `:`. diff --git a/packages/ai-harness/src/http.ts b/packages/ai-harness/src/http.ts index 1d6870dfb9..2ebcb085ff 100644 --- a/packages/ai-harness/src/http.ts +++ b/packages/ai-harness/src/http.ts @@ -492,6 +492,37 @@ export function createHarnessHandler

( reader.abort(), ) + /** + * The answer of a route that opened `session`, and the release of that + * handle when the answer is done. A stream answer calls `keep` for its + * watch: the handle is released when the stream ends or the client leaves. + */ + const withSession = async ( + session: HarnessSession, + answer: ( + keep: ( + watch: (reader: AbortController) => void, + ) => (reader: AbortController) => void, + ) => Promise | Response, + ) => { + let isKept = false + const keep = + (watch: (reader: AbortController) => void) => + (reader: AbortController) => { + isKept = true + // Both stream helpers abort the reader when the stream ends. + reader.signal.addEventListener('abort', () => void session.close(), { + once: true, + }) + watch(reader) + } + try { + return await answer(keep) + } finally { + if (!isKept) void session.close() + } + } + // ponytail: without `media.secret` (only allowed without a log store), a // random key per handler, so signed URLs stop working after a restart. const secret = @@ -566,7 +597,9 @@ export function createHarnessHandler

( : undefined const session = await openToRead(user, payload.threadId) if (!session) return json({ error: 'not found' }, 404) - return mediaResponse(request, session, payload.id) + return withSession(session, () => + mediaResponse(request, session, payload.id), + ) } /** @@ -592,17 +625,20 @@ export function createHarnessHandler

( if (running?.operation(runId)) { const { threadId } = running const session = await openToRead(principal, threadId) - const turn = session?.operation(runId) - if (!session || !turn) return json({ error: 'unknown run' }, 404) - const startedCursor = session - .snapshot() - .activeOperations.find((item) => item.id === runId)?.startedCursor - const from = resumeFrom ?? startedCursor ?? '0' - return cursorResponse( - (signal) => turn.events({ from, signal }), - true, - closeOnRevoke(principal, threadId), - ) + if (!session) return json({ error: 'unknown run' }, 404) + return withSession(session, (keep) => { + const turn = session.operation(runId) + if (!turn) return json({ error: 'unknown run' }, 404) + const startedCursor = session + .snapshot() + .activeOperations.find((item) => item.id === runId)?.startedCursor + const from = resumeFrom ?? startedCursor ?? '0' + return cursorResponse( + (signal) => turn.events({ from, signal }), + true, + keep(closeOnRevoke(principal, threadId)), + ) + }) } // Another instance runs the turn. A durable host finds its thread in the // run store, and its events in the shared log. Not a child run or an @@ -623,14 +659,16 @@ export function createHarnessHandler

( if (!shared) 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. - return cursorResponse( - (signal) => - followOperation( - shared.events({ from: resumeFrom ?? '0', signal }), - runId, - ), - true, - closeOnRevoke(principal, threadId), + return withSession(shared, (keep) => + cursorResponse( + (signal) => + followOperation( + shared.events({ from: resumeFrom ?? '0', signal }), + runId, + ), + true, + keep(closeOnRevoke(principal, threadId)), + ), ) } @@ -700,54 +738,59 @@ export function createHarnessHandler

( const ids = { inputId: params.runId, runId: params.runId, ...user } const lastEventId = request.headers.get('Last-Event-ID') const watch = closeOnRevoke(principal, params.threadId) - if (params.resume && params.resume.length > 0) { - const answers = { resume: params.resume, ...ids } - // Only the session class takes the run id. `host.open` gives one. - const receipt = await (session instanceof SessionImpl - ? session.resolveForRun(answers) - : session.resolve(answers)) + return withSession(session, async (keep) => { + if (params.resume && params.resume.length > 0) { + const answers = { resume: params.resume, ...ids } + // Only the session class takes the run id. `host.open` gives one. + const receipt = await (session instanceof SessionImpl + ? session.resolveForRun(answers) + : session.resolve(answers)) + if (receipt.status === 'rejected') { + return json({ error: receipt.reason }, 409) + } + if (!receipt.operationId) return json({ error: 'rejected' }, 409) + const { operationId } = receipt + return cursorResponse( + (signal) => + followOperation( + session.events({ + ...(durable && lastEventId ? { from: lastEventId } : {}), + signal, + }), + operationId, + ), + durable !== undefined, + keep(watch), + ) + } + if (message === undefined) { + return json({ error: 'no user message' }, 400) + } + // AG-UI `forwardedProps` are the input's context. + const { forwardedProps } = params + const operation = session.send(message, { + ...ids, + ...(Object.keys(forwardedProps).length > 0 + ? { context: forwardedProps } + : {}), + }) + // A turn that the host refused throws, and the answer is 403. + const receipt = await turnReceipt(operation) if (receipt.status === 'rejected') { - return json({ error: receipt.reason }, 409) + return json({ error: receipt.reason ?? 'rejected' }, 409) } - if (!receipt.operationId) return json({ error: 'rejected' }, 409) - const { operationId } = receipt + // A durable host sends each log cursor as the event id. A client cut + // off sends the same request with `Last-Event-ID`, maybe to another + // instance: the run id finds the same turn, and it streams from + // there. Without a log a cursor means nothing on another instance, + // so no id. + const from = (durable && lastEventId) || '0' return cursorResponse( - (signal) => - followOperation( - session.events({ - ...(durable && lastEventId ? { from: lastEventId } : {}), - signal, - }), - operationId, - ), + (signal) => operation.events({ from, signal }), durable !== undefined, - watch, + keep(watch), ) - } - if (message === undefined) return json({ error: 'no user message' }, 400) - // AG-UI `forwardedProps` are the input's context. - const { forwardedProps } = params - const operation = session.send(message, { - ...ids, - ...(Object.keys(forwardedProps).length > 0 - ? { context: forwardedProps } - : {}), }) - // A turn that the host refused throws, and the answer is 403. - const receipt = await turnReceipt(operation) - if (receipt.status === 'rejected') { - return json({ error: receipt.reason ?? 'rejected' }, 409) - } - // A durable host sends each log cursor as the event id. A client cut - // off sends the same request with `Last-Event-ID`, maybe to another - // instance: the run id finds the same turn, and it streams from there. - // Without a log a cursor means nothing on another instance, so no id. - const from = (durable && lastEventId) || '0' - return cursorResponse( - (signal) => operation.events({ from, signal }), - durable !== undefined, - watch, - ) } if (request.method === 'GET' && route === 'run') { @@ -762,7 +805,9 @@ export function createHarnessHandler

( if (!session) { return json({ messages: [], activeRun: null, interrupts: null }) } - return json(await hydration(session, runsOf(principal))) + return withSession(session, async () => + json(await hydration(session, runsOf(principal))), + ) } if (request.method === 'GET' && route === 'events') { @@ -774,15 +819,20 @@ export function createHarnessHandler

( request.headers.get('Last-Event-ID') ?? url.searchParams.get('from') ?? undefined - return sseResponse( - async function* (signal) { - const events = session.events({ ...(from ? { from } : {}), signal }) - for await (const entry of events) { - const frame: HostFrame = { type: 'harness.event', ...entry } - yield { id: entry.cursor, data: frame } - } - }, - closeOnRevoke(principal, threadId), + return withSession(session, (keep) => + sseResponse( + async function* (signal) { + const events = session.events({ + ...(from ? { from } : {}), + signal, + }) + for await (const entry of events) { + const frame: HostFrame = { type: 'harness.event', ...entry } + yield { id: entry.cursor, data: frame } + } + }, + keep(closeOnRevoke(principal, threadId)), + ), ) } @@ -805,7 +855,9 @@ 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) - return json(await applyInput({ harness, session, input, ...user })) + return withSession(session, async () => + json(await applyInput({ harness, session, input, ...user })), + ) } if ( @@ -816,10 +868,12 @@ export function createHarnessHandler

( if (!threadId) return json({ error: 'threadId is required' }, 400) const session = await openToRead(principal, threadId) if (!session) return json({ error: 'not found' }, 404) - return json( - route === 'transcript' - ? await session.transcript() - : describeForClient(harness, session), + return withSession(session, async () => + json( + route === 'transcript' + ? await session.transcript() + : describeForClient(harness, session), + ), ) } @@ -831,7 +885,9 @@ export function createHarnessHandler

( return json({ error: 'subagentRunId is required' }, 400) const session = await openToRead(principal, threadId) if (!session) return json({ error: 'not found' }, 404) - return json(await session.agentTranscript(subagentRunId)) + return withSession(session, async () => + json(await session.agentTranscript(subagentRunId)), + ) } if (request.method === 'GET' && route === 'snapshot') { @@ -839,7 +895,7 @@ export function createHarnessHandler

( if (!threadId) return json({ error: 'threadId is required' }, 400) const session = await openToRead(principal, threadId) if (!session) return json({ error: 'not found' }, 404) - return json(session.snapshot()) + return withSession(session, () => json(session.snapshot())) } if (request.method === 'GET' && route === 'sessions') { @@ -942,12 +998,15 @@ export function createHarnessHandler

( if (!threadId || !name) return json({ error: 'threadId and name are required' }, 400) const session = await openFor(principal, threadId) - const record = await session.putMedia({ - body: request.body ?? new Uint8Array(), - mimeType: request.headers.get('Content-Type') ?? '', - name, - }) - return json(record) + return withSession(session, async () => + json( + await session.putMedia({ + body: request.body ?? new Uint8Array(), + mimeType: request.headers.get('Content-Type') ?? '', + name, + }), + ), + ) } if ( @@ -960,24 +1019,29 @@ export function createHarnessHandler

( return json({ error: 'threadId and id are required' }, 400) const session = await openToRead(principal, threadId) if (!session) return json({ error: 'not found' }, 404) - if (route === 'media') return await mediaResponse(request, session, id) - if (!(await session.getMedia(id))) - return json({ error: 'not found' }, 404) - const expiresAt = Date.now() + SIGNED_URL_TTL_MS - const payload: SignedMedia = { - tenantId: principal?.tenantId ?? '', - threadId, - id, - userId: principal?.id ?? '', - exp: String(expiresAt), - } - // An empty field stays out of the URL. The served URL reads it as ''. - const fields = Object.entries(payload).filter(([, value]) => value !== '') - const query = new URLSearchParams([ - ...fields, - ['sig', await sign(payload)], - ]) - return json({ path: `media?${query}`, expiresAt }) + return withSession(session, async () => { + if (route === 'media') return mediaResponse(request, session, id) + if (!(await session.getMedia(id))) { + return json({ error: 'not found' }, 404) + } + const expiresAt = Date.now() + SIGNED_URL_TTL_MS + const payload: SignedMedia = { + tenantId: principal?.tenantId ?? '', + threadId, + id, + userId: principal?.id ?? '', + exp: String(expiresAt), + } + // An empty field stays out of the URL. The served URL reads it as ''. + const fields = Object.entries(payload).filter( + ([, value]) => value !== '', + ) + const query = new URLSearchParams([ + ...fields, + ['sig', await sign(payload)], + ]) + return json({ path: `media?${query}`, expiresAt }) + }) } return json({ error: 'not found' }, 404) } @@ -1127,6 +1191,10 @@ export function handleHarnessSocket

( }) const session = await opened if (!session) return refuse('forbidden') + // The socket holds the handle until it closes. + const release = () => void session.close() + if (reader.signal.aborted) release() + else reader.signal.addEventListener('abort', release, { once: true }) watchRevoke( host, { principal, threadId, signal: reader.signal }, diff --git a/packages/ai-harness/src/session/base.ts b/packages/ai-harness/src/session/base.ts index 12e52eab9f..96454b2a2c 100644 --- a/packages/ai-harness/src/session/base.ts +++ b/packages/ai-harness/src/session/base.ts @@ -387,6 +387,8 @@ export abstract class SessionBase { /** The last `recover()`. The calls run one at a time. */ protected recovery: Promise = Promise.resolve() protected readonly onClose: () => void + /** Tells the host each time the thread goes idle. See `SessionDependencies`. */ + protected readonly notifyIdle: (() => void) | undefined protected checkpoint: AnyChatMiddleware | undefined /** Gives each durable tool call its steps in the log. One per session. */ protected bindTool: DurableBind | undefined @@ -493,6 +495,7 @@ export abstract class SessionBase { this.eventListeners, ) this.onClose = deps.onClose + this.notifyIdle = deps.onIdle this.hostId = deps.hostId this.index = deps.index this.lease = deps.lease @@ -1214,10 +1217,11 @@ export abstract class SessionBase { } /** - * No input being stored, no queued or running turn, no waiting steer, and - * no running agent. + * @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. */ - protected isIdle(): boolean { + isIdle(): boolean { // A follow-up that waits for its run keeps the thread busy. const hasFollowUps = [...this.agentChains.values()].some( (chain) => chain.followUps.length > 0, @@ -1269,6 +1273,12 @@ export abstract class SessionBase { this.draining ? undefined : this.drainInbox(), ) for (const listener of [...this.idleListeners]) listener() + this.notifyIdle?.() + } + + /** @internal `close()` ran: the session takes no new work. */ + isClosing(): boolean { + return this.closing !== undefined } /** diff --git a/packages/ai-harness/src/session/handle.ts b/packages/ai-harness/src/session/handle.ts index b1ce2e2a79..973f9b9ccd 100644 --- a/packages/ai-harness/src/session/handle.ts +++ b/packages/ai-harness/src/session/handle.ts @@ -5,6 +5,8 @@ import type { AnyHarness } from '../define' import type { AgentSendOptions, HarnessInput, + Operation, + OperationKind, Principal, SendOptions, UserInput, @@ -20,6 +22,51 @@ export interface HandleGuard { isRevoked: () => boolean } +/** How a handle reaches the shared session of its thread. The host gives one. */ +export interface HandleLease { + /** The open session of the thread, or `undefined` when none is open now. */ + current: () => SessionImpl | undefined + /** + * Open the session of the thread again, with the checks of `host.open`. + * Throws at once when the host is closed. + */ + reopen: () => Promise> + /** One more holder: the session does not close on idle while it is held. */ + hold: () => void + /** + * One holder less. `dropped`: the garbage collector dropped the handle, + * so the host warns. Resolves after the close that the release starts. + */ + release: (options?: { dropped?: boolean }) => Promise +} + +/** What the finalizer of a handle needs. It must not hold the handle. */ +interface HandleToken { + isReleased: boolean + release: (options: { dropped: boolean }) => Promise +} + +/** Release a handle that its holder dropped without `close()`. */ +function releaseDropped(token: HandleToken) { + if (token.isReleased) return + token.isReleased = true + void token.release({ dropped: true }) +} + +const dropped = new FinalizationRegistry(releaseDropped) +const tokens = new WeakMap() + +/** + * @internal Run the finalizer of `handle`, as if the garbage collector + * dropped it. Tests use it, as a test cannot force a collection. + */ +export function finalizeHandle(handle: object) { + const token = tokens.get(handle) + if (!token) return + dropped.unregister(token) + releaseDropped(token) +} + /** The action that `authorize` gets for an input. */ function actionOf(input: HarnessInput): AccessAction { switch (input.op) { @@ -50,32 +97,235 @@ function actionOf(input: HarnessInput): AccessAction { } } +/** A listener that waits for the operation of a session that opens again. */ +interface Waiting { + member: 'onEvent' | 'onMedia' + args: Array + stop: (() => void) | undefined +} + +/** + * The operation of a call on a session that opens again. `call` runs on the + * session once it opened, and each member waits for that. `id` is known only + * then: await `receipt` before you read it. + */ +function laterOperation>({ + kind, + agent, + opening, + call, +}: { + kind: OperationKind + agent: string | undefined + opening: Promise + call: (session: TSession) => TOperation +}) { + let bound: TOperation | undefined + const waiting = new Set() + // Wrapped: an operation is thenable, so a promise of it would wait for its + // result. + const started = opening.then((session) => { + const operation = call(session) + bound = operation + for (const entry of waiting) { + entry.stop = Reflect.apply(operation[entry.member], operation, entry.args) + } + return { operation } + }) + /** `use` of the operation, once it started. */ + const ready = (use: (operation: TOperation) => R) => + started.then(({ operation }) => use(operation)) + const receipt = ready((operation) => operation.receipt) + // A caller that never awaits a failed open gets no unhandled rejection. + receipt.catch(() => {}) + const listen = + (member: Waiting['member']) => + (...args: Array) => { + if (bound) return Reflect.apply(bound[member], bound, args) + const entry: Waiting = { member, args, stop: undefined } + waiting.add(entry) + return () => { + waiting.delete(entry) + entry.stop?.() + } + } + async function* events( + options?: Parameters['events']>[0], + ) { + const { operation } = await started + yield* operation.events(options) + } + async function* stream( + options?: Parameters['stream']>[0], + ) { + const { operation } = await started + yield* operation.stream(options) + } + return { + get id() { + if (!bound) { + throw new Error( + 'The session of this operation opens again. Await operation.receipt before you read operation.id.', + ) + } + return bound.id + }, + get kind() { + return bound?.kind ?? kind + }, + get agent() { + return bound?.agent ?? agent + }, + status: () => bound?.status() ?? 'accepted', + receipt, + then: ( + onfulfilled?: (value: unknown) => TResult1 | PromiseLike, + onrejected?: (reason: unknown) => TResult2 | PromiseLike, + ) => ready((operation) => operation).then(onfulfilled, onrejected), + events, + stream, + onEvent: listen('onEvent'), + onMedia: listen('onMedia'), + cancel: (reason?: string) => ready((operation) => operation.cancel(reason)), + // An agent run takes messages. + ...(kind === 'agent' + ? { + send: (...args: Array) => + ready((operation) => + 'send' in operation && typeof operation.send === 'function' + ? Reflect.apply(operation.send, operation, args) + : Promise.reject(new Error('This operation takes no message.')), + ), + } + : {}), + } +} + +/** + * The members that read the stored state of the session. On a session that + * closed, they open it again. The other members that return a value at once + * read the last state of the session. + */ +const REOPENING_READS: ReadonlySet = new Set([ + 'transcript', + 'agentTranscript', + 'entry', + 'settled', +]) + type Session = SessionImpl /** - * The session that `host.open` returns, bound to `principal`: a call that - * passes no `principal` acts as it. It has every member of the session, and - * `this` stays the session. + * A handle to the shared session of a thread, that `host.open` returns, + * bound to `principal`: a call that passes no `principal` acts as it. It has + * every member of the session, and `this` stays the session. + * + * `close()` (or `await using`) releases only this handle: `lease` counts the + * handles, and the host closes the session when none holds it and it is + * idle. A handle that its holder drops is released by the garbage + * collector, with a warning. A call on a handle whose session closed opens + * it again through `lease`. A call through a released handle holds the + * session while it runs. * * With `guard` (the host has access), each action asks `authorize` first. An * input is checked before the session stores it. A call that passes another * principal throws. After `host.revoke`, every call throws * `AccessDeniedError`. Reads ask `authorize` once, when the host opens the - * handle. + * handle, and again when it opens the session again. */ export function sessionHandle({ session, principal, guard, + lease, }: { session: Session principal: Principal | undefined guard?: HandleGuard + lease: HandleLease }) { const { threadId } = session + /** The last session that this handle used. */ + let last = session + const token: HandleToken = { + isReleased: false, + release: (options) => lease.release(options), + } const live = (action: AccessAction) => { if (guard?.isRevoked()) throw new AccessDeniedError({ action, threadId }) } + /** The open session, or `undefined` when it closed. */ + const open = () => { + const current = lease.current() + if (current) last = current + return current + } + /** + * A released handle holds the session while a call runs. The function it + * returns lets go, once. + */ + const holdForCall = () => { + if (!token.isReleased) return () => {} + lease.hold() + let isHeld = true + return () => { + if (!isHeld) return + isHeld = false + void lease.release() + } + } + /** Run `call` on the open session. On a session that closed, open it again. */ + const run = async (call: (target: Session) => Promise) => { + const done = holdForCall() + try { + const target = open() ?? (last = await lease.reopen()) + return await call(target) + } finally { + done() + } + } + /** + * `run` for a call that returns an operation at once. On a session that + * closed, the operation waits for the open. + */ + const runOperation = >( + kind: OperationKind, + call: (target: Session) => T, + agent?: string, + ) => { + const done = holdForCall() + const target = open() + try { + if (target) { + const operation = call(target) + operation.then(done, done) + return operation + } + const opening = lease.reopen().then((reopened) => { + last = reopened + return reopened + }) + opening.catch(done) + return laterOperation({ + kind, + agent, + opening, + call: (reopened) => { + try { + const operation = call(reopened) + operation.then(done, done) + return operation + } catch (error) { + done() + throw error + } + }, + }) + } catch (error) { + done() + throw error + } + } /** Who acts in a call that passes `passed`. */ const actor = (passed: Principal | undefined) => { if (passed === undefined) return principal @@ -92,102 +342,182 @@ export function sessionHandle({ const acting = actor(options.principal) return acting ? { ...options, principal: acting } : options } - /** A call that stores an input: the session checks the input first. */ - const storing = (action: AccessAction, call: () => T) => { + /** `call` on `target`, which stores an input: the session checks it first. */ + const checkedOn = ( + target: Session, + call: (target: Session) => T, + ) => + guard + ? target.checked( + ({ input, run: accessRun }) => + guard.authorize(actionOf(input), accessRun), + () => call(target), + ) + : call(target) + /** A call that stores an input and resolves to its receipt. */ + const storing = ( + action: AccessAction, + call: (target: Session) => Promise, + ) => { live(action) - if (!guard) return call() - return session.checked( - ({ input, run }) => guard.authorize(actionOf(input), run), - call, - ) + return run((target) => checkedOn(target, call)) + } + /** A call that stores an input and returns its operation. */ + const storingOperation = >( + kind: OperationKind, + call: (target: Session) => T, + agent?: string, + ) => { + live('send') + return runOperation(kind, (target) => checkedOn(target, call), agent) } /** A call that stores no input: it runs after `authorize` allows `action`. */ - const acting = async (action: AccessAction, call: () => Promise) => { + const acting = ( + action: AccessAction, + call: (target: Session) => Promise, + ) => { live(action) - await guard?.authorize(action) - return call() + return run(async (target) => { + await guard?.authorize(action) + return call(target) + }) } const agentOf = (name: string) => ({ run: (input?: unknown, options?: AgentRunOptions) => - storing('send', () => - session.agentHandle(name, principal).run(input, options), + storingOperation( + 'agent', + (target) => target.agentHandle(name, principal).run(input, options), + name, ), start: (input?: unknown, options?: AgentStartOptions) => - storing('send', () => - session - .agentHandle(name, actor(options?.principal)) - .start(input, options), + storingOperation( + 'agent', + (target) => + target + .agentHandle(name, actor(options?.principal)) + .start(input, options), + name, ), }) const agents = new Proxy(session.agents, { get: (_target, name) => typeof name === 'string' ? agentOf(name) : undefined, }) - const methods = new Map( - Object.entries({ + /** Release this handle. A second call does nothing. */ + const close = () => { + if (token.isReleased) return Promise.resolve() + token.isReleased = true + dropped.unregister(token) + return lease.release() + } + type EventsOptions = Parameters['events']>[0] + /** `events` of a session that closed, or of a released handle. */ + async function* heldEvents(options: EventsOptions) { + // A released handle holds the session while it reads. + const done = holdForCall() + try { + const target = open() ?? (last = await lease.reopen()) + yield* target.events(options) + } finally { + done() + } + } + const events = (options?: EventsOptions) => { + live('read') + const target = open() + const isHeld = target !== undefined && !token.isReleased + return isHeld ? target.events(options) : heldEvents(options) + } + const methods = new Map([ + ...Object.entries({ send: (message: UserInput, options?: SendOptions | AgentSendOptions) => - storing('send', () => - options?.to === undefined - ? session.send(message, as(options ?? {})) - : session.send(message, as(options)), + storingOperation( + options?.to === undefined ? 'chat' : 'agent', + (target) => + options?.to === undefined + ? target.send(message, as(options ?? {})) + : target.send(message, as(options)), ), tool: (options: Parameters['tool']>[0]) => - storing('send', () => session.tool(as(options))), + storingOperation('chat', (target) => target.tool(as(options))), continue: (options?: Parameters['continue']>[0]) => - storing('send', () => session.continue(as(options ?? {}))), + storingOperation('chat', (target) => + target.continue(as(options ?? {})), + ), command: (options: Parameters['command']>[0]) => - storing('send', () => session.command(as(options))), + storingOperation('command', (target) => target.command(as(options))), reset: (options?: Parameters['reset']>[0]) => - storing('send', () => session.reset(as(options ?? {}))), + storing('send', (target) => target.reset(as(options ?? {}))), setDelivery: (options: Parameters['setDelivery']>[0]) => - storing('send', () => session.setDelivery(as(options))), + storing('send', (target) => target.setDelivery(as(options))), resolve: (options: Parameters['resolve']>[0]) => - storing('answer', () => session.resolve(as(options))), + storing('answer', (target) => target.resolve(as(options))), resolveForRun: ( options: Parameters['resolveForRun']>[0], - ) => storing('answer', () => session.resolveForRun(as(options))), + ) => storing('answer', (target) => target.resolveForRun(as(options))), answer: (options: Parameters['answer']>[0]) => - storing('answer', () => session.answer(as(options))), + storing('answer', (target) => target.answer(as(options))), cancel: (options?: Parameters['cancel']>[0]) => - storing('cancel', () => session.cancel(as(options ?? {}))), + storing('cancel', (target) => target.cancel(as(options ?? {}))), cancelInput: ( - target: string, + inputId: string, options?: Parameters['cancelInput']>[1], ) => - storing('cancel', () => session.cancelInput(target, as(options ?? {}))), + storing('cancel', (target) => + target.cancelInput(inputId, as(options ?? {})), + ), configure: (options: Parameters['configure']>[0]) => - storing('configure', () => session.configure(as(options))), + storing('configure', (target) => target.configure(as(options))), setConfig: (options: Parameters['setConfig']>[0]) => - storing('configure', () => session.setConfig(as(options))), + storing('configure', (target) => target.setConfig(as(options))), agents, agent: (name: string) => - session.agent(name) === undefined ? undefined : agentOf(name), + (open() ?? last).agent(name) === undefined ? undefined : agentOf(name), // These store no input, so the handle asks `authorize` itself. revert: (messageId: string) => - acting('send', () => session.revert(messageId)), - unrevert: () => acting('send', () => session.unrevert()), + acting('send', (target) => target.revert(messageId)), + unrevert: () => acting('send', (target) => target.unrevert()), background: (toolCallId?: string) => - acting('send', () => session.background(toolCallId)), + acting('send', (target) => target.background(toolCallId)), append: (records: Parameters['append']>[0]) => - acting('send', () => session.append(records)), + acting('send', (target) => target.append(records)), putMedia: (upload: Parameters['putMedia']>[0]) => - acting('send', () => session.putMedia(upload)), + acting('send', (target) => target.putMedia(upload)), updateEntry: (update: Parameters['updateEntry']>[0]) => - acting('configure', () => session.updateEntry(update)), - reload: () => acting('configure', () => session.reload()), + acting('configure', (target) => target.updateEntry(update)), + reload: () => acting('configure', (target) => target.reload()), + events, + close, }), - ) + [Symbol.asyncDispose, close], + ]) // ponytail: a Proxy, so the handle has every member of the session and - // stays an instance of it. Only the members above change. - return new Proxy(session, { - get(target, key) { + // stays an instance of it. Only the members above change. The other + // members run on the open session, or on the last one when it closed. + const handle = new Proxy(session, { + get(_target, key) { if (methods.has(key)) return methods.get(key) + const target = open() ?? last const value: unknown = Reflect.get(target, key) if (typeof value !== 'function') return value + if (REOPENING_READS.has(key)) { + return (...args: Array) => { + live('read') + return run(async (reopened) => { + const member: unknown = Reflect.get(reopened, key) + return typeof member === 'function' + ? Reflect.apply(member, reopened, args) + : member + }) + } + } return (...args: Array) => { live('read') return Reflect.apply(value, target, args) } }, }) + tokens.set(handle, token) + dropped.register(handle, token, token) + return handle } diff --git a/packages/ai-harness/src/session/harness-session.ts b/packages/ai-harness/src/session/harness-session.ts index 49fcb68d81..a26150d8b9 100644 --- a/packages/ai-harness/src/session/harness-session.ts +++ b/packages/ai-harness/src/session/harness-session.ts @@ -32,8 +32,10 @@ import type { HarnessInput, SettlementOutcome } from '../types' import type { WaitingSteer } from './types' /** - * A live harness session: one conversation (`threadId`) with its plugins, - * operations, inbox, and event stream. Open one with `host.open()`. + * A handle to a live harness session: one conversation (`threadId`) with its + * plugins, operations, inbox, and event stream. Open one with `host.open()`. + * Each open gives its own handle. Release it with `close()`, or with + * `await using`, when you are done. */ export interface HarnessSession< THarness extends AnyHarness = AnyHarness, @@ -47,7 +49,39 @@ export interface HarnessSession< | 'checked' | 'cancelWorkOf' | 'agentHandle' -> {} + | 'isIdle' + | 'isClosing' + // A handle closes itself, not the session. + | 'close' +> { + /** + * Release this handle. The session stays open while another handle holds + * it. When no handle holds it and it is idle, the host closes it after + * `lifecycle.idleCloseMs`. A second call does nothing. A call on the handle + * after the session closed opens it again. To stop the work of every + * session, call `host.close()`. + * + * @example + * ```ts + * const session = await host.open(assistant, { threadId }) + * try { + * await session.send('Hello.') + * } finally { + * await session.close() + * } + * ``` + */ + close: () => Promise + /** + * `close()`, for `await using`. + * + * @example + * ```ts + * await using session = await host.open(assistant, { threadId }) + * ``` + */ + [Symbol.asyncDispose]: () => Promise +} /** * The session class. The package does not export it: apps get the @@ -61,6 +95,14 @@ export class SessionImpl< return this } + /** + * @internal Close the session, for `await using` on the class. A handle of + * `host.open` releases itself instead. + */ + [Symbol.asyncDispose](): Promise { + return this.close() + } + /** @internal Mount session plugins and replay inputs left in the inbox. */ async open(): Promise { // First: plugins publish events while they mount. diff --git a/packages/ai-harness/src/session/requests.ts b/packages/ai-harness/src/session/requests.ts index 1bd80330af..d148df3b3c 100644 --- a/packages/ai-harness/src/session/requests.ts +++ b/packages/ai-harness/src/session/requests.ts @@ -902,19 +902,17 @@ export abstract class SessionRequests< private readonly followers = new Set>() /** - * Stop every running operation, wait for them, then dispose session plugins. - * Safe to call twice. + * @internal Close the session for every handle: stop every running + * operation, wait for them, then dispose session plugins. Safe to call + * twice. The host calls it (`host.close`, the idle close, and + * `host.sessions.delete`). The `close()` of a handle releases only that + * handle. * * With `recoverable`, the running turns and agent runs stop as on a host * crash: the log gets none of their later writes, they do not settle, and * the run store gets no aborted state. They give their leases back, so the - * next host that opens the thread runs them again at once. Use it when a - * durable host shuts down, for example for a deploy. - * - * @example - * ```ts - * await session.close({ recoverable: true }) - * ``` + * next host that opens the thread runs them again at once. `host.close` + * passes it when a durable host shuts down, for example for a deploy. */ close(options?: { recoverable?: boolean }): Promise { this.closing ??= (async () => { diff --git a/packages/ai-harness/src/session/types.ts b/packages/ai-harness/src/session/types.ts index 3fa0986a67..9fca9982a5 100644 --- a/packages/ai-harness/src/session/types.ts +++ b/packages/ai-harness/src/session/types.ts @@ -247,6 +247,11 @@ export interface SessionDependencies { /** Gets each event the session publishes, after the feed has it. */ onEvent: (event: StreamChunk) => void onClose: () => void + /** + * Called each time the thread goes idle. The host starts the idle close of + * a session that no handle holds. + */ + onIdle?: () => void /** * The host checks access (`access` on the host). Then an entry with no * owner stays without one, and provider keys read env vars only with diff --git a/packages/ai-harness/src/worker.ts b/packages/ai-harness/src/worker.ts index 5612913bd4..7d4b113d6e 100644 --- a/packages/ai-harness/src/worker.ts +++ b/packages/ai-harness/src/worker.ts @@ -93,6 +93,8 @@ export async function runHarnessWorker( } } finally { reader.abort() + // The worker holds one handle for its life. + await session?.close() await host.close() } } diff --git a/packages/ai-harness/tests/agent-runs.test.ts b/packages/ai-harness/tests/agent-runs.test.ts index 134164d51f..f8a34cd24f 100644 --- a/packages/ai-harness/tests/agent-runs.test.ts +++ b/packages/ai-harness/tests/agent-runs.test.ts @@ -185,11 +185,10 @@ describe('follow-ups to an agent run', () => { const receipt = await run.send('Add a title.', { mode: 'followUp' }) const followUp = session.operation(receipt.operationId ?? '') - await session.close() + await host.close() await expect(followUp).rejects.toThrow('Session closed.') expect(model.calls).toHaveLength(1) - await host.close() }) it('cancels a queued follow-up when a log failure stops the session', async () => { diff --git a/packages/ai-harness/tests/auth.test.ts b/packages/ai-harness/tests/auth.test.ts index 04703afed3..4ad9eaef10 100644 --- a/packages/ai-harness/tests/auth.test.ts +++ b/packages/ai-harness/tests/auth.test.ts @@ -310,9 +310,8 @@ describe('commands on close', () => { ) const command = session.command({ name: 'stuck' }) await running - await session.close() - expect(command.status()).toBe('cancelled') await host.close() + expect(command.status()).toBe('cancelled') }) }) diff --git a/packages/ai-harness/tests/fork.test.ts b/packages/ai-harness/tests/fork.test.ts index fb5e89ff29..d6148e3888 100644 --- a/packages/ai-harness/tests/fork.test.ts +++ b/packages/ai-harness/tests/fork.test.ts @@ -11,7 +11,7 @@ import { import { mediaIdOf, mediaPart } from '../src/media-entry' import { messageTexts, mockAdapter, text, toolCall } from './helpers' import type { ModelMessage } from '@tanstack/ai' -import type { HarnessPersistence } from '../src' +import type { HarnessPersistence, HostLifecycleOptions } from '../src' const SOURCE = 'source' const FORK = 'fork' @@ -56,14 +56,17 @@ describe.each([ { name: 'a host with a message store', durable: false }, { name: 'a durable host', durable: true }, ])('host.fork on $name', ({ durable }) => { - const setup = (persistence: HarnessPersistence) => { + const setup = ( + persistence: HarnessPersistence, + lifecycle: HostLifecycleOptions = {}, + ) => { const { adapter, calls } = mockAdapter([ () => text('A1'), () => text('A2'), () => text('A3'), ]) const harness = defineHarness({ name: 'test/fork', adapter }) - const host = createHarnessHost({ persistence }) + const host = createHarnessHost({ persistence, lifecycle }) return { harness, host, calls } } @@ -108,15 +111,17 @@ describe.each([ it('closes a source that only the fork opened, and keeps one the app has open', async () => { const persistence = durable ? durablePersistence() : memoryPersistence() - const { harness, host } = setup(persistence) + // A session that no handle holds closes at once when it is idle. + const { harness, host } = setup(persistence, { idleCloseMs: 0 }) const source = await host.open(harness, { threadId: SOURCE }) await source.send('one') await host.fork(harness, { threadId: SOURCE, newThreadId: FORK }) // The app opened the source: the fork left it open. expect(host.isOpen(harness, { threadId: SOURCE })).toBe(true) await source.close() + expect(host.isOpen(harness, { threadId: SOURCE })).toBe(false) await host.fork(harness, { threadId: SOURCE, newThreadId: 'other' }) - // The fork opened the source itself, so it closed it again. + // The fork opened the source itself, and released it at the end. expect(host.isOpen(harness, { threadId: SOURCE })).toBe(false) await host.close() }) diff --git a/packages/ai-harness/tests/forwarding.test.ts b/packages/ai-harness/tests/forwarding.test.ts index 50aad7da15..f0bd370d5d 100644 --- a/packages/ai-harness/tests/forwarding.test.ts +++ b/packages/ai-harness/tests/forwarding.test.ts @@ -403,12 +403,10 @@ describe('an input that reaches a host that does not own the thread', () => { ownerId: 'busy-host', until: Date.now() + 60_000, }) - const other = await createHarnessHost({ persistence }).open( - harnessWith(mockAdapter([]).adapter), - { - threadId: THREAD, - }, - ) + const otherHost = createHarnessHost({ persistence }) + const other = await otherHost.open(harnessWith(mockAdapter([]).adapter), { + threadId: THREAD, + }) let last = '' let operation: Operation | undefined for (let index = 0; index < 510; index++) { @@ -419,7 +417,7 @@ describe('an input that reaches a host that does not own the thread', () => { expect(forwarded.size).toBe(500) expect(forwarded.has(last)).toBe(true) // The owner runs the input. The follower here ends with the session. - await other.close() + await otherHost.close() await expect(operation).rejects.toThrow('Session closed.') }) diff --git a/packages/ai-harness/tests/handles.test.ts b/packages/ai-harness/tests/handles.test.ts new file mode 100644 index 0000000000..5ca42eef61 --- /dev/null +++ b/packages/ai-harness/tests/handles.test.ts @@ -0,0 +1,275 @@ +import { afterEach, describe, expect, it, vi } from 'vitest' +import { memoryPersistence } from '@tanstack/ai-persistence' +import { createHarnessHost, defineHarness } from '../src' +import { createHarnessHandler } from '../src/server' +import { hostCountsOf } from '../src/host' +import { finalizeHandle } from '../src/session/handle' +import { mockAdapter, text } from './helpers' +import type { HostLifecycleOptions } from '../src' + +const BASE = 'http://x/api/harness' + +/** A memory host with `lifecycle`, and a harness that answers each turn. */ +function setup(lifecycle: HostLifecycleOptions = {}) { + const { adapter, calls } = mockAdapter(() => text('Hello back.')) + const harness = defineHarness({ name: 'test/handles', adapter }) + const host = createHarnessHost({ + persistence: memoryPersistence(), + lifecycle, + }) + const isOpen = (threadId: string) => host.isOpen(harness, { threadId }) + const open = (threadId: string) => host.open(harness, { threadId }) + return { host, harness, calls, isOpen, open } +} + +/** The warnings of dropped handles that `console.warn` got. */ +function droppedWarnings() { + const warn = vi.spyOn(console, 'warn').mockImplementation(() => {}) + return () => + warn.mock.calls.filter(([message]) => + String(message).includes('dropped without close()'), + ) +} + +/** The warnings of a host past `maxOpen` that `console.warn` got. */ +function fullWarnings() { + const warn = vi.spyOn(console, 'warn').mockImplementation(() => {}) + return () => + warn.mock.calls.filter(([message]) => + String(message).includes('lifecycle.maxOpen'), + ) +} + +afterEach(() => { + vi.useRealTimers() + vi.restoreAllMocks() +}) + +describe('handle.close', () => { + it('keeps the session open while another handle holds it, and closes it with the last handle', async () => { + const { host, isOpen, open } = setup({ idleCloseMs: 0 }) + const first = await open('t1') + const second = await open('t1') + + await first.close() + // A second close of the same handle does not release the other one. + await first.close() + expect(isOpen('t1')).toBe(true) + + await second.close() + expect(isOpen('t1')).toBe(false) + expect(hostCountsOf(host)).toMatchObject({ sessions: 0, handles: 0 }) + await host.close() + }) + + it('releases the handle with await using', async () => { + const { host, isOpen, open } = setup({ idleCloseMs: 0 }) + { + await using session = await open('t1') + await session.send('Hello.') + expect(isOpen('t1')).toBe(true) + } + expect(isOpen('t1')).toBe(false) + await host.close() + }) +}) + +describe('the idle close', () => { + it('closes a session that no handle holds after idleCloseMs', async () => { + const { host, isOpen, open } = setup({ idleCloseMs: 1000 }) + const session = await open('t1') + vi.useFakeTimers() + + await session.close() + await vi.advanceTimersByTimeAsync(999) + expect(isOpen('t1')).toBe(true) + await vi.advanceTimersByTimeAsync(1) + expect(isOpen('t1')).toBe(false) + await host.close() + }) + + it('stops the wait when a new handle opens', async () => { + const { host, isOpen, open } = setup({ idleCloseMs: 1000 }) + const first = await open('t1') + vi.useFakeTimers() + await first.close() + await vi.advanceTimersByTimeAsync(500) + + const second = await open('t1') + await vi.advanceTimersByTimeAsync(2000) + expect(isOpen('t1')).toBe(true) + + await second.close() + await vi.advanceTimersByTimeAsync(1000) + expect(isOpen('t1')).toBe(false) + await host.close() + }) +}) + +describe('a dropped handle', () => { + it('is released when the garbage collector drops it, with one warning per session', async () => { + const warnings = droppedWarnings() + const { host, isOpen, open } = setup({ idleCloseMs: 0 }) + const first = await open('t1') + const second = await open('t1') + + finalizeHandle(first) + expect(isOpen('t1')).toBe(true) + finalizeHandle(second) + + expect(isOpen('t1')).toBe(false) + expect(warnings()).toHaveLength(1) + expect(String(warnings()[0]?.[0])).toContain('"t1"') + await host.close() + }) + + it('gives no warning and no second release after close()', async () => { + const warnings = droppedWarnings() + const { host, isOpen, open } = setup({ idleCloseMs: 0 }) + const closed = await open('t1') + const held = await open('t1') + + await closed.close() + finalizeHandle(closed) + + expect(isOpen('t1')).toBe(true) + expect(warnings()).toHaveLength(0) + await held.close() + await host.close() + }) +}) + +describe('a call after the session closed', () => { + it('opens the session again, and the turn runs', async () => { + const { host, calls, isOpen, open } = setup({ idleCloseMs: 0 }) + const session = await open('t1') + await session.close() + expect(isOpen('t1')).toBe(false) + + const turn = session.send('Hello.') + // The operation exists once the session opened again. + expect(() => turn.id).toThrow('Await operation.receipt') + const receipt = await turn.receipt + + expect(receipt.status).toBe('accepted') + expect(turn.id).toMatch(/^op-chat-/) + expect((await turn).text).toBe('Hello back.') + expect(calls).toHaveLength(1) + const transcript = await session.transcript() + expect(transcript.map((message) => message.content)).toEqual([ + 'Hello.', + 'Hello back.', + ]) + // The released handle held the session only for its calls. + await vi.waitFor(() => expect(isOpen('t1')).toBe(false)) + await host.close() + }) +}) + +describe('maxOpen', () => { + it('closes the least recently used idle session that no handle holds', async () => { + const { host, isOpen, open } = setup({ maxOpen: 2 }) + await (await open('a')).close() + await (await open('b')).close() + // `a` is used again, so `b` is the least recently used. + await (await open('a')).close() + + const c = await open('c') + + expect(isOpen('a')).toBe(true) + expect(isOpen('b')).toBe(false) + expect(isOpen('c')).toBe(true) + await c.close() + await host.close() + }) + + it('opens past the limit with one warning when every session is held', async () => { + const warnings = fullWarnings() + const { host, isOpen, open } = setup({ maxOpen: 1 }) + const held = [await open('a')] + + held.push(await open('b'), await open('c')) + + expect(['a', 'b', 'c'].map(isOpen)).toEqual([true, true, true]) + expect(warnings()).toHaveLength(1) + await Promise.all(held.map((session) => session.close())) + await host.close() + }) +}) + +describe('host.close', () => { + it('refuses later opens and the calls of open handles', async () => { + const { host, harness, open } = setup() + const session = await open('t1') + + await host.close() + + await expect(host.open(harness, { threadId: 't2' })).rejects.toThrow( + 'The host is closed.', + ) + expect(() => session.send('Hello.')).toThrow('The host is closed.') + await expect(session.transcript()).rejects.toThrow('The host is closed.') + }) +}) + +describe('the HTTP handler', () => { + it('releases the handles of 10 000 reads, so no session stays open', async () => { + const { host, harness } = setup({ idleCloseMs: 0 }) + const handler = createHarnessHandler({ host, harness, openAccess: true }) + const statuses = new Set() + + for (let index = 0; index < 10_000; index += 1) { + const response = await handler( + new Request(`${BASE}/snapshot?threadId=t-${index}`), + ) + statuses.add(response.status) + } + + expect([...statuses]).toEqual([200]) + await vi.waitFor(() => + expect(hostCountsOf(host)).toEqual({ + sessions: 0, + handles: 0, + holds: 0, + statuses: 0, + closing: 0, + }), + ) + await host.close() + }, 60_000) + + it('releases the handle of an event stream when the client leaves', async () => { + const { host, harness, isOpen } = setup({ idleCloseMs: 0 }) + const handler = createHarnessHandler({ host, harness, openAccess: true }) + const response = await handler(new Request(`${BASE}/events?threadId=t1`)) + expect(hostCountsOf(host).handles).toBe(1) + + await response.body?.cancel() + + await vi.waitFor(() => expect(isOpen('t1')).toBe(false)) + expect(hostCountsOf(host).handles).toBe(0) + await host.close() + }) + + it('releases the handle of a run stream when the turn ends', async () => { + const { host, harness, isOpen } = setup({ idleCloseMs: 0 }) + const handler = createHarnessHandler({ host, harness, openAccess: true }) + const response = await handler( + new Request(`${BASE}/run`, { + method: 'POST', + headers: { 'content-type': 'application/json' }, + body: JSON.stringify({ + threadId: 't1', + runId: 'run-1', + messages: [{ id: 'u1', role: 'user', content: 'Hello.' }], + tools: [], + context: [], + }), + }), + ) + + expect(await response.text()).toContain('Hello back.') + await vi.waitFor(() => expect(isOpen('t1')).toBe(false)) + await host.close() + }) +}) diff --git a/packages/ai-harness/tests/harness-text.test.ts b/packages/ai-harness/tests/harness-text.test.ts index ced5e8cf5d..07daaa724c 100644 --- a/packages/ai-harness/tests/harness-text.test.ts +++ b/packages/ai-harness/tests/harness-text.test.ts @@ -70,7 +70,11 @@ describe('harnessText', () => { adapter: inner.adapter, plugins: () => [watch], }) - const host = createHarnessHost({ persistence: memoryPersistence() }) + // A session that no handle holds closes at once when it is idle. + const host = createHarnessHost({ + persistence: memoryPersistence(), + lifecycle: { idleCloseMs: 0 }, + }) const adapter = harnessText(studio, { host, principal: { id: 'u-1' }, @@ -99,9 +103,13 @@ describe('harnessText', () => { it('leaves a thread open that the caller opened', async () => { const inner = mockAdapter([() => text('kept')]) const studio = defineHarness({ name: 'test/inner', adapter: inner.adapter }) - const host = createHarnessHost({ persistence: memoryPersistence() }) + // Only a handle keeps an idle session open. + const host = createHarnessHost({ + persistence: memoryPersistence(), + lifecycle: { idleCloseMs: 0 }, + }) const threadId = 'outer-kept:test/inner' - await host.open(studio, { threadId }) + const held = await host.open(studio, { threadId }) for await (const _chunk of chat({ adapter: harnessText(studio, { host }), @@ -111,6 +119,8 @@ describe('harnessText', () => { // drain } expect(host.isOpen(studio, { threadId })).toBe(true) + await held.close() + expect(host.isOpen(studio, { threadId })).toBe(false) await host.close() }) diff --git a/packages/ai-harness/tests/oauth-edges.test.ts b/packages/ai-harness/tests/oauth-edges.test.ts index 427787762c..cd14c74ce2 100644 --- a/packages/ai-harness/tests/oauth-edges.test.ts +++ b/packages/ai-harness/tests/oauth-edges.test.ts @@ -419,11 +419,10 @@ describe('oauthConnector edges', () => { const asked = nextSignInUrl(session) const connecting = session.command({ name: 'connect:svc' }) const redirect = (await asked).searchParams.get('redirect_uri') ?? '' - await session.close() + await host.close() expect(connecting.status()).toBe('cancelled') // The sign-in stopped too: nothing listens for the redirect now. await expect(fetch(redirect)).rejects.toThrow() - await host.close() }) }) diff --git a/packages/ai-harness/tests/reload.test.ts b/packages/ai-harness/tests/reload.test.ts index 5d6098bc9c..9e0ed15417 100644 --- a/packages/ai-harness/tests/reload.test.ts +++ b/packages/ai-harness/tests/reload.test.ts @@ -172,7 +172,7 @@ describe('session.reload', () => { const reload = session.reload() // The reload waits for the running turn now. await new Promise((resolve) => setTimeout(resolve, 10)) - await session.close() + await host.close() await reload await running.then( () => undefined, diff --git a/packages/ai-harness/tests/retention.test.ts b/packages/ai-harness/tests/retention.test.ts index 9792913706..334c9de14d 100644 --- a/packages/ai-harness/tests/retention.test.ts +++ b/packages/ai-harness/tests/retention.test.ts @@ -109,9 +109,12 @@ describe('retention', () => { // Every event is longer than 1 byte, so only the newest one stays. const { host, session } = await openWith({ eventBytes: 1 }) await session.send('Hi') - await session.close() + // The read gets the kept events at once, then ends at the abort. + const reading = new AbortController() + setTimeout(() => reading.abort(), 0) - expect(await collect(session.events({ from: '0' }))).toHaveLength(1) + const events = session.events({ from: '0', signal: reading.signal }) + expect(await collect(events)).toHaveLength(1) await host.close() }) diff --git a/packages/ai-harness/tests/session-plugins.test.ts b/packages/ai-harness/tests/session-plugins.test.ts index f4c06b7929..61b3f00a75 100644 --- a/packages/ai-harness/tests/session-plugins.test.ts +++ b/packages/ai-harness/tests/session-plugins.test.ts @@ -289,9 +289,11 @@ describe('questions', () => { await vi.waitFor(() => expect(session.snapshot().pendingQuestions).toHaveLength(1), ) - await session.close() - await expect(open).rejects.toThrow('Session closed.') await host.close() + // The close cancels the command that waits for its answer. + await expect(open).rejects.toBeDefined() + expect(open.status()).toBe('cancelled') + expect(session.snapshot().pendingQuestions).toEqual([]) }) it('rejects the questions of a turn when the turn is cancelled', async () => { diff --git a/packages/ai-harness/tests/session.test.ts b/packages/ai-harness/tests/session.test.ts index f79e462792..13ea3ce2cf 100644 --- a/packages/ai-harness/tests/session.test.ts +++ b/packages/ai-harness/tests/session.test.ts @@ -33,9 +33,10 @@ import type { Reply } from './helpers' /** The key of the thread `t1` in the stores, in the default tenant. */ const T1_KEY = logIdOf({ logId: 't1' }) +/** A memory host. A session closes as soon as its last handle closes and it is idle. */ function setup() { const persistence = memoryPersistence() - const host = createHarnessHost({ persistence }) + const host = createHarnessHost({ persistence, lifecycle: { idleCloseMs: 0 } }) return { persistence, host } } @@ -560,10 +561,9 @@ describe('inbox', () => { await host.close() }) - it('rejects inputs that come while the session closes', async () => { + it('refuses inputs that come while the host closes, and stops the running turn', async () => { const { host } = setup() - const release = gate() - const { adapter } = mockAdapter([untilAborted(), after(release.opened, '')]) + const { adapter } = mockAdapter([untilAborted()]) const session = await host.open( defineHarness({ name: 'test/closing', adapter }), { @@ -573,22 +573,13 @@ describe('inbox', () => { const running = session.send('First.') await running.receipt - const closed = session.close() - const late = session.send('Late.') - expect(await late.receipt).toMatchObject({ - status: 'rejected', - reason: 'closed', - }) - await expect(late).rejects.toThrow() - expect( - await session.send('Later.', { mode: 'followUp' }).receipt, - ).toMatchObject({ - status: 'rejected', - reason: 'closed', - }) + const closed = host.close() + expect(() => session.send('Late.')).toThrow('The host is closed.') + expect(() => session.send('Later.', { mode: 'followUp' })).toThrow( + 'The host is closed.', + ) await expect(running).rejects.toThrow() await closed - release.open() }) it('warns when the inbox cannot mark an input rejected', async () => { diff --git a/packages/ai-harness/tests/sessions-index.test.ts b/packages/ai-harness/tests/sessions-index.test.ts index 8eefa44544..ec2698a7fa 100644 --- a/packages/ai-harness/tests/sessions-index.test.ts +++ b/packages/ai-harness/tests/sessions-index.test.ts @@ -182,17 +182,18 @@ describe('host.sessions', () => { }) it('keeps the first owner when another user opens the thread again', async () => { - const { host, harness } = hostWith() + const { persistence, host, harness } = hostWith() await host.open(harness, { threadId: 't1', principal: { id: 'alice' } }) await host.close() + const next = createHarnessHost({ persistence }) - await host.open(harness, { threadId: 't1', principal: { id: 'bob' } }) + await next.open(harness, { threadId: 't1', principal: { id: 'bob' } }) - expect(await host.sessions.get({ threadId: 't1' })).toMatchObject({ + expect(await next.sessions.get({ threadId: 't1' })).toMatchObject({ ownerId: 'alice', members: [{ userId: 'alice', role: 'owner' }], }) - await host.close() + await next.close() }) it('lists newest first, and only top-level sessions by default', async () => { @@ -408,7 +409,7 @@ describe('host.sessions', () => { it('deletes the entry for good, so the next open writes a new entry', async () => { vi.useFakeTimers({ toFake: ['Date'] }) - const { host, harness } = hostWith() + const { persistence, host, harness } = hostWith() vi.setSystemTime(1000) await host.open(harness, { threadId: 't1', principal: { id: 'u1' } }) await host.sessions.rename({ threadId: 't1', title: 'Old title' }) @@ -417,16 +418,17 @@ describe('host.sessions', () => { expect(await host.sessions.get({ threadId: 't1' })).toBeUndefined() await host.close() vi.setSystemTime(2000) - await host.open(harness, { threadId: 't1' }) + const next = createHarnessHost({ persistence }) + await next.open(harness, { threadId: 't1' }) - expect(await host.sessions.get({ threadId: 't1' })).toEqual({ + expect(await next.sessions.get({ threadId: 't1' })).toEqual({ threadId: 't1', members: [], harness: 'test/sessions', createdAt: 2000, updatedAt: 2000, }) - await host.close() + await next.close() }) it('refuses to open a child thread: its root thread reads it', async () => { diff --git a/packages/ai-harness/tests/shared-log.test.ts b/packages/ai-harness/tests/shared-log.test.ts index 4d25b3c0f5..7f11bfda23 100644 --- a/packages/ai-harness/tests/shared-log.test.ts +++ b/packages/ai-harness/tests/shared-log.test.ts @@ -255,24 +255,22 @@ describe('sessions that share one log', () => { await expect(root.append([{ type: 'app.signal' }])).rejects.toThrow( 'Log conflict', ) - await expect(child.append([{ type: 'app.signal' }])).rejects.toThrow( - 'Log conflict', - ) + const isOpen = (threadId: string) => + host.isOpen(harnessWith([]).harness, { threadId }) + await vi.waitFor(() => { + expect(isOpen(LOG)).toBe(false) + expect(isOpen('task-1')).toBe(false) + }) - // The host replaces a stopped session on the next open. - await vi.waitFor(async () => - expect( - await host.open(harnessWith([]).harness, { threadId: LOG }), - ).not.toBe(root), - ) - await vi.waitFor(async () => - expect( - await host.open(harnessWith([]).harness, { - threadId: 'task-1', - logId: LOG, - }), - ).not.toBe(child), - ) + // Each handle opens a new session on its next call. It folds the log + // again, so it appends after the record of the other writer. + await child.append([{ type: 'app.child' }]) + await root.append([{ type: 'app.root' }]) + const records = await persistence.stores.log.read({ logId: STORE_LOG }) + const types = records + .map((entry) => entry.record.type) + .filter((type) => type.startsWith('app.')) + expect(types).toEqual(['app.signal', 'app.other', 'app.child', 'app.root']) await host.close().catch(() => {}) }) diff --git a/packages/ai-mcp/src/harness.ts b/packages/ai-mcp/src/harness.ts index 9331ee69a6..224f9ac500 100644 --- a/packages/ai-mcp/src/harness.ts +++ b/packages/ai-mcp/src/harness.ts @@ -277,7 +277,8 @@ export async function createHarnessMcpServer

( /** * The execute of a tool that works on the session of its thread. The * session opens for the user of the call, so the host checks each action. - * A refused action becomes a tool error that says so. + * A refused action becomes a tool error that says so. The call releases + * its handle at the end, and the session closes on idle. */ const onSession = ( @@ -290,7 +291,11 @@ export async function createHarnessMcpServer

( async (args: unknown, ctx: ToolExecutionContext) => { try { const session = await open(threadOf(args, defaultThread), ctx.context) - return await work(session, args, ctx.context) + try { + return await work(session, args, ctx.context) + } finally { + await session.close() + } } catch (error) { if (error instanceof AccessDeniedError) throw deniedError(error) throw error @@ -597,6 +602,8 @@ export async function createHarnessMcpServer

( const commands = (session?.describe().commands ?? []).filter(({ name }) => exposedCommands.includes(name), ) + // The server read the agents and the commands, so it holds no session. + await session?.close() const commandTools = toolNames('command', commands).map( ({ item: command, toolName, description }) => toolDefinition({ @@ -695,12 +702,16 @@ export async function createHarnessMcpServer

( }).read(async (_uri, { threadId, id }, ctx) => { try { const target = await open(threadId, ctx.context) - const record = await target.getMedia(id) - if (record === null) { - throw new Error(`Media ${id} was not found in thread ${threadId}.`) + try { + const record = await target.getMedia(id) + if (record === null) { + throw new Error(`Media ${id} was not found in thread ${threadId}.`) + } + const blob = toBase64(await target.loadMedia({ id })) + return { blob, mimeType: record.mimeType } + } finally { + await target.close() } - const blob = toBase64(await target.loadMedia({ id })) - return { blob, mimeType: record.mimeType } } catch (error) { if (error instanceof AccessDeniedError) throw deniedError(error) throw error diff --git a/packages/ai-mcp/tests/harness.test.ts b/packages/ai-mcp/tests/harness.test.ts index d1351c6e83..57f14238c1 100644 --- a/packages/ai-mcp/tests/harness.test.ts +++ b/packages/ai-mcp/tests/harness.test.ts @@ -522,6 +522,33 @@ async function connect( } describe('createHarnessMcpServer', () => { + it('holds no session after the start or after a tool call', async () => { + // A session that no handle holds closes at once when it is idle. + const host = createHarnessHost({ + persistence: memoryPersistence(), + lifecycle: { idleCloseMs: 0 }, + }) + cleanups.push(() => host.close()) + const harness = defineHarness({ + name: 'test/mcp-handles', + adapter: scripted([]).adapter, + }) + const isOpen = () => host.isOpen(harness, { threadId: 'main' }) + const server = await createHarnessMcpServer({ + host, + harness, + openAccess: true, + }) + // The server read the agents and the commands of the session. + expect(isOpen()).toBe(false) + const { client } = await connect(server) + + const status = await client.callTool({ name: 'status', arguments: {} }) + + expect(status.isError).not.toBe(true) + expect(isOpen()).toBe(false) + }) + it('lists the control tools, one tool per exposed agent, and one per exposed command', async () => { const { server } = await harnessServer([]) const { client } = await connect(server) From 5c1b958aee043b73503342b0d1c2d252bf3dd785 Mon Sep 17 00:00:00 2001 From: Alem Tuzlak Date: Sat, 17 Oct 2026 13:52:29 +0200 Subject: [PATCH 05/12] docs(ai-harness): lifecycle and retention, handles, observe, signals, and OAuth cancel --- .changeset/ai-harness-cli.md | 1 + .changeset/ai-harness-core.md | 8 +- .changeset/ai-harness-transport.md | 4 +- .changeset/ai-mcp-harness.md | 6 +- .changeset/harness-acp-agent.md | 2 +- docs/config.json | 9 +- docs/harness/auth.md | 71 +++++++++++ docs/harness/cli.md | 37 +++++- docs/harness/connect.md | 4 +- docs/harness/custom-ui.md | 2 + docs/harness/durable-sessions.md | 6 +- docs/harness/lifecycle.md | 185 +++++++++++++++++++++++++++++ docs/harness/mcp.md | 4 +- docs/harness/overview.md | 6 +- docs/harness/plugins.md | 41 ++++++- docs/harness/sessions.md | 10 +- 16 files changed, 369 insertions(+), 27 deletions(-) create mode 100644 docs/harness/lifecycle.md diff --git a/.changeset/ai-harness-cli.md b/.changeset/ai-harness-cli.md index 8569cf66b8..6484a6b25d 100644 --- a/.changeset/ai-harness-cli.md +++ b/.changeset/ai-harness-cli.md @@ -8,4 +8,5 @@ Add `@tanstack/ai-harness-cli`: `runCli(harness)` runs a harness from the termin - `-p "prompt"` runs one prompt for scripts and CI, and `--output ndjson` prints every AG-UI event. - `--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. +- 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. diff --git a/.changeset/ai-harness-core.md b/.changeset/ai-harness-core.md index 7eaa5b371c..b4ca313aab 100644 --- a/.changeset/ai-harness-core.md +++ b/.changeset/ai-harness-core.md @@ -5,12 +5,16 @@ Add `@tanstack/ai-harness`: build an agent harness on TanStack AI. A harness keeps one conversation open across many turns. - `defineHarness()` names an agent: its adapter, tools, system prompts, middleware, agents, subagents, routing, plugins, turn hooks, and durability options. -- `createHarnessHost()` opens sessions. 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()` 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. +- `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 claims nothing. +- `host.sessions.delete` closes the sessions of a thread and of its child threads. Then it deletes their logs or messages, index entries, uploaded media, generated artifacts, settings, config, interrupted turn, usage, and revert data. +- 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. - 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. -- `definePlugin()` adds tools, prompts, middleware, agents, commands, config options, state, and extension points to a session or to one turn. +- `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()`. - 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) and `agents.start()` (a child that counts against the subagent budget). - With `stores.log` from `@tanstack/ai-persistence`, a session is durable: inputs, tool results, and steps are in the log, and the next host recovers a turn that a crash stopped. `durableTool()` runs side effects once. `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. diff --git a/.changeset/ai-harness-transport.md b/.changeset/ai-harness-transport.md index 93f73d6597..ee46421619 100644 --- a/.changeset/ai-harness-transport.md +++ b/.changeset/ai-harness-transport.md @@ -4,8 +4,8 @@ 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. +- `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. - `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. -- `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. +- `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. diff --git a/.changeset/ai-mcp-harness.md b/.changeset/ai-mcp-harness.md index fa23fd493d..a475b51da5 100644 --- a/.changeset/ai-mcp-harness.md +++ b/.changeset/ai-mcp-harness.md @@ -5,9 +5,9 @@ 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 kind of interrupt in one call. -- 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. -- `@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. `oauth: true` signs in to a server with `/connect `. -- `@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). +- 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. +- `@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. +- `@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. - `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. `@tanstack/ai-harness` is an optional peer dependency. Only the `./harness` and `./connector` entries need it. diff --git a/.changeset/harness-acp-agent.md b/.changeset/harness-acp-agent.md index 5797f4917b..69065bd902 100644 --- a/.changeset/harness-acp-agent.md +++ b/.changeset/harness-acp-agent.md @@ -2,4 +2,4 @@ '@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). Tool approvals become permission requests, 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`. `@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). Tool approvals become permission requests, 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. diff --git a/docs/config.json b/docs/config.json index 6bfa335aee..bea2790407 100644 --- a/docs/config.json +++ b/docs/config.json @@ -973,7 +973,7 @@ "label": "Durable sessions", "to": "harness/durable-sessions", "addedAt": "2026-10-13", - "updatedAt": "2026-10-15" + "updatedAt": "2026-10-17" }, { "label": "List, rename, and fork sessions", @@ -981,6 +981,11 @@ "addedAt": "2026-10-13", "updatedAt": "2026-10-17" }, + { + "label": "Lifecycle and retention", + "to": "harness/lifecycle", + "addedAt": "2026-10-17" + }, { "label": "Run side effects once", "to": "harness/durable-tools", @@ -1122,7 +1127,7 @@ "label": "Use MCP servers", "to": "harness/mcp", "addedAt": "2026-10-13", - "updatedAt": "2026-10-16" + "updatedAt": "2026-10-17" }, { "label": "Use from any MCP client", diff --git a/docs/harness/auth.md b/docs/harness/auth.md index 7efe944084..721208f7a9 100644 --- a/docs/harness/auth.md +++ b/docs/harness/auth.md @@ -58,7 +58,9 @@ export const assistant = defineHarness({ - The user runs `/connect github`. - With `login: 'loopback'`, the CLI opens the browser. The harness listens on `127.0.0.1` on a random port for one callback, with PKCE and a random `state`. +- With `oauth.issuer` set, the sign-in refuses a code that the service sends with another `iss` (RFC 9207). - The token goes into the credential store. `/disconnect github` deletes it. +- A cancel of the command, or a close of the session, stops a sign-in that waits. The listener on `127.0.0.1` closes at once. `login` is required. Use `login: 'device'` for a harness on a server, and for SSH sessions, containers, and CI. The user then gets a code to enter on the service's page. Device login needs `oauth.deviceUrl`. @@ -198,8 +200,77 @@ Your plugin reads the credentials of the user of each run: See [Act as the user of a run](./plugins#act-as-the-user-of-a-run). +## Write your own sign-in command + +`oauthConnector` covers most services. For a sign-in of your own, use the helpers from `@tanstack/ai-harness/oauth`. Pass the `signal` of the command, so a cancel stops the sign-in at once: + +```ts group=harness-auth-own +import { defineCommand, definePlugin } from '@tanstack/ai-harness' +import { deviceLogin, isExpired, refreshCredential } from '@tanstack/ai-harness/oauth' +import type { OAuthConfig } from '@tanstack/ai-harness/oauth' + +const gitlab: OAuthConfig = { + authorizationUrl: 'https://gitlab.com/oauth/authorize', + tokenUrl: 'https://gitlab.com/oauth/token', + deviceUrl: 'https://gitlab.com/oauth/authorize_device', + clientId: 'your-gitlab-app-client-id', + scopes: ['read_api'], +} + +export const gitlabLogin = definePlugin({ + name: 'acme/gitlab-login', + setup: (ctx) => ({ + commands: { + 'connect:gitlab': defineCommand({ + description: 'Sign in to GitLab', + run: async (_input, { credentials, signal }) => { + const credential = await deviceLogin(gitlab, { + onCode: ({ userCode, verificationUri }) => + ctx.session.authRequired({ connector: 'gitlab', url: verificationUri, userCode }), + signal, + }) + await credentials.set({ id: 'gitlab', credential }) + return 'Connected to GitLab.' + }, + }), + }, + }), +}) +``` + +A tool refreshes an expired token with the `signal` of its run, so the refresh stops with the run: + +```ts group=harness-auth-own +import { toolDefinition } from '@tanstack/ai' +import type { HarnessRunContext } from '@tanstack/ai-harness' + +export const gitlabProjects = toolDefinition({ + name: 'gitlab_projects', + description: 'List my GitLab projects', +}).server(async (_args, ctx) => { + const { credentials, signal } = ctx.context.run + let credential = await credentials.require('gitlab') + if (isExpired(credential)) { + credential = await refreshCredential(gitlab, credential, { signal }) + await credentials.set({ id: 'gitlab', credential }) + } + const token = credential.type === 'api_key' ? credential.value : credential.accessToken + const response = await fetch('https://gitlab.com/api/v4/projects?membership=true', { + headers: { authorization: `Bearer ${token}` }, + signal, + }) + return response.json() +}) +``` + +- `deviceLogin`, `loopbackLogin`, `exchangeCode`, and `refreshCredential` take a `signal`. When it aborts, the call stops at once and rejects with the reason of the signal. +- The `signal` of a command aborts on a cancel of the command and on a close of the session. +- `loopbackLogin` takes `timeoutMs`. The default is 10 minutes. This timer does not keep the process alive. +- When another library runs the OAuth flow, `startLoopbackReceiver({ timeoutMs, signal })` only waits for the callback. `waitForCode(state)` gives the code and the `iss` of the callback. + ## What you have now - Browser and device-code sign-in for any OAuth service. - Tools that get a fresh token without the model seeing it. - Turns that wait for a sign-in and then finish their work. +- Sign-ins of your own that stop at once on a cancel. diff --git a/docs/harness/cli.md b/docs/harness/cli.md index 3e89a48f77..7ebbcbe252 100644 --- a/docs/harness/cli.md +++ b/docs/harness/cli.md @@ -70,7 +70,39 @@ Line mode reads one message or command per line and waits for each turn. It work | 0 | The turn finished. | | 1 | The turn failed. | | 2 | The turn waits for an approval. | -| 130 | The turn was cancelled. | +| 130 | The turn was cancelled, or SIGINT (Ctrl+C) stopped the CLI. | +| 143 | SIGTERM cancelled the turn or stopped the CLI. | + +The codes 130 and 143 apply in every mode. See [Stop with Ctrl+C](#stop-with-ctrlc). + +## Stop with Ctrl+C + +Ctrl+C works the same in every mode: + +- The first Ctrl+C (SIGINT) or SIGTERM cancels the running turn. +- A second one, or one while no turn runs, stops the CLI. `runCli` resolves to 130 for SIGINT, and to 143 for SIGTERM. +- When the CLI made the host, it closes the host with `recoverable: true`. A durable host then runs the cut work again. +- When you gave the CLI your own `host`, the CLI releases only its session handle. Your host stays open. + +The CLI hears the signals on `process`. To send them yourself, for example in a test, pass a `signals` source. An `EventEmitter` works: + +```ts group=harness-cli-signals +import { EventEmitter } from 'node:events' +import { defineHarness } from '@tanstack/ai-harness' +import { runCli } from '@tanstack/ai-harness-cli' +import { openaiText } from '@tanstack/ai-openai' + +const assistant = defineHarness({ + name: 'acme/assistant', + adapter: openaiText('gpt-6.1-sol'), +}) + +const signals = new EventEmitter() +const exit = runCli(assistant, { argv: ['-p', 'Write a long story.'], signals }) +// No turn runs yet, so this stops the CLI. +signals.emit('SIGTERM') +console.log(await exit) // 143 +``` ## Serve over HTTP safely @@ -108,7 +140,7 @@ await host.close() ``` - You own the host, so the CLI does not close it. Close it yourself when `runCli` returns. Without `host`, the CLI builds a host and closes it. -- Give `host` or `persistence`, not both. With both, `runCli` writes `Give runCli a host or persistence, not both.` and returns code 1. +- 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. @@ -185,6 +217,7 @@ For an Ink screen with approvals, questions, sign-ins, and child agents, copy [` - One entry file that runs your harness as a terminal app, a script step, an editor agent, an MCP server, or an HTTP server. - A CLI that uses your durable host and runs as your user. +- A Ctrl+C that cancels the turn first, then stops with exit code 130. - Your own terminal screen on the same session, with any TUI library. - Files that you send with `@path`, and a folder with the media that your agents make. diff --git a/docs/harness/connect.md b/docs/harness/connect.md index 9701c97d3a..bfd4ff7262 100644 --- a/docs/harness/connect.md +++ b/docs/harness/connect.md @@ -72,7 +72,9 @@ For a list of sessions: - `POST sessions`: change a session with `{ op, threadId }`. The ops are `rename`, `update`, `delete`, `fork`, `share`, and `unshare`. See [List, rename, and fork sessions](./sessions) and [Share the thread](./shared-threads#2-share-the-thread). - `GET host-events`: the status changes and index changes of the sessions that this user can read, as SSE. -A read of a thread that has no session answers `404`, and it starts no session. Only `GET run?threadId=` answers an empty thread, because `useChat` reads it before the first message. +A read of a thread that does not exist answers `404`. Only `GET run?threadId=` answers an empty thread, because `useChat` reads it before the first message. + +The read paths read the thread without a session, so they set up no plugins and claim nothing. `GET describe` opens the session, because the plugins give the commands and the settings. For a read route of your own, use `host.observe`. See [Read a thread without a session](./lifecycle#read-a-thread-without-a-session). ## Cursors diff --git a/docs/harness/custom-ui.md b/docs/harness/custom-ui.md index fb49b21a3d..cf333c9536 100644 --- a/docs/harness/custom-ui.md +++ b/docs/harness/custom-ui.md @@ -487,6 +487,8 @@ export const handler = createHarnessHandler({ host, harness: assistant }) The handler also serves `GET transcript` and `GET describe`. The view reads them when it starts. For the other routes, see [Connect clients](./connect). +`GET transcript`, `GET snapshot`, `GET events`, and the other read paths read the thread without a session, so they start no plugins on the server. `GET describe` opens the session, because the plugins give the commands and the settings. To read a thread in a route of your own, use `host.observe`. See [Read a thread without a session](./lifecycle#read-a-thread-without-a-session). + ### In the browser Make the view one time, in code that runs in the browser. Then read it in your components: diff --git a/docs/harness/durable-sessions.md b/docs/harness/durable-sessions.md index 805b68afeb..18094dbffc 100644 --- a/docs/harness/durable-sessions.md +++ b/docs/harness/durable-sessions.md @@ -77,7 +77,7 @@ A background agent that you started with `resume: true` runs again on the next h ## Stop a host for a deploy -`host.close()` ends each running turn `aborted`, so no host runs it again. When the server stops only for a deploy, close it with `recoverable: true`. The running turns stop, but they stay open in the log, and the next host continues them at once: +`host.close()` closes every session of the host and refuses later opens. It ends each running turn `aborted`, so no host runs it again. When the server stops only for a deploy, close it with `recoverable: true`. The running turns stop, but they stay open in the log, and the next host continues them at once: ```ts group=harness-durable process.on('SIGTERM', async () => { @@ -89,7 +89,7 @@ process.on('SIGTERM', async () => { - Running turns and background agents stop. The host writes no end for them and gives their leases back, so the next host does not wait for the lease to expire. - A tool that the stop cut gets the same result as after a crash. See [Mark tools that are safe to run again](#mark-tools-that-are-safe-to-run-again). - A turn that a user cancelled still ends `aborted`. -- `session.close({ recoverable: true })` does the same for one session. +- `session.close()` releases only your handle. It does not stop the work of the session. See [Lifecycle and retention](./lifecycle#open-and-release-a-session). During a rolling deploy, old and new hosts run at the same time. A host that reads a log record type from a newer version does not write to that log. Its session stops with `LogVersionError`, and the new hosts continue the thread. @@ -165,7 +165,7 @@ It resolves with the threads that it opened, as `{ threadId, harness }` entries. 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. -- `'never'`: the session stays open, like any session that you open. +- `'never'`: the session closes like any session that no handle holds, after `lifecycle.idleCloseMs`. See [When a session closes](./lifecycle#when-a-session-closes). ### Watch for claim warnings diff --git a/docs/harness/lifecycle.md b/docs/harness/lifecycle.md new file mode 100644 index 0000000000..7b60a0f609 --- /dev/null +++ b/docs/harness/lifecycle.md @@ -0,0 +1,185 @@ +--- +title: Lifecycle and retention +id: harness-lifecycle +order: 4 +description: "Release session handles with await using, let idle sessions close, cap the open sessions, read a thread without a session, bound what each session keeps in memory, and delete a thread for good." +keywords: + - tanstack ai + - harness + - lifecycle + - retention + - await using + - idle close + - delete +--- + +Your server opens a session for each request. Each open session holds its plugins, its MCP processes, and its recent events in memory. A server that never lets them go grows until it runs out of memory. The host closes the sessions that nobody uses, caps how many stay open, and bounds what each session keeps. + +## Open and release a session + +Each `host.open` gives you a handle: your own reference to the live session of the thread. When you are done with the handle, release it. With `await using`, the release runs at the end of the block: + +```ts group=harness-lifecycle +import { createHarnessHost, defineHarness } from '@tanstack/ai-harness' +import { memoryPersistence } from '@tanstack/ai-persistence' +import { openaiText } from '@tanstack/ai-openai' + +export const assistant = defineHarness({ + name: 'acme/assistant', + adapter: openaiText('gpt-6.1-sol'), +}) + +const host = createHarnessHost({ persistence: memoryPersistence() }) + +export async function reply(threadId: string, message: string) { + await using session = await host.open(assistant, { threadId }) + const turn = await session.send(message) + return turn.text +} +``` + +- The handles of one thread share one live session. A release lets go of this handle only. +- `await session.close()` also releases the handle. A second call does nothing. +- A release does not stop work. A turn that runs goes on until it ends. +- If you drop a handle with no release, the garbage collector releases it, and the host logs one warning. +- A call on a handle after its session closed opens the session again. On a host with [`access`](./access), `authorize` checks that call again. + +## When a session closes + +The host closes a session when no handle holds it and it was idle for `lifecycle.idleCloseMs`. Set the limits on the host: + +```ts group=harness-lifecycle +const busyHost = createHarnessHost({ + persistence: memoryPersistence(), + lifecycle: { idleCloseMs: 30_000, maxOpen: 200 }, +}) +``` + +- `idleCloseMs`: how long an idle session with no handle stays open, in milliseconds. The default is 60 000. `0` closes it at once. +- `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. +- 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). + +## Stop the host + +`host.close()` closes every session of the host, for every handle. Then `host.open` throws: + +```ts group=harness-lifecycle +await busyHost.close() +``` + +- Without options, the running work is cancelled. +- With `recoverable: true`, the running work stops, and the next host that opens the thread runs it again. Use it on a durable host that stops for a deploy. See [Stop a host for a deploy](./durable-sessions#stop-a-host-for-a-deploy). +- After the close, a call on a handle throws. `snapshot()` still gives the last state. +- The close also ends the event reads of each [view](#read-a-thread-without-a-session), and `host.observe` throws from then on. + +## Read a thread without a session + +A preview route or a status page only reads a thread. An open sets up every plugin, can start MCP processes, and claims the thread for work. `host.observe` reads the thread and opens no session. + +On the server, read the transcript in your route: + +```ts group=harness-lifecycle-observe +import { createHarnessHost, defineHarness } from '@tanstack/ai-harness' +import { memoryPersistence } from '@tanstack/ai-persistence' +import { openaiText } from '@tanstack/ai-openai' + +const assistant = defineHarness({ + name: 'acme/assistant', + adapter: openaiText('gpt-6.1-sol'), +}) +const host = createHarnessHost({ persistence: memoryPersistence() }) + +export async function GET(request: Request) { + const threadId = new URL(request.url).searchParams.get('threadId') + if (!threadId) return new Response('Missing threadId', { status: 400 }) + const view = await host.observe(assistant, { threadId }) + return Response.json(await view.transcript()) +} +``` + +In the browser, fetch the route: + +```ts group=harness-lifecycle-observe-client +const response = await fetch('/api/preview?threadId=chat-1') +console.log(await response.json()) +``` + +On a host with `access`, pass the `principal` of the request. The thread needs the `read` action, else `observe` throws `AccessDeniedError`. A thread with no index entry gets the same error, so a read creates no thread. + +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. On a host with no log, only an open session on this host has events. +- `agentTranscript(id)`: 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. +- `usage()`: the token totals of the thread. +- `entry()`: the index entry of the thread. +- `logState()`: the fold of the log, on a durable host. + +The view sets up no plugins, writes nothing, and claims nothing. When the session is open on this host, the view reads through it, but it never holds the session open. The read paths of [`createHarnessHandler`](./connect) use the view, so a client that only reads starts no plugins. `GET describe` opens the session, because the plugins give the commands and the settings. + +## Bound what a session keeps in memory + +A thread can run for weeks. Each turn and each agent run that settles leaves its receipt, its result, and its events in memory. `lifecycle.retention` keeps a window of each and drops the oldest: + +```ts group=harness-lifecycle +const leanHost = createHarnessHost({ + persistence: memoryPersistence(), + lifecycle: { + retention: { settledOperations: 200, settledAgentRuns: 100, eventBytes: 2_000_000 }, + }, +}) +``` + +| Limit | Default | Past the limit | +| --- | --- | --- | +| `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. | + +On a durable host (`stores.log`), the log keeps what memory drops: + +- `session.settled(inputId)` of an old input reads the log. +- `agentRun(id)` and `agentRuns()` also find the runs in the log. +- `events({ from })` reads the older events from the log. +- `eventBytes` also bounds the log tail that the host keeps in memory. + +On a host without a log, a message to an agent run that left the window gets `not_running`. + +## Delete a thread for good + +A user deletes a chat, and its data must go with it: + +```ts group=harness-lifecycle +await leanHost.sessions.delete({ threadId: 'chat-1' }) +``` + +First, `delete` closes the sessions of the thread and of its child threads, for every handle. The host finds the children through the `parentThreadId` of their index entries. Then it deletes this data of each thread: + +- The log. On a host without a log, the messages. +- The index entry. +- The uploaded media and the generated artifacts. +- The thread settings, the plugin config, and the interrupted turn. +- The usage totals and the revert data. + +These stay in their stores: plugin state, run records, interrupts, the inbox, and work claims. To remove them, delete them in your own stores. + +- 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. +- `host.events()` yields a `session-deleted` event. See [Show a live status](./sessions#show-a-live-status). + +## What you have now + +- Handles that release at the end of a block, so idle sessions close by themselves. +- A cap on the open sessions of a host. +- Read routes that start no plugins and claim no thread. +- Sessions that keep a bounded window in memory, also after weeks of turns. +- A delete that removes the data of a thread and its children. diff --git a/docs/harness/mcp.md b/docs/harness/mcp.md index 3430546b44..7cae04bc5d 100644 --- a/docs/harness/mcp.md +++ b/docs/harness/mcp.md @@ -63,9 +63,9 @@ 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. -- `stdio`: the server runs as a child process on this machine (Node only). It stops when the session closes. +- `stdio`: the server runs as a child process on this machine (Node only). The host starts it one time, and every session of the host 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`. -- `oauth: true`: the user signs in with `/connect linear`. See [Sign in to an OAuth server](#3-sign-in-to-an-oauth-server). +- `oauth: true`: the user signs in with `/connect linear`. Each user connects with their own sign-in. See [Sign in to an OAuth server](#3-sign-in-to-an-oauth-server). ## 2. Check the servers diff --git a/docs/harness/overview.md b/docs/harness/overview.md index 6592a5da0d..d1d2545c02 100644 --- a/docs/harness/overview.md +++ b/docs/harness/overview.md @@ -102,7 +102,7 @@ const host = createHarnessHost({ persistence: memoryPersistence() }) const session = await host.open(assistant, { threadId: 'thread-1' }) ``` -Opening the same `threadId` again returns the live session. +Opening the same `threadId` again gives another handle to the same live session. Release a handle with `await session.close()`, or open it with `await using`, so an idle session can close. See [Lifecycle and retention](./lifecycle). Does the host serve many users? Give it `access`, so each user reaches only their own sessions. See [Access and tenants](./access). @@ -191,7 +191,7 @@ export const withClock = defineHarness({ }) ``` -`plugins` is a function, so every session gets its own plugin instances. Set `lifetime: 'turn'` on a plugin to set it up again for each turn. +`plugins` is a function, so every session gets its own plugin instances. Set `lifetime: 'turn'` on a plugin to set it up again for each turn, or `lifetime: 'host'` to share one setup between all the sessions of the host. If two plugins add a tool with the same name, `host.open` fails and names both plugins. To write your own, see [Write a harness plugin](./plugins). @@ -230,7 +230,7 @@ const support = await host.open(assistant, { }) ``` -If the `threadId` already has a live session, `host.open` returns that session with its first `promptCache` value. The agents that the model calls as tools use the retention of the turn. `session.usage()` shows the cache reads and writes, see [Track usage and cost](./usage). For the options and the cost, see [Prompt Caching](../advanced/prompt-caching). +If the `threadId` already has a live session, `host.open` gives a handle to that session, which keeps its first `promptCache` value. The agents that the model calls as tools use the retention of the turn. `session.usage()` shows the cache reads and writes, see [Track usage and cost](./usage). For the options and the cost, see [Prompt Caching](../advanced/prompt-caching). To change it for one prompt, pass `overrides.promptCache`. See [Give one prompt its own settings](./turn-control#give-one-prompt-its-own-settings). diff --git a/docs/harness/plugins.md b/docs/harness/plugins.md index 73123d25a0..bc1befac84 100644 --- a/docs/harness/plugins.md +++ b/docs/harness/plugins.md @@ -24,7 +24,7 @@ export const today = definePlugin({ }) ``` -Add it with `plugins: () => [today]` in `defineHarness`. `setup` runs once per session. +Add it with `plugins: () => [today]` in `defineHarness`. `setup` runs once per session. To share one setup between all the sessions of a host, see [Share one setup between sessions](#share-one-setup-between-sessions). ## What `setup` can return @@ -258,6 +258,7 @@ Where each part of a plugin reads the run: - A command: `run` of the command context. Its `credentials` is the same as `run.credentials`. - A plugin with `lifetime: 'turn'`: `ctx.turn.run`. Its `ctx.credentials` and `ctx.keys` belong to the same user. - 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`. The run has these fields: @@ -401,10 +402,46 @@ Read `ctx.collect` at run time (in a command, a tool, or a middleware hook). Dur ## Clean up resources -Open resources with `ctx.resources.acquire(open, close)`. They close when the session closes, newest first. If a later plugin fails in `setup`, every plugin set up so far is cleaned up, and no model request is sent. +Open resources with `ctx.resources.acquire(open, close)`. They close when the plugin is cleaned up, newest first: when the session closes, or after the turn for `lifetime: 'turn'`. If a later plugin fails in `setup`, every plugin set up so far is cleaned up, and no model request is sent. Set `lifetime: 'turn'` to set a plugin up again for each turn. +## Share one setup between sessions + +A plugin that opens a database pool, or starts a process, does that work again for each session. A host with 300 open sessions then holds 300 pools. Set `lifetime: 'host'`, and the host sets the plugin up one time and shares it with every session: + +```ts group=harness-plugins-host +import { toolDefinition } from '@tanstack/ai' +import { definePlugin } from '@tanstack/ai-harness' + +// Your shared client, for example a database pool. +async function connect() { + return { countOrders: async () => 42, end: async () => {} } +} + +export const orders = definePlugin({ + name: 'acme/orders', + lifetime: 'host', + setup: async (ctx) => { + const db = await ctx.resources.acquire(connect, (pool) => pool.end()) + return { + tools: [ + toolDefinition({ name: 'count_orders', description: 'Count the orders' }).server( + async () => db.countOrders(), + ), + ], + } + }, +}) +``` + +- The host sets the plugin up one time for each harness. Every session of that harness shares this setup. +- The host cleans it up when the last session that uses it closes, or at `host.close()`. +- Its setup context has no session, no credentials, no keys, and no commands. +- A tool of the plugin reads the user of its call from `ctx.context.run`. See [Act as the user of a run](#act-as-the-user-of-a-run). + +The [MCP plugin](./mcp) works this way for `stdio` servers: one process serves every session of the host. + ## Reload plugins You changed a plugin, an agent file, or the list that `plugins()` returns, and an open session still runs the old setup. `session.reload()` sets the plugins up again: diff --git a/docs/harness/sessions.md b/docs/harness/sessions.md index 4d4b912b90..27df96ec31 100644 --- a/docs/harness/sessions.md +++ b/docs/harness/sessions.md @@ -43,6 +43,8 @@ await session.send('Plan a trip to Lisbon.') console.log((await host.sessions.get({ threadId: 'chat-1' }))?.ownerId) // 'ada' ``` +`host.open` gives you a handle. When you are done with it, release it with `await session.close()`. Or open it with `await using`, which releases it at the end of the block. Then the session can close when it goes idle. See [Lifecycle and retention](./lifecycle#open-and-release-a-session). + What an entry holds: - `title`: the name of the chat. Set it with `host.sessions.rename`, or from a plugin with `ctx.session.updateEntry({ title })`. @@ -62,7 +64,7 @@ const durableHost = createHarnessHost({ }) ``` -- Without `stores.sessions`, `list` is empty, `get` returns `undefined`, and `rename`, `update`, and `delete` write nothing. +- Without `stores.sessions`, `list` is empty, `get` returns `undefined`, and `rename` and `update` write nothing. - The store needs `getVersioned` and `upsertIf`, else `createHarnessHost` throws. The host changes an entry with compare-and-set, so two hosts never lose each other's changes. - To write your own store, read the [`SessionIndexStore` contract](../persistence/store-reference#sessionindexstore). @@ -160,10 +162,10 @@ await host.sessions.delete({ threadId: 'chat-1' }) ``` - `rename` resolves to the changed entry. For a thread with no entry, it resolves to `undefined`. -- `delete` removes the index entry for good. The transcript, the log, the media, and the settings of the thread stay in their stores. -- After `delete`, the turns of the open session do not write the entry again. If the thread opens again after a restart, the host writes a new entry for it, with the user who opens it as the owner. +- `delete` closes the sessions of the thread and of its child threads, for every handle. Then it deletes their data: the log or the messages, the index entry, the media, the settings, and the usage. +- The next user who opens the same `threadId` starts empty. The host writes a new entry for the thread, with that user as the owner. -`delete` closes the open session and removes the entry and the transcript, so the next user who opens the same `threadId` starts empty. Uploaded media and plugin data of the thread stay in their stores. +For the full list of what `delete` removes, and what stays, see [Delete a thread for good](./lifecycle#delete-a-thread-for-good). ## Keep app data on a session From 4faa68796a1792aa5e8059d6f5d19640a174ddbd Mon Sep 17 00:00:00 2001 From: Alem Tuzlak Date: Sat, 17 Oct 2026 14:04:10 +0200 Subject: [PATCH 06/12] feat(ai-harness, ai-mcp, ai-acp): host.observe, a read-only session view, and read routes on it --- packages/ai-acp/src/agent/index.ts | 40 +- packages/ai-acp/tests/agent-handlers.test.ts | 29 ++ packages/ai-harness/src/host.ts | 112 +++++- packages/ai-harness/src/http.ts | 372 +++++++++++------- packages/ai-harness/src/index.ts | 1 + packages/ai-harness/src/log.ts | 109 ++++- .../src/session/observed-session.ts | 324 +++++++++++++++ packages/ai-harness/src/session/state.ts | 9 +- packages/ai-harness/tests/handles.test.ts | 13 +- packages/ai-harness/tests/observe.test.ts | 340 ++++++++++++++++ packages/ai-mcp/src/harness.ts | 59 ++- 11 files changed, 1215 insertions(+), 193 deletions(-) create mode 100644 packages/ai-harness/src/session/observed-session.ts create mode 100644 packages/ai-harness/tests/observe.test.ts diff --git a/packages/ai-acp/src/agent/index.ts b/packages/ai-acp/src/agent/index.ts index 982800fca8..7c0b830206 100644 --- a/packages/ai-acp/src/agent/index.ts +++ b/packages/ai-acp/src/agent/index.ts @@ -251,7 +251,8 @@ const APPROVAL_OPTIONS = [ * * 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 when it is idle. `session/resume` only checks the session index: + * the session opens at its first prompt. * * ACP v2 is a draft, so this API is experimental too. */ @@ -298,22 +299,30 @@ export function createAcpAgent

( } /** - * The session of an id that this agent or the session index knows. 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. + * Throw unless this agent or the session index knows the id. 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 known = async (sessionId: string) => { + const checkKnown = async (sessionId: string) => { + if (handles.has(sessionId)) return try { - if (!handles.has(sessionId)) { - const entry = await host.sessions.get({ - threadId: sessionId, - principal, - }) - const isOwnSession = - entry?.harness === harness.name && entry.parentThreadId === undefined - if (!isOwnSession) throw RequestError.resourceNotFound(sessionId) + const entry = await host.sessions.get({ threadId: sessionId, principal }) + const isOwnSession = + entry?.harness === harness.name && entry.parentThreadId === undefined + if (!isOwnSession) throw RequestError.resourceNotFound(sessionId) + } catch (error) { + if (error instanceof AccessDeniedError) { + throw RequestError.resourceNotFound(sessionId) } + throw error + } + } + + /** The session of an id that this agent or the session index knows. */ + const known = async (sessionId: string) => { + await checkKnown(sessionId) + try { return await live(sessionId) } catch (error) { if (error instanceof AccessDeniedError) { @@ -532,8 +541,9 @@ export function createAcpAgent

( nextCursor: page.truncated ? page.cursor : null, } }) + // The session opens at its first prompt: a resume mounts nothing. .onRequest('session/resume', async ({ params }) => { - await known(params.sessionId) + await checkKnown(params.sessionId) return {} }) .onRequest('session/prompt', async ({ params, client }) => { diff --git a/packages/ai-acp/tests/agent-handlers.test.ts b/packages/ai-acp/tests/agent-handlers.test.ts index 6755b79fb1..2275da7594 100644 --- a/packages/ai-acp/tests/agent-handlers.test.ts +++ b/packages/ai-acp/tests/agent-handlers.test.ts @@ -798,6 +798,35 @@ describe('ACP agent handlers', () => { }) }) + it('resumes a session without opening it: it opens at its first prompt', async () => { + const { harness, host, editor, idle, textOf } = setup({ + lifecycle: { idleCloseMs: 0 }, + }) + close = () => host.close() + await (await host.open(harness, { threadId: 'kept-session' })).close() + const isOpen = () => host.isOpen(harness, { threadId: 'kept-session' }) + await editor.connectWith(createAcpAgent({ host, harness }), async (ctx) => { + await ctx.request('initialize', { + protocolVersion: 2, + info: { name: 'editor', version: '1.0.0' }, + }) + + await ctx.request('session/resume', { + sessionId: 'kept-session', + cwd: '/tmp', + }) + expect(isOpen()).toBe(false) + + await ctx.request('session/prompt', { + sessionId: 'kept-session', + prompt: [{ type: 'text', text: 'hello' }], + }) + await vi.waitFor(() => expect(idle('kept-session')).toBeDefined()) + expect(textOf('kept-session')).toBe('echo: hello') + expect(isOpen()).toBe(true) + }) + }) + describe('on a host with access', () => { const olga = { id: 'olga', tenantId: 'acme' } const bea = { id: 'bea', tenantId: 'acme' } diff --git a/packages/ai-harness/src/host.ts b/packages/ai-harness/src/host.ts index 109df9152f..2456913f57 100644 --- a/packages/ai-harness/src/host.ts +++ b/packages/ai-harness/src/host.ts @@ -7,6 +7,7 @@ import { SharedLog, loadLogState, logIdOf, parseLogId } from './log' import { LEASE } from './resume' import { SessionImpl } from './session/harness-session' import { sessionHandle } from './session/handle' +import { observedSession } from './session/observed-session' import { HARNESS_EVENTS } from './types' import { isRecord } from './utils' import { changedEntry } from './session-entry' @@ -42,6 +43,7 @@ import type { SessionEntryChange, SessionEntryUpdate } from './session-entry' import type { ReduceOptions } from './log' import type { LeaseOptions } from './resume' import type { HarnessSession } from './session/index' +import type { ObservedSession } from './session/observed-session' import type { ForkPoint, Principal } from './types' type SharedStores = { @@ -509,6 +511,32 @@ export interface HarnessHost

{ harness: THarness, options: OpenSessionOptions

, ) => Promise> + /** + * A read-only view of a thread of `harness`, for `principal`: its events, + * transcript, snapshot, media, usage, and index entry (see + * `ObservedSession`). It mounts no plugins, writes nothing (no index entry, no + * metadata), claims nothing, and runs no recovery, so a read of an idle + * thread costs only store reads. When the session of the thread is open on + * this host, the view reads through it, and it never holds it open. On a + * durable host, the views of one log share its fold and its events. Use + * `open` for a write, or for what needs plugins (`describe`). + * + * With access, `principal` is required, and the thread needs the `read` + * action, else the call throws `AccessDeniedError`. A child thread is + * checked through its root thread. A thread with no index entry gets the + * same answer as a refusal: a read creates no thread. Throws after + * `host.close()`. + * + * @example + * ```ts + * const view = await host.observe(assistant, { threadId: 't-1', principal }) + * const snapshot = await view.snapshot() + * ``` + */ + observe: ( + harness: THarness, + options: Pick, 'threadId' | 'principal' | 'logId'>, + ) => Promise> /** * Is a session of `harness` open on `threadId` on this host, in the tenant * of `principal`? A session that closed on idle is not open, also when a @@ -576,12 +604,14 @@ export interface HarnessHost

{ */ reload: (harness: AnyHarness) => Promise /** - * Close every live session, for every handle, and refuse later opens. - * After it, a call on a handle throws, except a read that gives a value at - * once (`snapshot()`): it reads the last state. With `recoverable`, the - * running work stops with no settlement, and the next host that opens a - * thread runs it again: use it when a durable host shuts down, for - * example for a deploy. Without it, the running work is cancelled. + * Close every live session, for every handle, end the event reads of + * every `observe` view, and refuse later opens and views. After it, a call + * on a view throws, and so does a call on a handle, except a read that + * gives a value at once (`snapshot()`): it reads the last state. With + * `recoverable`, the running work stops with no settlement, and the next + * host that opens a thread runs it again: use it when a durable host shuts + * down, for example for a deploy. Without it, the running work is + * cancelled. * * @example * ```ts @@ -591,10 +621,11 @@ export interface HarnessHost

{ close: (options?: { recoverable?: boolean }) => Promise /** * The `durability.reduce` fold of `harness` for a log that has an open - * session in this host. `logId` is the `logId` of `open`, in the tenant of - * `principal`. `undefined` when no session has the log open, or when the - * reduce of another harness folded the log (the harness of its first - * session folds a shared log). + * session or a running `observe` read in this host. `logId` is the `logId` + * of `open`, in the tenant of `principal`. `undefined` when nothing has the + * log open, or when the reduce of another harness folded the log (the + * harness of its first session folds a shared log). `view.logState()` of + * `host.observe` reads a log that is not open too. * * @example * ```ts @@ -1728,6 +1759,28 @@ export function createHarnessHost

( } } + /** + * A reader of the log `ids` for the thread `threadId`, on a durable host, + * and the reduce that folds the log. It shares the log of the open + * sessions and the other readers. Close it after the read. + */ + const readLog = async ( + store: LogStore, + ids: { logId: string; rootThread: string }, + harness: AnyHarness, + threadId: string, + ) => { + let log = await sharedLog(store, ids, harness) + // The log closed while this read waited, and the host dropped it. + while (!log.isOpen) log = await sharedLog(store, ids, harness) + // `host.close()` ran while the log loaded: it stopped the readers. + if (isClosed) throw hostClosed() + return { + reader: log.reader(threadId), + reduce: resolved.get(ids.logId)?.reduce, + } + } + const host: HarnessHost

= { async open(harness, openOptions) { if (isClosed) throw hostClosed() @@ -1743,6 +1796,43 @@ export function createHarnessHost

( throw error } }, + async observe(harness, { threadId, principal, logId: openLogId }) { + if (isClosed) throw hostClosed() + const tenantId = principal?.tenantId + if (access) { + const user = need(principal, 'host.observe') + const entry = await index.get({ tenantId, threadId }) + // 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. + await allow({ principal: user, action: 'read', threadId, entry }) + } + const key = sessionKey(harness.name, tenantId, threadId) + const logId = openLogId ?? threadId + const ids = { logId: logIdOf({ tenantId, logId }), rootThread: logId } + return observedSession({ + harness, + threadId, + tenantId, + live: () => { + const open = isClosed ? undefined : live.get(key) + // The key names the harness, so the session is a session of + // `harness`, as in `openSession`. + return open && !open.isClosing() + ? (open as SessionImpl) + : undefined + }, + check: () => { + if (isClosed) throw hostClosed() + }, + index, + persistence, + media, + readLog: logStore + ? () => readLog(logStore, ids, harness, threadId) + : undefined, + }) + }, isOpen(harness, { threadId, principal }) { if (access) need(principal, 'host.isOpen') return sessions.has( @@ -1871,6 +1961,8 @@ export function createHarnessHost

( // First: no open starts while the sessions close. isClosed = true for (const key of holds.keys()) cancelIdle(key) + // The event reads of `host.observe` views end. + for (const { log } of resolved.values()) log.stopReaders() await Promise.allSettled(closing.values()) const opened = await Promise.allSettled(sessions.values()) await Promise.all( diff --git a/packages/ai-harness/src/http.ts b/packages/ai-harness/src/http.ts index 2ebcb085ff..dffc9ad975 100644 --- a/packages/ai-harness/src/http.ts +++ b/packages/ai-harness/src/http.ts @@ -25,6 +25,7 @@ import { import { base64url, isRecord } from './utils' import { SessionMetadataError } from './session-entry' import { SessionImpl } from './session/harness-session' +import { liveSessionOf } from './session/observed-session' import type { ModelMessage, RunStore, @@ -36,7 +37,8 @@ import type { AnyHarness } from './define' import type { HarnessHost, HostEvent } from './host' import type { HostFrame } from './protocol' import type { HarnessSession } from './session/index' -import type { Principal, SessionEvent, UserInput } from './types' +import type { ObservedSession } from './session/observed-session' +import type { Cursor, Principal, SessionEvent, UserInput } from './types' /** * Who a handler or a socket serves when its host has no `access`. `true`: @@ -281,10 +283,10 @@ function fromBase64url(value: string) { */ async function mediaResponse( request: Request, - session: HarnessSession, + view: ObservedSession, id: string, ) { - const record = await session.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') { @@ -293,7 +295,7 @@ async function mediaResponse( headers: { ...MEDIA_HEADERS, 'Content-Range': `bytes */${record.size}` }, }) } - const bytes = await session.loadMedia({ id, range }) + const bytes = await view.loadMedia({ id, range }) const headers = { ...MEDIA_HEADERS, 'Content-Type': record.mimeType, @@ -340,6 +342,30 @@ function requestUsers

( } } +/** + * The view of `host.observe` for a read of the thread `threadId`, or + * `undefined` for a thread with no index entry on a host with access: the + * host refuses such a thread as it refuses a stranger, and a read of it is + * not found. Throws `AccessDeniedError` for a stranger. + */ +async function observeToRead

( + host: HarnessHost

, + harness: AnyHarness, + { threadId, principal }: { threadId: string; principal: P | undefined }, +) { + const user = principal ? { principal } : {} + try { + return await host.observe(harness, { threadId, ...user }) + } catch (error) { + if (!(error instanceof AccessDeniedError)) throw error + // `sessions.get` refuses a stranger too, and has no entry to give for a + // thread that does not exist. Only the refused path reads it. + const entry = await host.sessions.get({ threadId, ...user }) + if (entry === undefined) return undefined + throw error + } +} + /** `event` is the revoke of `principal` on the thread `threadId`. */ const isRevokeOf = ( event: HostEvent, @@ -434,6 +460,11 @@ function watchRevoke

( * read, as SSE: status changes, session index changes, and the revokes of * this principal. First the current status of each open session. * + * The `GET` routes of a thread, except `describe`, read with `host.observe`: + * they open no session, mount no plugins, and write nothing. On a host with + * no log, `GET events` of a thread whose session is not open on this + * instance ends at once, and the client connects again. + * * With host `access`, `access.identify` reads the user of each request * (`null` is 401), and the host checks each action: a refused action is 403. * A read of a thread with no entry is 404, and it does not start a session. @@ -471,9 +502,14 @@ export function createHarnessHandler

( durable?.runs && tenantStores({ runs: durable.runs }, principal?.tenantId).runs + /** The view of a thread for a read route. See `observeToRead`. */ + const viewFor = (principal: P | undefined, threadId: string) => + observeToRead(host, harness, { threadId, principal }) + /** - * Open a thread for a read. With access, a thread with no entry is not - * found, so a read never starts a session: `undefined`. + * Open a thread for `describe`, which needs its plugins. With access, a + * thread with no entry is not found, so a read never starts a session: + * `undefined`. */ const openToRead = async (principal: P | undefined, threadId: string) => { const isMissing = @@ -564,7 +600,7 @@ export function createHarnessHandler

( } /** - * Serve a signed media URL. The URL names its user: the handle opens as + * Serve a signed media URL. The URL names its user: the read is checked as * that user, so a user who lost access (unshare, revoke) is refused. */ const serveSigned = async (request: Request, url: URL) => { @@ -595,11 +631,9 @@ export function createHarnessHandler

( const user = userId ? ({ id: userId, ...(tenantId ? { tenantId } : {}) } as P) : undefined - const session = await openToRead(user, payload.threadId) - if (!session) return json({ error: 'not found' }, 404) - return withSession(session, () => - mediaResponse(request, session, payload.id), - ) + const view = await viewFor(user, payload.threadId) + if (!view) return json({ error: 'not found' }, 404) + return mediaResponse(request, view, payload.id) } /** @@ -622,23 +656,22 @@ export function createHarnessHandler

( tenantId: principal?.tenantId, operationId: runId, }) - if (running?.operation(runId)) { + const turn = running?.operation(runId) + if (running && turn) { const { threadId } = running - const session = await openToRead(principal, threadId) - if (!session) return json({ error: 'unknown run' }, 404) - return withSession(session, (keep) => { - const turn = session.operation(runId) - if (!turn) return json({ error: 'unknown run' }, 404) - const startedCursor = session - .snapshot() - .activeOperations.find((item) => item.id === runId)?.startedCursor - const from = resumeFrom ?? startedCursor ?? '0' - return cursorResponse( - (signal) => turn.events({ from, signal }), - true, - keep(closeOnRevoke(principal, threadId)), - ) - }) + // The view checks the read. The turn keeps its session busy, so the + // session does not close on idle while the turn streams. + const view = await viewFor(principal, threadId) + if (!view) return json({ error: 'unknown run' }, 404) + const startedCursor = running + .snapshot() + .activeOperations.find((item) => item.id === runId)?.startedCursor + const from = resumeFrom ?? startedCursor ?? '0' + return cursorResponse( + (signal) => turn.events({ from, signal }), + true, + closeOnRevoke(principal, threadId), + ) } // Another instance runs the turn. A durable host finds its thread in the // run store, and its events in the shared log. Not a child run or an @@ -650,25 +683,22 @@ export function createHarnessHandler

( record.parentRunId === undefined if (!record || !isChatTurn) return json({ error: 'unknown run' }, 404) const { threadId } = record - const user = principal ? { principal } : {} - const entry = await host.sessions.get({ threadId, ...user }) + const view = await viewFor(principal, threadId) + if (!view) return json({ error: 'unknown run' }, 404) + const entry = await view.entry() if (entry?.harness !== undefined && entry.harness !== harness.name) { return json({ error: 'unknown run' }, 404) } - const shared = await openToRead(principal, threadId) - if (!shared) 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. - return withSession(shared, (keep) => - cursorResponse( - (signal) => - followOperation( - shared.events({ from: resumeFrom ?? '0', signal }), - runId, - ), - true, - keep(closeOnRevoke(principal, threadId)), - ), + return cursorResponse( + (signal) => + followOperation( + view.events({ from: resumeFrom ?? '0', signal }), + runId, + ), + true, + closeOnRevoke(principal, threadId), ) } @@ -799,40 +829,35 @@ export function createHarnessHandler

( if (runId) return await joinTurn(request, principal, runId) const threadId = url.searchParams.get('threadId') if (!threadId) return json({ error: 'threadId is required' }, 400) - const session = await openToRead(principal, threadId) + const view = await viewFor(principal, threadId) // A ChatClient hydrates a new thread before its first message, and a // non-OK answer is an error there. A new thread has nothing to load. - if (!session) { + if (!view) { return json({ messages: [], activeRun: null, interrupts: null }) } - return withSession(session, async () => - json(await hydration(session, runsOf(principal))), - ) + return json(await hydration(view, runsOf(principal))) } if (request.method === 'GET' && route === 'events') { const threadId = url.searchParams.get('threadId') if (!threadId) return json({ error: 'threadId is required' }, 400) - const session = await openToRead(principal, threadId) - if (!session) return json({ error: 'not found' }, 404) + const view = await viewFor(principal, threadId) + if (!view) return json({ error: 'not found' }, 404) const from = request.headers.get('Last-Event-ID') ?? url.searchParams.get('from') ?? undefined - return withSession(session, (keep) => - sseResponse( - async function* (signal) { - const events = session.events({ - ...(from ? { from } : {}), - signal, - }) - for await (const entry of events) { - const frame: HostFrame = { type: 'harness.event', ...entry } - yield { id: entry.cursor, data: frame } - } - }, - keep(closeOnRevoke(principal, threadId)), - ), + // On a host with no log, a thread whose session is not open here has + // nothing to stream: the stream ends, and the client connects again. + return sseResponse( + async function* (signal) { + const events = view.events({ ...(from ? { from } : {}), signal }) + for await (const entry of events) { + const frame: HostFrame = { type: 'harness.event', ...entry } + yield { id: entry.cursor, data: frame } + } + }, + closeOnRevoke(principal, threadId), ) } @@ -860,20 +885,22 @@ export function createHarnessHandler

( ) } - if ( - request.method === 'GET' && - (route === 'transcript' || route === 'describe') - ) { + if (request.method === 'GET' && route === 'transcript') { + const threadId = url.searchParams.get('threadId') + if (!threadId) return json({ error: 'threadId is required' }, 400) + const view = await viewFor(principal, threadId) + if (!view) return json({ error: 'not found' }, 404) + return json(await view.transcript()) + } + + // The only read that opens the session: it needs the plugins. + if (request.method === 'GET' && route === 'describe') { const threadId = url.searchParams.get('threadId') if (!threadId) return json({ error: 'threadId is required' }, 400) const session = await openToRead(principal, threadId) if (!session) return json({ error: 'not found' }, 404) - return withSession(session, async () => - json( - route === 'transcript' - ? await session.transcript() - : describeForClient(harness, session), - ), + return withSession(session, () => + json(describeForClient(harness, session)), ) } @@ -883,19 +910,17 @@ export function createHarnessHandler

( if (!threadId) return json({ error: 'threadId is required' }, 400) if (!subagentRunId) return json({ error: 'subagentRunId is required' }, 400) - const session = await openToRead(principal, threadId) - if (!session) return json({ error: 'not found' }, 404) - return withSession(session, async () => - json(await session.agentTranscript(subagentRunId)), - ) + const view = await viewFor(principal, threadId) + if (!view) return json({ error: 'not found' }, 404) + return json(await view.agentTranscript(subagentRunId)) } if (request.method === 'GET' && route === 'snapshot') { const threadId = url.searchParams.get('threadId') if (!threadId) return json({ error: 'threadId is required' }, 400) - const session = await openToRead(principal, threadId) - if (!session) return json({ error: 'not found' }, 404) - return withSession(session, () => json(session.snapshot())) + const view = await viewFor(principal, threadId) + if (!view) return json({ error: 'not found' }, 404) + return json(await view.snapshot()) } if (request.method === 'GET' && route === 'sessions') { @@ -1017,31 +1042,27 @@ export function createHarnessHandler

( const id = url.searchParams.get('id') if (!threadId || !id) return json({ error: 'threadId and id are required' }, 400) - const session = await openToRead(principal, threadId) - if (!session) return json({ error: 'not found' }, 404) - return withSession(session, async () => { - if (route === 'media') return mediaResponse(request, session, id) - if (!(await session.getMedia(id))) { - return json({ error: 'not found' }, 404) - } - const expiresAt = Date.now() + SIGNED_URL_TTL_MS - const payload: SignedMedia = { - tenantId: principal?.tenantId ?? '', - threadId, - id, - userId: principal?.id ?? '', - exp: String(expiresAt), - } - // An empty field stays out of the URL. The served URL reads it as ''. - const fields = Object.entries(payload).filter( - ([, value]) => value !== '', - ) - const query = new URLSearchParams([ - ...fields, - ['sig', await sign(payload)], - ]) - return json({ path: `media?${query}`, expiresAt }) - }) + 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))) { + return json({ error: 'not found' }, 404) + } + const expiresAt = Date.now() + SIGNED_URL_TTL_MS + const payload: SignedMedia = { + tenantId: principal?.tenantId ?? '', + threadId, + id, + userId: principal?.id ?? '', + exp: String(expiresAt), + } + // An empty field stays out of the URL. The served URL reads it as ''. + const fields = Object.entries(payload).filter(([, value]) => value !== '') + const query = new URLSearchParams([ + ...fields, + ['sig', await sign(payload)], + ]) + return json({ path: `media?${query}`, expiresAt }) } return json({ error: 'not found' }, 404) } @@ -1071,24 +1092,24 @@ export function createHarnessHandler

( * waits for. On a durable host, `runs` finds a turn that another instance * runs. */ -async function hydration(session: HarnessSession, runs?: RunStore) { - const { pendingInterrupts, activeOperations } = session.snapshot() +async function hydration(view: ObservedSession, runs?: RunStore) { + const { pendingInterrupts, activeOperations } = await view.snapshot() const [first] = pendingInterrupts const runId = first ? readInterruptBinding(first)?.interruptedRunId : undefined + // Only a session that is open here runs operations. + const session = liveSessionOf(view) const running = activeOperations.find( (item) => item.kind === 'chat' && - session.operation(item.id)?.status() === 'running', + session?.operation(item.id)?.status() === 'running', ) - const remote = running - ? undefined - : await runs?.findActiveRun(session.threadId) + const remote = running ? undefined : await runs?.findActiveRun(view.threadId) const runningId = running?.id ?? (remote?.kind !== 'agent' ? remote?.runId : undefined) return { - messages: modelMessagesToUIMessages(await session.transcript()), + messages: modelMessagesToUIMessages(await view.transcript()), activeRun: runningId ? { runId: runningId } : null, interrupts: runId ? { runId, pending: pendingInterrupts } : null, } @@ -1140,6 +1161,11 @@ const errorMessage = (error: unknown) => * input. When `host.revoke` or `unshare` takes the access of the user, the * socket sends a `forbidden` error frame and closes. * + * The subscribe and the snapshots read with `host.observe`, so they open no + * session. The first input opens the session, and the socket holds it until + * it closes. On a host with no log, a thread whose session is not open has + * no events: the stream starts when the first input opens the session. + * * Throws at the start for a host without `access` and without `openAccess`, * and for a host with both. * @@ -1153,11 +1179,85 @@ export function handleHarnessSocket

( ) { const { host, harness, socket, request } = options const usersOf = requestUsers('handleHarnessSocket', host, options.openAccess) + const isDurable = durableStoresOf(host) !== undefined const reader = new AbortController() 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 the open. - let opened: Promise | undefined + // refused and an input sent right after the subscribe waits for it. It + // resolves to `undefined` when the host refused the thread. The view is + // `undefined` for a thread with no index entry yet: its first input + // makes it. + let subscribed: + | Promise< + { threadId: string; view: ObservedSession | undefined } | undefined + > + | undefined + /** The session that the first input opens. The socket holds it. */ + let held: Promise | undefined + /** `held`, once it opened. */ + let holding: HarnessSession | undefined + let onHeld: (session: HarnessSession | undefined) => void = () => {} + const firstHeld = new Promise((resolve) => { + onHeld = resolve + }) + reader.signal.addEventListener('abort', () => onHeld(undefined), { + once: true, + }) + /** The held session, opened by the first input that needs it. */ + const hold = (threadId: string, principal: P | undefined) => { + if (held) return held + const opening = host + .open(harness, { threadId, ...(principal ? { principal } : {}) }) + .then((session) => { + // The socket holds the handle until it closes. + const release = () => void session.close() + if (reader.signal.aborted) release() + else reader.signal.addEventListener('abort', release, { once: true }) + holding = session + onHeld(session) + return session + }) + held = opening + // A failed open holds nothing, so the next input opens again. + opening.catch(() => { + if (held === opening) held = undefined + }) + return opening + } + /** + * Send the events of the thread: the view first. A durable view follows + * the log while the socket is open. On a host with no log, the view reads + * the open session, and it ends when no session is open. Then the events + * of the session that the first input opens come next. That session has a + * new feed: its cursors start again. A view of the held session ends with + * that session. + */ + const stream = async ( + view: ObservedSession | undefined, + from: Cursor | undefined, + ) => { + let cursor = from + const forward = async (events: AsyncIterable) => { + for await (const entry of events) { + cursor = entry.cursor + send({ type: 'harness.event', ...entry }) + } + } + const readOptions = () => ({ + ...(cursor ? { from: cursor } : {}), + signal: reader.signal, + }) + if (view) await forward(view.events(readOptions())) + const isDone = + reader.signal.aborted || + holding !== undefined || + (view !== undefined && isDurable) + if (isDone) return + const session = await firstHeld + if (!session) return + if (!isDurable) cursor = undefined + await forward(session.events(readOptions())) + } const send = (frame: HostFrame) => { try { socket.send(JSON.stringify(frame)) @@ -1181,20 +1281,20 @@ export function handleHarnessSocket

( const { principal } = user const frame = parseControlFrame(String(message.data)) if (frame.type === 'harness.subscribe') { - if (opened) throw new Error('Already subscribed.') + if (subscribed) throw new Error('Already subscribed.') const { threadId } = frame - opened = host - .open(harness, { threadId, ...(principal ? { principal } : {}) }) - .catch((error: unknown) => { + subscribed = observeToRead(host, harness, { + threadId, + principal, + }).then( + (view) => ({ threadId, view }), + (error: unknown) => { if (!(error instanceof AccessDeniedError)) throw error return undefined - }) - const session = await opened - if (!session) return refuse('forbidden') - // The socket holds the handle until it closes. - const release = () => void session.close() - if (reader.signal.aborted) release() - else reader.signal.addEventListener('abort', release, { once: true }) + }, + ) + const subscription = await subscribed + if (!subscription) return refuse('forbidden') watchRevoke( host, { principal, threadId, signal: reader.signal }, @@ -1205,14 +1305,7 @@ export function handleHarnessSocket

( v: HARNESS_PROTOCOL_VERSION, threadId: frame.threadId, }) - const events = session.events({ - ...(frame.from ? { from: frame.from } : {}), - signal: reader.signal, - }) - void (async () => { - for await (const entry of events) - send({ type: 'harness.event', ...entry }) - })() + void stream(subscription.view, frame.from) .catch((error: unknown) => send({ type: 'harness.error', message: errorMessage(error) }), ) @@ -1223,12 +1316,17 @@ export function handleHarnessSocket

( }) return } - const session = await opened - if (!session) throw new Error('Send harness.subscribe first.') + const subscription = await subscribed + if (!subscription) throw new Error('Send harness.subscribe first.') if (frame.type === 'harness.snapshot') { - send({ type: 'harness.snapshot', snapshot: session.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 }) return } + const session = await hold(subscription.threadId, principal) const receipt = await applyInput({ harness, session, diff --git a/packages/ai-harness/src/index.ts b/packages/ai-harness/src/index.ts index 5fab83b942..790c1c5829 100644 --- a/packages/ai-harness/src/index.ts +++ b/packages/ai-harness/src/index.ts @@ -141,6 +141,7 @@ export type { SessionInspection, SessionSnapshot, } from './session/index' +export type { ObservedSession } from './session/observed-session' export type { SessionUsage, UsageCounts } from './usage' export type { OnHarnessEvent } from './feed' diff --git a/packages/ai-harness/src/log.ts b/packages/ai-harness/src/log.ts index a402b53b1d..87e4743778 100644 --- a/packages/ai-harness/src/log.ts +++ b/packages/ai-harness/src/log.ts @@ -98,6 +98,21 @@ export function revertOf(value: unknown): RevertState | undefined { } } +/** + * The messages that the transcript shows: while `revert` stands, the + * messages after its message are hidden. + */ +export function visibleMessages( + messages: ReadonlyArray, + revert: RevertState | undefined, +) { + const at = + revert === undefined + ? -1 + : messages.findIndex((message) => message.id === revert.messageId) + return at < 0 ? [...messages] : messages.slice(0, at + 1) +} + const HARNESS_RECORD_TYPES = new Set([ 'harness.event', 'harness.transcript', @@ -843,6 +858,29 @@ export interface LogWriter extends EventFeed { catchUp: () => Promise } +/** + * A read view of one session of a shared log, from {@link SharedLog.reader}. + * It follows what the writers append, and it never writes. + */ +export interface LogReader { + /** The fold of the session. Empty when the log has no record of it. */ + readonly state: LogState + /** The `reduce` fold of the whole log. */ + readonly reduced: TState | undefined + /** The log position. */ + head: () => Cursor + /** + * The events of the session after `from` (exclusive), then each new one, + * until `signal` aborts, the reader closes, or the log fails. + */ + read: (options: { + from?: Cursor + signal?: AbortSignal + }) => AsyncIterable + /** Let go of the log. A second call does nothing. */ + close: () => void +} + /** * Writes one log for every session that shares it, and keeps its fold. * All writes run one after another. Each batch goes at the next position. @@ -899,6 +937,10 @@ export class SharedLog { private readonly waiters = new Set<() => void>() /** The threads with an open view. */ private readonly threads = new Set() + /** The open readers. Like the views, they keep the log open. */ + private readers = 0 + /** `stopReaders` ran: every read of a reader ends. */ + private readersStopped = false private idle = false private writing = false private hasWritten = false @@ -952,13 +994,59 @@ export class SharedLog { } /** - * False after the last view closed or a write failed. The host then opens - * a new log. + * False after the last view and the last reader closed, or after a write + * failed. The host then opens a new log. */ get isOpen() { return !this.idle } + /** + * A read view of session `threadId`: its fold and its events, which follow + * what the writers append. It never writes, and it is not a view of the + * thread, so a session of the thread can open while it reads. The readers + * of one log share its fold and its tail of events. Close it after the + * read: the log stays open while a view or a reader is open. + */ + reader(threadId: string): LogReader { + if (this.idle) { + throw new Error(`The log ${JSON.stringify(this.logId)} is closed.`) + } + this.readers += 1 + const shared = this.state + let closed = false + return { + get state() { + // Not `sessionOf`: a read adds no session to the fold. + return shared.sessions.get(threadId) ?? emptyLogState() + }, + get reduced() { + return shared.reduced + }, + head: () => this.head(), + read: (options) => + this.read(threadId, { + ...options, + isClosed: () => + closed || this.readersStopped || this.failure !== undefined, + }), + close: () => { + if (closed) return + closed = true + this.readers -= 1 + // Its readers end. The others wait again. + this.wake() + this.idleWhenUnused() + }, + } + } + + /** End the reads of every reader, for a host that closes. */ + stopReaders() { + this.readersStopped = true + this.wake() + } + /** The view of session `threadId`. Close it when the session closes. */ view(threadId: string, onFailure: (error: unknown) => void): LogWriter { if (this.idle) { @@ -1022,12 +1110,23 @@ export class SharedLog { } // The readers of the closed view end. The others wait again. this.wake() - // Idle only after the last append landed. Until then the host keeps - // this log, so an open of the thread gets it, not a fold without it. + this.idleWhenUnused() + } + + /** + * Close the log when no view and no reader uses it. Only after the last + * append landed: until then the host keeps this log, so an open of the + * thread gets it, not a fold without it. + */ + private idleWhenUnused() { void this.flush() .catch(() => {}) .then(() => { if (this.threads.size > 0) return + // No session of this host writes any more, so the readers follow the + // records of another host, as a host that only reads does. + this.hasWritten = false + if (this.readers > 0) return this.idle = true this.unsubscribe() this.onIdle() @@ -1216,6 +1315,8 @@ export class SharedLog { this.failure = error this.idle = true for (const listener of [...this.failureListeners]) listener(error) + // The readers end: their log stopped. + this.wake() this.onIdle() } diff --git a/packages/ai-harness/src/session/observed-session.ts b/packages/ai-harness/src/session/observed-session.ts new file mode 100644 index 0000000000..c3ca0542e4 --- /dev/null +++ b/packages/ai-harness/src/session/observed-session.ts @@ -0,0 +1,324 @@ +import { loadChildTranscript } from '@tanstack/ai-persistence' +import { logIdOf, logMessageStore, revertOf, visibleMessages } from '../log' +import { createMediaStore } from '../media' +import { tenantStores } from '../tenant-stores' +import { emptyUsage, isSessionUsage } from '../usage' +import { INTERRUPTED, USAGE, isInterruptedTurn } from './helpers' +import type { ModelMessage } from '@tanstack/ai' +import type { BlobRange, SessionIndexEntry } from '@tanstack/ai-persistence' +import type { LogReader, ReduceOptions } from '../log' +import type { MediaPersistence } from '../media' +import type { SessionUsage } from '../usage' +import type { + AnyHarness, + HarnessLogStateOf, + SessionMetadataOf, +} from '../define' +import type { HarnessPersistence, SessionIndexWriter } from '../host' +import type { Cursor, MediaRecord, SessionEvent } from '../types' +import type { SessionImpl } from './harness-session' +import type { SessionSnapshot } from './types' + +/** + * A read-only view of one thread, from `host.observe`. It mounts no plugins, + * claims nothing, and writes nothing, so a read of an idle thread costs only + * store reads. When the session of the thread is open on this host, the view + * reads through it. The view never holds that session: it closes on idle as + * if the view did not exist. Each call reads again, so a view can live as + * long as you need it. + * + * @example + * ```ts + * const view = await host.observe(assistant, { threadId, principal }) + * const messages = await view.transcript() + * ``` + */ +export interface ObservedSession { + readonly threadId: string + /** + * The events of the thread after `from` (exclusive), then each new event, + * until `signal` aborts. + * + * - A durable host (`stores.log`) reads the log. The read follows what the + * sessions of every host append, also when no session is open. + * - A host with no log has the events of an open session only. When the + * session is not open on this host, the read ends at once, as there is + * nothing to stream. When the session closes, the read ends. + * + * Stop the read with `signal`, or by leaving the loop: on a durable host, + * a read keeps the log open until it stops. + */ + events: (options?: { + from?: Cursor + signal?: AbortSignal + }) => AsyncIterable + /** + * The saved messages of the thread, oldest first, as + * `session.transcript()` gives them: while a revert stands, the messages + * after its message are hidden. + */ + transcript: () => Promise> + /** + * The snapshot of the thread. When the session is open on this host, it is + * `session.snapshot()`. Else it is built from the stores, with no plugin + * setup: + * + * - `status` is `requires_action` while interrupts wait, else `idle`. + * - `pendingInterrupts` come from `stores.metadata`, and `usage` from the + * log or `stores.metadata`. + * - `plugins` has only the plugin state that `stores.metadata` keeps. The + * state of a plugin that never saved one is left out, because its first + * value comes from the plugin setup. + * - There are no active operations, queued turns, waiting inputs, or + * questions: those live in an open session. + * - `cursor` is the log position on a durable host, else `'0'`. + */ + snapshot: () => Promise + /** + * 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> + /** The record of a media file of the thread, or `null` when it is not found. */ + getMedia: (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. + */ + loadMedia: (options: { + id: string + range?: BlobRange + }) => Promise> + /** The token usage of the thread, as `session.usage()` gives it. */ + usage: () => Promise + /** The session index entry of the thread, or `undefined` when it has none. */ + entry: () => Promise< + SessionIndexEntry> | undefined + > + /** + * The `durability.reduce` fold of the harness for the log of the thread, + * as `host.logState` gives it. `undefined` on a host with no log, and when + * the reduce of another harness folds the log. + */ + logState: () => Promise | undefined> +} + +/** What a view reads, from the host that makes it. */ +export interface ObservedSources { + harness: THarness + threadId: string + tenantId: string | undefined + /** The open session of the thread on this host, or `undefined`. */ + live: () => SessionImpl | undefined + /** Throws when the host is closed. */ + check: () => void + index: SessionIndexWriter + /** The stores of the host, in all tenants. */ + persistence: HarnessPersistence + media: MediaPersistence + /** + * On a durable host: a reader of the log of the thread, and the `reduce` + * that folds the log. The caller closes the reader. + */ + readLog: + | (() => Promise<{ + reader: LogReader + reduce: ReduceOptions | undefined + }>) + | undefined +} + +/** The open session of each view, for the HTTP handler. Not public API. */ +const liveSessions = new WeakMap< + object, + () => Pick | undefined +>() + +/** + * @internal The session that `view` reads through now, or `undefined` when + * it is not open on this host. + */ +export function liveSessionOf(view: ObservedSession) { + return liveSessions.get(view)?.() +} + +/** The view of `host.observe`. See {@link ObservedSession}. */ +export function observedSession( + sources: ObservedSources, +) { + const { harness, threadId, tenantId, live, check, index, readLog } = sources + const { stores } = sources.persistence + const { metadata } = stores + const storeThreadId = logIdOf({ tenantId, logId: threadId }) + const tenant = tenantId !== undefined ? { tenantId } : {} + const messages = tenantStores(stores, tenantId).messages + const project = harness.durability?.project + // The children of a durable session each have their own log. + const childMessages = stores.log + ? logMessageStore({ + store: stores.log, + ...(project ? { project } : {}), + ...tenant, + }) + : messages + const mediaStore = createMediaStore({ + persistence: sources.media, + threadId, + ...tenant, + options: harness.media, + }) + + /** The open session of the thread, once the host is known to be open. */ + const open = () => { + check() + return live() + } + + /** `read` on a reader of the log from `logReader`, which closes after it. */ + const fromLog = async ( + logReader: NonNullable, + read: (reader: LogReader, reduce: ReduceOptions | undefined) => T, + ) => { + const { reader, reduce } = await logReader() + try { + return read(reader, reduce) + } finally { + reader.close() + } + } + + /** The usage totals and the cursor of a session that is not open. */ + const storedTotals = async () => { + if (readLog) { + return fromLog(readLog, (reader) => ({ + usage: structuredClone(reader.state.usage), + cursor: reader.head(), + })) + } + const stored = await metadata?.get(USAGE, storeThreadId) + return { + usage: isSessionUsage(stored) ? stored : emptyUsage(), + cursor: '0', + } + } + + /** The plugin state that `stores.metadata` keeps, by plugin name. */ + const storedPluginState = async () => { + const state = new Map() + if (!metadata) return Object.fromEntries(state) + // The plugin list only: no plugin is set up. + const names = (harness.plugins?.() ?? []).map((plugin) => plugin.name) + const saved = await Promise.all( + names.map(async (name) => ({ + name, + value: await metadata.get(`plugin:${name}`, storeThreadId), + })), + ) + for (const { name, value } of saved) { + if (value !== null && value !== undefined) state.set(name, value) + } + // fromEntries defines own keys, so a plugin named `__proto__` stays a key. + return Object.fromEntries(state) + } + + async function* events(options: { from?: Cursor; signal?: AbortSignal }) { + check() + if (readLog) { + const { reader } = await readLog() + try { + yield* reader.read(options) + } finally { + reader.close() + } + return + } + // A host with no log: only an open session has events. + const session = live() + if (session) yield* session.events(options) + } + + const view: ObservedSession = { + threadId, + events: (options) => events(options ?? {}), + async transcript() { + const session = open() + if (session) return session.transcript() + if (readLog) { + return fromLog(readLog, ({ state }) => + visibleMessages(state.messages, state.revert), + ) + } + const [stored, revert] = await Promise.all([ + messages?.loadThread(threadId), + metadata?.get('harness:revert', storeThreadId), + ]) + return visibleMessages(stored ?? [], revertOf(revert)) + }, + async snapshot() { + const session = open() + if (session) return session.snapshot() + const [interrupted, plugins, totals] = await Promise.all([ + metadata?.get(INTERRUPTED, storeThreadId), + storedPluginState(), + storedTotals(), + ]) + const pendingInterrupts = isInterruptedTurn(interrupted) + ? interrupted.interrupts + : [] + const snapshot: SessionSnapshot = { + threadId, + status: pendingInterrupts.length > 0 ? 'requires_action' : 'idle', + activeOperations: [], + queuedTurns: 0, + waitingInputs: [], + pendingInterrupts, + pendingQuestions: [], + plugins, + ...totals, + } + return snapshot + }, + async agentTranscript(subagentRunId) { + check() + if (!childMessages) return [] + return ( + (await loadChildTranscript(childMessages, subagentRunId, threadId)) ?? + [] + ) + }, + getMedia(id) { + check() + return mediaStore.get(id) + }, + loadMedia({ id, range }) { + check() + return mediaStore.load(id, range) + }, + async usage() { + const session = open() + if (session) return session.usage() + return (await storedTotals()).usage + }, + async entry() { + check() + // The schema checked the metadata on each write. + return (await index.get({ tenantId, threadId })) as + | SessionIndexEntry> + | undefined + }, + async logState() { + check() + if (!readLog) return undefined + return fromLog(readLog, (reader, reduce) => + // The reduce of another harness folded this log: its state has + // another type. Else the reduce of `harness` folded it, so it has + // that type. A map of logs cannot carry a state type per log. + reduce === harness.durability?.reduce + ? (reader.reduced as HarnessLogStateOf | undefined) + : undefined, + ) + }, + } + liveSessions.set(view, live) + return view +} diff --git a/packages/ai-harness/src/session/state.ts b/packages/ai-harness/src/session/state.ts index 1c068173e3..0e4384ef46 100644 --- a/packages/ai-harness/src/session/state.ts +++ b/packages/ai-harness/src/session/state.ts @@ -8,7 +8,7 @@ import { loadChildTranscript } from '@tanstack/ai-persistence' import { checkConfigValue } from '../config' import { isRecord } from '../utils' import { changedEntry } from '../session-entry' -import { threadOf } from '../log' +import { threadOf, visibleMessages } from '../log' import { mediaIdOf, mediaOfMessage, mediaPart } from '../media-ref' import { BackgroundJobs, @@ -491,9 +491,10 @@ export abstract class SessionState< * the messages after its message are hidden. */ async transcript() { - const messages = [...(await this.messages.loadThread(this.threadId))] - const at = this.revertedAt(messages) - return at < 0 ? messages : messages.slice(0, at + 1) + return visibleMessages( + await this.messages.loadThread(this.threadId), + this.reverted, + ) } /** diff --git a/packages/ai-harness/tests/handles.test.ts b/packages/ai-harness/tests/handles.test.ts index 5ca42eef61..f616f3230f 100644 --- a/packages/ai-harness/tests/handles.test.ts +++ b/packages/ai-harness/tests/handles.test.ts @@ -238,15 +238,20 @@ describe('the HTTP handler', () => { await host.close() }, 60_000) - it('releases the handle of an event stream when the client leaves', async () => { - const { host, harness, isOpen } = setup({ idleCloseMs: 0 }) + it('holds no handle for an event stream, so the session still closes on idle', async () => { + const { host, harness, isOpen, open } = setup({ idleCloseMs: 0 }) const handler = createHarnessHandler({ host, harness, openAccess: true }) + const session = await open('t1') + await session.send('Hello.') const response = await handler(new Request(`${BASE}/events?threadId=t1`)) + // Only the handle of `open`. expect(hostCountsOf(host).handles).toBe(1) - await response.body?.cancel() + await session.close() - await vi.waitFor(() => expect(isOpen('t1')).toBe(false)) + // The stream read the open session, and it ends with it. + expect(await response.text()).toContain('Hello back.') + expect(isOpen('t1')).toBe(false) expect(hostCountsOf(host).handles).toBe(0) await host.close() }) diff --git a/packages/ai-harness/tests/observe.test.ts b/packages/ai-harness/tests/observe.test.ts new file mode 100644 index 0000000000..26e23e1d1f --- /dev/null +++ b/packages/ai-harness/tests/observe.test.ts @@ -0,0 +1,340 @@ +import { afterEach, describe, expect, it, vi } from 'vitest' +import { EventType } from '@tanstack/ai' +import { memoryLogStore, memoryPersistence } from '@tanstack/ai-persistence' +import { + AccessDeniedError, + createHarnessHost, + defineHarness, + definePlugin, + membersOnly, +} from '../src' +import { createHarnessHandler, handleHarnessSocket } from '../src/server' +import { hostCountsOf } from '../src/host' +import { fakeSocket, mockAdapter, text } from './helpers' +import type { HarnessHostOptions, Principal, SessionEvent } from '../src' + +const BASE = 'http://x/api/harness' +const THREAD = 't1' + +/** + * A plugin like the MCP plugin: its setup would start a stdio server. It + * counts its setups, and saves how many it saw as its plugin state. + */ +function stdioLikePlugin() { + const setups = { count: 0 } + const plugin = definePlugin({ + name: 'test/stdio-like', + setup: async (ctx) => { + setups.count += 1 + await ctx + .state({ mounts: 0 }) + .update((saved) => ({ mounts: saved.mounts + 1 })) + return {} + }, + }) + return { plugin, setups } +} + +/** The stores of a host with no log, or of a durable host. */ +function persistenceOf(shape: 'no log' | 'durable') { + const { messages, runs, metadata, sessions } = memoryPersistence().stores + const shared = { runs, metadata, sessions } + return shape === 'durable' + ? { stores: { log: memoryLogStore(), ...shared } } + : { stores: { messages, ...shared } } +} + +/** + * A host whose sessions close at once when no handle holds them, and a + * harness with the stdio-like plugin. `writes()` counts the index and + * metadata writes since the last `reset()`. + */ +function setup( + shape: 'no log' | 'durable', + options: Pick = {}, +) { + const persistence = persistenceOf(shape) + const { plugin, setups } = stdioLikePlugin() + const { adapter } = mockAdapter(() => text('Hello back.')) + const harness = defineHarness({ + name: 'test/observe', + adapter, + plugins: () => [plugin], + }) + const host = createHarnessHost({ + persistence, + log: { coalesceMs: 0 }, + lifecycle: { idleCloseMs: 0 }, + ...options, + }) + const { sessions, metadata } = persistence.stores + const spies = [ + vi.spyOn(sessions, 'upsertIf'), + vi.spyOn(metadata, 'set'), + vi.spyOn(metadata, 'setIf'), + vi.spyOn(metadata, 'delete'), + ] + const writes = () => + spies.reduce((total, spy) => total + spy.mock.calls.length, 0) + /** Open the thread, run one turn, and close the handle: the session closes. */ + const seed = async (principal?: Principal) => { + const session = await host.open(harness, { + threadId: THREAD, + ...(principal ? { principal } : {}), + }) + await session.send('Hello.') + await session.close() + } + const reset = () => { + setups.count = 0 + for (const spy of spies) spy.mockClear() + } + const isOpen = () => host.isOpen(harness, { threadId: THREAD }) + return { host, harness, persistence, setups, writes, seed, reset, isOpen } +} + +/** The text of the first text delta of `events`. */ +async function firstText(events: AsyncIterable) { + for await (const { event } of events) { + if (event.type === EventType.TEXT_MESSAGE_CONTENT) return event.delta + } + return undefined +} + +/** Every event of `events`, for a read that ends by itself. */ +async function collect(events: AsyncIterable) { + const all: Array = [] + for await (const entry of events) all.push(entry) + return all +} + +afterEach(() => { + vi.restoreAllMocks() +}) + +describe.each(['no log', 'durable'] as const)( + 'host.observe on a host with %s', + (shape) => { + it('reads an idle thread with no plugin setup and no store write', async () => { + const { host, harness, setups, writes, seed, reset, isOpen } = + setup(shape) + await seed() + reset() + + const view = await host.observe(harness, { threadId: THREAD }) + const transcript = await view.transcript() + const snapshot = await view.snapshot() + const entry = await view.entry() + + expect(transcript.map((message) => message.content)).toEqual([ + 'Hello.', + 'Hello back.', + ]) + // The saved plugin state, read with no setup. + expect(snapshot).toMatchObject({ + threadId: THREAD, + status: 'idle', + activeOperations: [], + pendingInterrupts: [], + plugins: { 'test/stdio-like': { mounts: 1 } }, + }) + expect(entry?.harness).toBe('test/observe') + expect(setups.count).toBe(0) + expect(writes()).toBe(0) + expect(isOpen()).toBe(false) + await host.close() + }) + }, +) + +describe('host.observe on a durable host', () => { + it('streams what open sessions write, shares one fold, and does not hold the session', async () => { + const { host, harness, persistence, seed, isOpen } = setup('durable') + const reads = vi.spyOn(persistence.stores.metadata, 'get') + const first = await host.observe(harness, { threadId: THREAD }) + const second = await host.observe(harness, { threadId: THREAD }) + const reading = new AbortController() + const replies: Array = [] + const streamed = (async () => { + const events = first.events({ signal: reading.signal }) + for await (const { event } of events) { + if (event.type === EventType.TEXT_MESSAGE_CONTENT) { + replies.push(event.delta) + } + } + })() + + await seed() + // The stream reads, and the session closed on idle all the same. + expect(isOpen()).toBe(false) + expect(hostCountsOf(host).sessions).toBe(0) + // A new session of the thread writes to the same log. + await seed() + + await vi.waitFor(() => + expect(replies).toEqual(['Hello back.', 'Hello back.']), + ) + const transcript = await second.transcript() + expect(transcript.map((message) => message.content)).toEqual([ + 'Hello.', + 'Hello back.', + 'Hello.', + 'Hello back.', + ]) + // Each fold of the log reads its checkpoint once: one fold for the two + // views and the two sessions. + const folds = reads.mock.calls.filter( + ([namespace]) => namespace === 'harness:log-checkpoint', + ) + expect(folds).toHaveLength(1) + reading.abort() + await streamed + await host.close() + }) + + it('ends the event reads of its views at host.close', async () => { + const { host, harness, seed } = setup('durable') + await seed() + const view = await host.observe(harness, { threadId: THREAD }) + const seen: Array = [] + const reading = (async () => { + for await (const entry of view.events()) seen.push(entry) + })() + // The read has the log open, and it waits for new events. + await vi.waitFor(() => expect(seen.length).toBeGreaterThan(0)) + + await host.close() + + await expect(reading).resolves.toBeUndefined() + await expect(view.transcript()).rejects.toThrow('The host is closed.') + }) +}) + +describe('host.observe on a host with no log', () => { + it('ends the events at once when the session is not open, and reads the open session', async () => { + const { host, harness, seed } = setup('no log') + await seed() + const view = await host.observe(harness, { threadId: THREAD }) + + expect(await collect(view.events())).toEqual([]) + + const session = await host.open(harness, { threadId: THREAD }) + await session.send('Again.') + expect(await firstText(view.events())).toBe('Hello back.') + await session.close() + await host.close() + }) +}) + +describe('host.observe with access', () => { + const olga = { id: 'olga', tenantId: 'acme' } + const cid = { id: 'cid', tenantId: 'acme' } + + it('refuses a stranger and a thread with no entry, and checks a child through its root', async () => { + const { host, harness, persistence, seed } = setup('no log', { + access: { identify: () => null, authorize: membersOnly() }, + }) + await seed(olga) + await persistence.stores.sessions.upsert({ + tenantId: 'acme', + threadId: 'subagent:run-1', + parentThreadId: THREAD, + harness: harness.name, + members: [], + createdAt: 1, + updatedAt: 1, + }) + const observe = (threadId: string, principal: Principal) => + host.observe(harness, { threadId, principal }) + + await expect(observe(THREAD, cid)).rejects.toBeInstanceOf(AccessDeniedError) + await expect(observe('subagent:run-1', cid)).rejects.toBeInstanceOf( + AccessDeniedError, + ) + // No entry is not found: the same answer as a refusal, and no entry. + await expect(observe('t-none', olga)).rejects.toBeInstanceOf( + AccessDeniedError, + ) + expect( + await host.sessions.get({ threadId: 't-none', principal: olga }), + ).toBeUndefined() + const child = await observe('subagent:run-1', olga) + expect(await child.transcript()).toEqual([]) + await host.close() + }) +}) + +describe('the read routes', () => { + it('mount nothing and write nothing for snapshot, transcript, and events', async () => { + const { host, harness, setups, writes, seed, reset, isOpen } = + setup('no log') + await seed() + reset() + const handler = createHarnessHandler({ host, harness, openAccess: true }) + const get = (route: string) => + handler(new Request(`${BASE}/${route}?threadId=${THREAD}`)) + + const snapshot = await (await get('snapshot')).json() + const transcript: Array<{ content: string }> = await ( + await get('transcript') + ).json() + const events = await (await get('events')).text() + + expect(snapshot.status).toBe('idle') + expect(transcript.map((message) => message.content)).toEqual([ + 'Hello.', + 'Hello back.', + ]) + // No session is open, so there is nothing to stream. + expect(events).toBe('') + expect(setups.count).toBe(0) + expect(writes()).toBe(0) + expect(isOpen()).toBe(false) + await host.close() + }) + + it('opens the session of a socket at its first input, not at subscribe or snapshot', async () => { + const { host, harness, setups, seed, reset } = setup('no log') + await seed() + reset() + const { socket, sent, closed, receive } = fakeSocket() + handleHarnessSocket({ + host, + harness, + socket, + request: new Request(`${BASE}/socket`), + openAccess: true, + }) + + receive({ type: 'harness.subscribe', threadId: THREAD }) + receive({ type: 'harness.snapshot' }) + await vi.waitFor(() => + expect(sent.map((frame) => frame.type)).toEqual([ + 'harness.hello', + 'harness.snapshot', + ]), + ) + expect(setups.count).toBe(0) + + receive({ + type: 'harness.input', + requestId: 'r-1', + input: { op: 'prompt', message: 'Again.' }, + }) + + // The stream goes on with the session that the input opened. + await vi.waitFor(() => + expect( + sent.some( + (frame) => + frame.type === 'harness.event' && + frame.event.type === EventType.TEXT_MESSAGE_CONTENT, + ), + ).toBe(true), + ) + expect(setups.count).toBe(1) + expect(closed).toEqual([]) + socket.close() + await host.close() + }) +}) diff --git a/packages/ai-mcp/src/harness.ts b/packages/ai-mcp/src/harness.ts index 224f9ac500..3ae5ae5fbb 100644 --- a/packages/ai-mcp/src/harness.ts +++ b/packages/ai-mcp/src/harness.ts @@ -34,8 +34,10 @@ import type { HarnessHost, HarnessSession, MediaRecord, + ObservedSession, Operation, Principal, + SessionSnapshot, } from '@tanstack/ai-harness' import type { OpenAccess } from '@tanstack/ai-harness/server' import type { ContentBlock } from '@modelcontextprotocol/server' @@ -188,7 +190,8 @@ const attachmentsSchema: JSONSchema = { * `harness.expose.tools`. * Every tool takes an optional `threadId`. Sessions open with * `host.open(harness, { threadId, principal })`, so the host checks each - * call. + * call. The `status` tool and the media resource only read: they use + * `host.observe`, so they open no session and mount no plugins. * * Who makes the calls: * @@ -274,6 +277,9 @@ export async function createHarnessMcpServer

( const open = (threadId: string, context: Record) => host.open(harness, { threadId, principal: principalOf(context) }) + const observe = (threadId: string, context: Record) => + host.observe(harness, { threadId, principal: principalOf(context) }) + /** * The execute of a tool that works on the session of its thread. The * session opens for the user of the call, so the host checks each action. @@ -302,6 +308,24 @@ export async function createHarnessMcpServer

( } } + /** + * The execute of a tool that only reads its thread. It reads with + * `host.observe`, so it opens no session and mounts no plugins. A refused + * read becomes a tool error that says so. + */ + const onView = + (work: (view: ObservedSession) => Promise) => + async (args: unknown, ctx: ToolExecutionContext) => { + try { + return await work( + await observe(threadOf(args, defaultThread), ctx.context), + ) + } catch (error) { + if (error instanceof AccessDeniedError) throw deniedError(error) + throw error + } + } + // Answers the approvals with the approval mode. `undefined` means that the // client must answer. One resolve answers every interrupt, and only the // client can answer an interrupt that is not an approval. @@ -345,7 +369,7 @@ export async function createHarnessMcpServer

( const outcome = await settle(session, operation, from) if (!outcome.done) { waiting.set(outcome.questionId, operation) - return { status: 'waiting', ...pending(session) } + return { status: 'waiting', ...pending(session.snapshot()) } } media.push(...(await mediaOf(operation))) if (operation.kind !== 'chat') { @@ -359,7 +383,7 @@ export async function createHarnessMcpServer

( const result = { status: operation.status(), text: turnText(outcome.value), - ...pending(session), + ...pending(session.snapshot()), } return withMedia(session, result, media) } @@ -532,7 +556,7 @@ export async function createHarnessMcpServer

( waiting.delete(questionId) // A question that no tool call waits on: nothing to continue. if (operation === undefined) { - return { status: 'answered', ...pending(session) } + return { status: 'answered', ...pending(session.snapshot()) } } return finish(session, operation, from, context) }), @@ -544,11 +568,11 @@ export async function createHarnessMcpServer

( 'Show what the harness does now: the status, the interrupts and questions that wait, the background agents, and the queued turns.', inputSchema: withThreadId(undefined), }).server( - onSession(async (session) => { - const snapshot = session.snapshot() + onView(async (view) => { + const snapshot = await view.snapshot() return { status: snapshot.status, - ...pending(session), + ...pending(snapshot), agents: snapshot.activeOperations .filter((operation) => operation.kind === 'agent') .map(({ id, agent }) => ({ id, agent })), @@ -701,17 +725,14 @@ export async function createHarnessMcpServer

( argsSchema: { parse: mediaAddress }, }).read(async (_uri, { threadId, id }, ctx) => { try { - const target = await open(threadId, ctx.context) - try { - const record = await target.getMedia(id) - if (record === null) { - throw new Error(`Media ${id} was not found in thread ${threadId}.`) - } - const blob = toBase64(await target.loadMedia({ id })) - return { blob, mimeType: record.mimeType } - } finally { - await target.close() + // A read: no session opens. + const view = await observe(threadId, ctx.context) + const record = await view.getMedia(id) + if (record === null) { + throw new Error(`Media ${id} was not found in thread ${threadId}.`) } + const blob = toBase64(await view.loadMedia({ id })) + return { blob, mimeType: record.mimeType } } catch (error) { if (error instanceof AccessDeniedError) throw deniedError(error) throw error @@ -862,8 +883,8 @@ async function questionAfter( return undefined } -function pending(session: HarnessSession) { - const snapshot = session.snapshot() +/** The interrupts and questions of `snapshot` that wait for the client. */ +function pending(snapshot: SessionSnapshot) { return { interrupts: snapshot.pendingInterrupts.map((interrupt) => ({ id: interrupt.id, From dc51be691260a81cfacdd4170775be7ec6db9c48 Mon Sep 17 00:00:00 2001 From: Alem Tuzlak Date: Sat, 17 Oct 2026 14:07:25 +0200 Subject: [PATCH 07/12] test(ai-e2e): delete purges media and messages, and reads work after idle close --- .../e2e/src/routes/api.harness-protocol.$.ts | 17 +- testing/e2e/tests/harness-lifecycle.spec.ts | 215 ++++++++++++++++++ 2 files changed, 231 insertions(+), 1 deletion(-) create mode 100644 testing/e2e/tests/harness-lifecycle.spec.ts diff --git a/testing/e2e/src/routes/api.harness-protocol.$.ts b/testing/e2e/src/routes/api.harness-protocol.$.ts index 23b5da7439..6981aac80d 100644 --- a/testing/e2e/src/routes/api.harness-protocol.$.ts +++ b/testing/e2e/src/routes/api.harness-protocol.$.ts @@ -70,6 +70,14 @@ const membersHost = createHarnessHost({ persistence: memoryPersistence(), access: { identify, authorize: membersOnly() }, }) +// A host that closes a session as soon as no handle holds it and it is idle. +// A request picks it with the `x-harness-idle: close` header, so a read +// after a turn finds a closed session. +const idleHost = createHarnessHost({ + persistence: memoryPersistence(), + access: tenantAccess, + lifecycle: { idleCloseMs: 0 }, +}) // On the durable host, an `app.signal` record that a client appends becomes // a user message. The host without a log does not read it. const signalProject: ProjectOptions = { @@ -309,9 +317,16 @@ function harnessFor(request: Request) { }) const isMembersOnly = request.headers.get('x-harness-access') === 'members' const isDurable = request.headers.get('x-harness-durable') === '1' + const isIdleClose = request.headers.get('x-harness-idle') === 'close' return { harness, - host: isMembersOnly ? membersHost : isDurable ? durableHost : host, + host: isMembersOnly + ? membersHost + : isDurable + ? durableHost + : isIdleClose + ? idleHost + : host, } } diff --git a/testing/e2e/tests/harness-lifecycle.spec.ts b/testing/e2e/tests/harness-lifecycle.spec.ts new file mode 100644 index 0000000000..0b20206a46 --- /dev/null +++ b/testing/e2e/tests/harness-lifecycle.spec.ts @@ -0,0 +1,215 @@ +import { test, expect } from './fixtures' +import { openHarnessThread } from './helpers' +import type { APIRequestContext } from '@playwright/test' +import type { MediaRecord } from '@tanstack/ai-harness' + +/** + * The life of a thread on the harness protocol: `POST sessions` with + * `op: 'delete'` purges the thread (its messages, its index entry, and its + * media files), and a read of a thread whose session closed on idle still + * works. The `x-harness-idle: close` header picks a host in the route that + * closes an idle session at once. + */ + +const BASE = '/api/harness-protocol' +const ANSWER = 'Hello over the harness protocol.' +const NOT_FOUND = { status: 404, body: { error: 'not found' } } +// A 1x1 red PNG. +const RED_DOT = Buffer.from( + 'iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAIAAACQd1PeAAAADElEQVR42mP4z8AAAAMBAQD3A0FDAAAAAElFTkSuQmCC', + 'base64', +) + +/** The routes of the protocol handler, as the e2e user. */ +function harnessApi( + request: APIRequestContext, + { + testId, + aimockPort, + extraHeaders = {}, + }: { + testId: string + aimockPort: number + extraHeaders?: Record + }, +) { + const headers = { + authorization: 'Bearer e2e-token', + 'x-test-id': testId, + 'x-aimock-port': String(aimockPort), + ...extraHeaders, + } + const json = { ...headers, 'content-type': 'application/json' } + /** A GET of a thread route, as its status and its JSON body. */ + const read = async (route: string, query: Record) => { + const response = await request.get( + `${BASE}/${route}?${new URLSearchParams(query)}`, + { headers }, + ) + return { status: response.status(), body: await response.json() } + } + return { + headers, + read, + /** Send the hello prompt of the protocol fixture, and wait for its end. */ + async send(threadId: string, runId: string) { + const response = await request.post(`${BASE}/run`, { + headers: json, + data: { + threadId, + // The run id is also the input id, so each run needs its own. + runId, + messages: [ + { id: 'u1', role: 'user', content: '[harness-protocol] hello' }, + ], + tools: [], + context: [], + state: {}, + forwardedProps: {}, + }, + }) + // The stream ends when the turn ends. + return response.text() + }, + /** The role of each saved message, and the text of each answer. */ + async transcript(threadId: string) { + const { status, body } = await read('transcript', { threadId }) + const messages: Array<{ role: string; content: unknown }> = body + return { + status, + roles: messages.map((message) => message.role), + answers: messages + .filter((message) => message.role === 'assistant') + .map((message) => message.content), + } + }, + /** Upload `bytes` as a PNG file of `threadId`. */ + async upload(threadId: string, bytes: Buffer) { + const query = new URLSearchParams({ threadId, name: 'dot.png' }) + const response = await request.post(`${BASE}/media?${query}`, { + headers: { ...headers, 'content-type': 'image/png' }, + data: bytes, + }) + const record: MediaRecord = await response.json() + return record + }, + /** A GET of the file `id` of `threadId`, with the login of the user. */ + media: (threadId: string, id: string) => + request.get(`${BASE}/media?${new URLSearchParams({ threadId, id })}`, { + headers, + }), + /** The signed URL of a file, from the media-url route. */ + async signedUrl(record: MediaRecord) { + const { body } = await read('media-url', { + threadId: record.threadId, + id: record.id, + }) + const signed: { path: string } = body + return `${BASE}/${signed.path}` + }, + /** Delete the thread with `POST sessions`. */ + async remove(threadId: string) { + const response = await request.post(`${BASE}/sessions`, { + headers: json, + data: { op: 'delete', threadId }, + }) + return response.status() + }, + } +} + +test.describe('harness lifecycle', () => { + test('a delete removes the media files of the thread, also from a signed URL', async ({ + request, + testId, + aimockPort, + }) => { + const api = harnessApi(request, { testId, aimockPort }) + const threadId = `lifecycle-media-${testId}` + const record = await api.upload(threadId, RED_DOT) + const url = await api.signedUrl(record) + + const before = await api.media(threadId, record.id) + expect(before.status()).toBe(200) + expect(await before.body()).toEqual(RED_DOT) + + expect(await api.remove(threadId)).toBe(204) + + const after = await api.media(threadId, record.id) + expect({ status: after.status(), body: await after.json() }).toEqual( + NOT_FOUND, + ) + // The signature is still good, but the thread is gone. + const signed = await request.get(url) + expect({ status: signed.status(), body: await signed.json() }).toEqual( + NOT_FOUND, + ) + + // The same thread id again: it has an entry, but not the old file. + await openHarnessThread(request, api.headers, threadId) + const reopened = await api.media(threadId, record.id) + expect({ status: reopened.status(), body: await reopened.json() }).toEqual( + NOT_FOUND, + ) + }) + + test('a delete removes the messages, and a new prompt on the same id starts an empty transcript', async ({ + request, + testId, + aimockPort, + }) => { + const api = harnessApi(request, { testId, aimockPort }) + const threadId = `lifecycle-delete-${testId}` + await api.send(threadId, `lifecycle-first-${testId}`) + expect(await api.transcript(threadId)).toEqual({ + status: 200, + roles: ['user', 'assistant'], + answers: [ANSWER], + }) + + expect(await api.remove(threadId)).toBe(204) + + // With host access, a read of a thread with no entry is 404. + expect(await api.read('transcript', { threadId })).toEqual(NOT_FOUND) + expect(await api.read('snapshot', { threadId })).toEqual(NOT_FOUND) + + await api.send(threadId, `lifecycle-again-${testId}`) + // Only the new turn: the first user message and answer are gone. + expect(await api.transcript(threadId)).toEqual({ + status: 200, + roles: ['user', 'assistant'], + answers: [ANSWER], + }) + }) + + test('a thread whose session closed on idle still answers its reads', async ({ + request, + testId, + aimockPort, + }) => { + const api = harnessApi(request, { + testId, + aimockPort, + extraHeaders: { 'x-harness-idle': 'close' }, + }) + const threadId = `lifecycle-idle-${testId}` + await api.send(threadId, `lifecycle-idle-${testId}`) + + // The turn ended, so the host closed the session. The reads load the + // thread from the store. + expect(await api.read('snapshot', { threadId })).toMatchObject({ + status: 200, + body: { + threadId, + status: 'idle', + activeOperations: [], + pendingInterrupts: [], + }, + }) + expect(await api.transcript(threadId)).toEqual({ + status: 200, + roles: ['user', 'assistant'], + answers: [ANSWER], + }) + }) +}) From 3edade71a03a7e2bb5a417bdd8631df8bbfb7759 Mon Sep 17 00:00:00 2001 From: Alem Tuzlak Date: Sat, 17 Oct 2026 14:20:58 +0200 Subject: [PATCH 08/12] feat(ai-harness, ai-mcp): host-lifetime plugins, MCP stdio servers once per host, and delete that purges a thread --- docs/harness/lifecycle.md | 3 +- docs/harness/mcp.md | 4 +- docs/harness/plugins.md | 28 +- packages/ai-harness/src/host-plugins.ts | 100 ++++++ packages/ai-harness/src/host.ts | 111 ++++-- packages/ai-harness/src/http.ts | 3 +- packages/ai-harness/src/index.ts | 5 + packages/ai-harness/src/log.ts | 3 +- packages/ai-harness/src/media.ts | 7 +- packages/ai-harness/src/plugins.ts | 315 +++++++++++------ packages/ai-harness/src/purge.ts | 62 ++++ packages/ai-harness/src/session/base.ts | 2 + .../ai-harness/src/session/harness-session.ts | 7 +- packages/ai-harness/src/session/turn.ts | 4 +- packages/ai-harness/src/session/types.ts | 3 + .../ai-harness/tests/host-plugins.test.ts | 227 +++++++++++++ packages/ai-harness/tests/purge.test.ts | 198 +++++++++++ packages/ai-harness/tests/types.test-d.ts | 2 +- packages/ai-mcp/src/harness-plugin.ts | 318 +++++++++++------- packages/ai-mcp/tests/harness-plugin.test.ts | 70 +++- 20 files changed, 1201 insertions(+), 271 deletions(-) create mode 100644 packages/ai-harness/src/host-plugins.ts create mode 100644 packages/ai-harness/src/purge.ts create mode 100644 packages/ai-harness/tests/host-plugins.test.ts create mode 100644 packages/ai-harness/tests/purge.test.ts diff --git a/docs/harness/lifecycle.md b/docs/harness/lifecycle.md index 7b60a0f609..cde5397ac4 100644 --- a/docs/harness/lifecycle.md +++ b/docs/harness/lifecycle.md @@ -168,12 +168,13 @@ First, `delete` closes the sessions of the thread and of its child threads, for - The index entry. - The uploaded media and the generated artifacts. - The thread settings, the plugin config, and the interrupted turn. -- The usage totals and the revert data. +- The usage totals, the revert data, and the log checkpoint. These stay in their stores: plugin state, run records, interrupts, the inbox, and work claims. To remove them, delete them in your own stores. - 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. - `host.events()` yields a `session-deleted` event. See [Show a live status](./sessions#show-a-live-status). ## What you have now diff --git a/docs/harness/mcp.md b/docs/harness/mcp.md index 7cae04bc5d..4cc1700887 100644 --- a/docs/harness/mcp.md +++ b/docs/harness/mcp.md @@ -63,8 +63,8 @@ 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. -- `stdio`: the server runs as a child process on this machine (Node only). The host starts it one time, and every session of the host 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`. +- `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. - `oauth: true`: the user signs in with `/connect linear`. Each user connects with their own sign-in. See [Sign in to an OAuth server](#3-sign-in-to-an-oauth-server). ## 2. Check the servers diff --git a/docs/harness/plugins.md b/docs/harness/plugins.md index bc1befac84..7e22ba4f62 100644 --- a/docs/harness/plugins.md +++ b/docs/harness/plugins.md @@ -437,9 +437,35 @@ export const orders = definePlugin({ - The host sets the plugin up one time for each harness. Every session of that harness shares this setup. - The host cleans it up when the last session that uses it closes, or at `host.close()`. -- Its setup context has no session, no credentials, no keys, and no commands. +- Its setup context has no session, no credentials, no keys, and no commands. It has `harness`, the name of the harness, and `resources`. - A tool of the plugin reads the user of its call from `ctx.context.run`. See [Act as the user of a run](#act-as-the-user-of-a-run). +Some parts need the session: a command, a setting, plugin state, or a question to the user. Return them from `session`. It runs for each session, with the same context as a session plugin: + +```ts group=harness-plugins-host +import { defineCommand } from '@tanstack/ai-harness' + +export const pool = definePlugin({ + name: 'acme/pool', + lifetime: 'host', + setup: async (ctx) => { + const db = await ctx.resources.acquire(connect, (shared) => shared.end()) + return { + session: (session) => ({ + commands: { + orders: defineCommand({ + description: 'Count the orders', + run: async () => `${await db.countOrders()} orders in ${session.session.threadId}`, + }), + }, + }), + } + }, +}) +``` + +The session gets the shared contributions and the `session` contributions together. The resources of `session` close with the session. + The [MCP plugin](./mcp) works this way for `stdio` servers: one process serves every session of the host. ## Reload plugins diff --git a/packages/ai-harness/src/host-plugins.ts b/packages/ai-harness/src/host-plugins.ts new file mode 100644 index 0000000000..e837fb759b --- /dev/null +++ b/packages/ai-harness/src/host-plugins.ts @@ -0,0 +1,100 @@ +import { ResourceScope, disposeAll } from './resources' +import type { + HostPluginContext, + HostPluginContributions, + HostPluginDefinition, +} from './plugins' + +/** One shared setup of a host plugin, and how to give it back. */ +export interface SharedHostPlugin { + contributions: HostPluginContributions | void + /** Give the setup back. The last release disposes it. Safe to call twice. */ + release: () => Promise +} + +interface Entry { + /** The sessions that hold the setup, and the ones that wait for it. */ + refs: number + scope: ResourceScope + ready: Promise | void> +} + +/** The context of a host plugin setup whose resources live in `scope`. */ +export function hostPluginContext( + harness: string, + scope: ResourceScope, +): HostPluginContext { + return { + harness, + resources: { + acquire: (open, close) => scope.acquire(open, close), + signal: scope.signal, + }, + } +} + +/** + * @internal The plugins with `lifetime: 'host'` of one host. Each one is + * set up once for each harness name and plugin name, and shared by a + * reference count. The last release disposes it, and so does `close()`. + */ +export class HostPlugins { + private readonly entries = new Map() + private closed = false + + /** The shared setup of `plugin` for the harness `harness`. */ + async acquire( + harness: string, + plugin: HostPluginDefinition, + ): Promise { + if (this.closed) { + throw new Error('The host is closed. Create a new host to open sessions.') + } + const key = `${harness}\u0000${plugin.name}` + let entry = this.entries.get(key) + if (!entry) { + const scope = new ResourceScope() + const ready = Promise.resolve().then(() => + plugin.setup(hostPluginContext(harness, scope)), + ) + const created: Entry = { refs: 0, scope, ready } + // A failed setup is not kept: the next session sets it up again. + ready.catch(() => { + if (this.entries.get(key) === created) this.entries.delete(key) + void scope.dispose().catch(() => {}) + }) + this.entries.set(key, created) + entry = created + } + const held = entry + held.refs += 1 + let contributions: HostPluginContributions | void + try { + contributions = await held.ready + } catch (error) { + held.refs -= 1 + throw error + } + let released = false + return { + contributions, + release: () => { + if (released) return Promise.resolve() + released = true + held.refs -= 1 + const isLast = held.refs <= 0 && this.entries.get(key) === held + if (!isLast) return Promise.resolve() + this.entries.delete(key) + return held.scope.dispose() + }, + } + } + + /** Dispose every shared setup. Later `acquire` calls throw. */ + close(): Promise { + this.closed = true + const scopes = [...this.entries.values()].map((entry) => entry.scope) + this.entries.clear() + return disposeAll(scopes) + } +} diff --git a/packages/ai-harness/src/host.ts b/packages/ai-harness/src/host.ts index 2456913f57..0c0a9fc1bd 100644 --- a/packages/ai-harness/src/host.ts +++ b/packages/ai-harness/src/host.ts @@ -4,6 +4,8 @@ import { memoryPersistence } from '@tanstack/ai-persistence' import type { CredentialStore } from '@tanstack/ai/byok' import { AccessDeniedError, membersOnlyPolicyOf } from './access' import { SharedLog, loadLogState, logIdOf, parseLogId } from './log' +import { HostPlugins } from './host-plugins' +import { purgeThreadData } from './purge' import { LEASE } from './resume' import { SessionImpl } from './session/harness-session' import { sessionHandle } from './session/handle' @@ -229,7 +231,8 @@ export interface OpenSessionOptions

{ * The session index of a host: one entry per tenant and thread, for a list * of sessions. The host writes the entry when a session opens and after each * turn. Without `stores.sessions`, `list` is empty, `get` is `undefined`, - * and `rename` and `delete` write nothing. + * and `rename` writes nothing. `delete` still deletes the data of the + * thread in the other stores. * * Each method takes a `principal`. Its `tenantId` picks the tenant (without * it, the default tenant), as in `host.open`. A method never reads or @@ -306,10 +309,22 @@ export interface HostSessions

{ principal?: P }) => Promise> | undefined> /** - * Remove the index entry of a thread. After it, `get` is `undefined`, and - * the next open of the thread writes a new entry. Only the entry goes: the - * transcript, the log, and the other data of the thread stay in their - * stores. + * Delete a thread and its child threads (the entries with this thread as + * `parentThreadId`, and their children). For each thread, it: + * + * - closes its open sessions on this host, of every harness, for every + * handle. A handle that calls again opens an empty thread. + * - deletes its log, or its messages on a host with no `stores.log`. + * - deletes its media: the uploads and the generated files, records and + * bytes. + * - deletes its thread settings, plugin config, usage, interrupted turn, + * revert state, and log checkpoint from `stores.metadata`. + * - removes its index entry. Then `get` is `undefined`, and the next open + * of the thread writes a new entry. + * + * A store the host does not have is skipped. Plugin state, runs, + * interrupts, the inbox, and work claims stay. Only the tenant of + * `principal` is touched. * * @example * ```ts @@ -1124,6 +1139,8 @@ export function createHarnessHost

( return loaded } const sessions = new Map>() + /** 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. */ const live = new Map() /** @@ -1365,6 +1382,7 @@ export function createHarnessHost

( hostId, index, retention, + hostPlugins, ...(logStore ? { log: { @@ -1670,22 +1688,77 @@ export function createHarnessHost

( }) } + /** The session keys of a thread of a tenant on this host, of any harness. */ + const keysOf = (tenantId: string | undefined, threadId: string) => { + const suffix = `\u0000${logIdOf({ tenantId, logId: threadId })}` + return [...sessions.keys()].filter((key) => key.endsWith(suffix)) + } + /** The open sessions of a thread of a tenant on this host, of any harness. */ const openSessionsOf = async ( tenantId: string | undefined, threadId: string, ) => { - const suffix = `\u0000${logIdOf({ tenantId, logId: threadId })}` const opened = await Promise.allSettled( - [...sessions].flatMap(([key, session]) => - key.endsWith(suffix) ? [session] : [], - ), + keysOf(tenantId, threadId).flatMap((key) => sessions.get(key) ?? []), ) return opened.flatMap((entry) => entry.status === 'fulfilled' ? [entry.value] : [], ) } + /** The thread ids of the index entries with `threadId` as parent. */ + const childrenOf = async (tenantId: string | undefined, threadId: string) => { + if (!sessionStore) return [] + const children: Array = [] + let cursor: string | undefined + do { + const page = await sessionStore.list({ + tenantId, + parentThreadId: threadId, + ...(cursor !== undefined ? { cursor } : {}), + }) + children.push(...page.entries.map((entry) => entry.threadId)) + cursor = page.truncated ? page.cursor : undefined + } while (cursor !== undefined) + return children + } + + /** + * `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. + */ + const purge = async ( + tenantId: string | undefined, + threadId: string, + seen: Set, + ): Promise => { + if (seen.has(threadId)) return + seen.add(threadId) + // First, so no session of the thread writes or starts a child while + // the purge runs. For every handle. + const keys = keysOf(tenantId, threadId) + for (const key of keys) { + const session = await sessions.get(key)?.catch(() => undefined) + if (session) await closeShared(key, session) + } + const children = await childrenOf(tenantId, threadId) + for (const child of children) await purge(tenantId, child, seen) + const storeId = logIdOf({ tenantId, logId: threadId }) + // A `host.observe` view can still hold the shared log of the thread. + // Its reads end, and the next open folds the log from the store. + resolved.get(storeId)?.log.stopReaders() + logs.delete(storeId) + resolved.delete(storeId) + await purgeThreadData({ + storeId, + stores: persistence.stores, + media: media.stores, + }) + await index.remove({ tenantId, threadId }) + } + /** * `host.fork` after its checks: the session class of the new thread, and * its key, which the fork holds. Give the hold to a handle, or release it. @@ -1970,6 +2043,8 @@ export function createHarnessHost

( .filter((entry) => entry.status === 'fulfilled') .map((entry) => entry.value.close(closeOptions)), ) + // After the sessions: each one gave its host plugins back. + await hostPlugins.close() }, logState: ({ harness, logId, principal }) => { if (access) need(principal, 'host.logState') @@ -2052,25 +2127,11 @@ export function createHarnessHost

( )) as SessionIndexEntry> | undefined }, async delete({ threadId, principal }) { - const tenantId = principal?.tenantId - const entry = await checkEntry('host.sessions.delete', 'delete', { + await checkEntry('host.sessions.delete', 'delete', { threadId, principal, }) - // The transcript goes with the entry, so the next opener of this - // thread id cannot read it. ponytail: the log or the messages only; - // PR 4 adds a purge of blobs, metadata, and children. - if (entry?.harness !== undefined) { - const key = sessionKey(entry.harness, tenantId, threadId) - const opening = sessions.get(key) - // For every handle. A handle that calls again opens a new thread. - if (opening) await closeShared(key, await opening) - } - const storeId = logIdOf({ tenantId, logId: threadId }) - const { log, messages } = persistence.stores - if (log) await log.deleteLog({ logId: storeId }) - else await messages?.saveThread(storeId, []) - await index.remove({ tenantId, threadId }) + await purge(principal?.tenantId, threadId, new Set()) }, async fork({ harness, threadId, at, members, principal }) { const tenantId = principal?.tenantId diff --git a/packages/ai-harness/src/http.ts b/packages/ai-harness/src/http.ts index dffc9ad975..b0fbdb5e4e 100644 --- a/packages/ai-harness/src/http.ts +++ b/packages/ai-harness/src/http.ts @@ -448,7 +448,8 @@ function watchRevoke

( * pinned?, metadata? }` (each field replaces the stored one, and the * harness `metadataSchema` checks `metadata`. Without a schema, a client * cannot write `metadata`: 403 `not_exposed`), `{ op: 'delete', threadId }` - * (it removes the index entry only), `{ op: 'fork', threadId, before }` + * (it deletes the thread and its child threads, as + * `host.sessions.delete` does), `{ op: 'fork', threadId, before }` * or `{ op: 'fork', threadId, through }` with a message id, and with * `membersOnly()`, `{ op: 'share', threadId, userId, role }` and * `{ op: 'unshare', threadId, userId }`. Rename, update, fork, share, and diff --git a/packages/ai-harness/src/index.ts b/packages/ai-harness/src/index.ts index 790c1c5829..5d4ebaf3d6 100644 --- a/packages/ai-harness/src/index.ts +++ b/packages/ai-harness/src/index.ts @@ -39,8 +39,12 @@ export type { export { definePlugin } from './plugins' export type { HarnessPlugin, + HostPluginContext, + HostPluginContributions, + HostPluginDefinition, PluginContributions, PluginDefinition, + PluginResources, PluginLifetime, PluginPrompt, AgentGroup, @@ -48,6 +52,7 @@ export type { PluginCommands, PluginSetupContext, PluginState, + SessionPluginDefinition, } from './plugins' export { createExtensionPoint, createPluginEvent } from './extensions' diff --git a/packages/ai-harness/src/log.ts b/packages/ai-harness/src/log.ts index 87e4743778..b1d38a52b2 100644 --- a/packages/ai-harness/src/log.ts +++ b/packages/ai-harness/src/log.ts @@ -617,7 +617,8 @@ export function foldLogEntry( ) } -const CHECKPOINT_NAMESPACE = 'harness:log-checkpoint' +/** @internal The metadata namespace of log checkpoints, keyed by store log id. */ +export const CHECKPOINT_NAMESPACE = 'harness:log-checkpoint' // ponytail: a checkpoint every 50 appends. Tune it when a cold fold is slow. const CHECKPOINT_EVERY = 50 const READ_PAGE = 500 diff --git a/packages/ai-harness/src/media.ts b/packages/ai-harness/src/media.ts index b3f26f2633..620897f39e 100644 --- a/packages/ai-harness/src/media.ts +++ b/packages/ai-harness/src/media.ts @@ -85,6 +85,11 @@ function uploadRunId(storeThreadId: string) { return `upload:${storeThreadId}` } +/** @internal The key prefix of the upload bytes of a thread in the blob store. */ +export function uploadBlobPrefix(storeThreadId: string) { + return `artifacts/${uploadRunId(storeThreadId)}/` +} + /** The media record of `record`, an artifact of the thread `threadId`. */ function toMediaRecord( record: ArtifactRecord, @@ -230,7 +235,7 @@ export function createMediaStore({ const id = crypto.randomUUID() const runId = uploadRunId(storeThreadId) - const blobKey = `artifacts/${runId}/${id}` + const blobKey = `${uploadBlobPrefix(storeThreadId)}${id}` // A stream has no length up front, so count it as the store reads it. let seen = 0 const counted = diff --git a/packages/ai-harness/src/plugins.ts b/packages/ai-harness/src/plugins.ts index e296d1bd9f..7f5cf19f68 100644 --- a/packages/ai-harness/src/plugins.ts +++ b/packages/ai-harness/src/plugins.ts @@ -1,4 +1,5 @@ import { CapabilityRegistry, createCapability } from '@tanstack/ai' +import { hostPluginContext } from './host-plugins' import { ResourceScope, disposeAll } from './resources' import { assertMetadataSchema } from './session-entry' import type { @@ -28,14 +29,18 @@ import type { AnyCommand, PluginSessionApi } from './commands' import type { ConfigOption } from './config' import type { ExtensionItem, ExtensionPoint, PluginEvent } from './extensions' import type { MetadataOf } from './session-entry' +import type { SharedHostPlugin } from './host-plugins' /** * How long a plugin's `setup` result and resources live: * - `session` (default): from session open to session close. Resources such * as an index, a watcher, or an LSP server stay up across turns. * - `turn`: set up for each chat turn, disposed when that turn ends. + * - `host`: set up once for each harness on a host, and shared by every + * session of that harness. Disposed when the last session that uses it + * closes, or at `host.close()`. See {@link HostPluginContext}. */ -export type PluginLifetime = 'session' | 'turn' +export type PluginLifetime = 'session' | 'turn' | 'host' /** * A prompt a plugin contributes. `id` must be unique in the session. A @@ -112,6 +117,53 @@ export interface PluginContributions { }) => ReadonlyArray | Promise> } +/** + * What the `setup` of a plugin with `lifetime: 'host'` returns. The host + * mounts it into each session that uses the plugin. It cannot add commands + * or config: return them from `session`. + */ +export interface HostPluginContributions< + TMetadata = Record, +> extends Omit { + /** + * Runs for each session, after the shared setup, with the context of a + * session plugin. Use it for the parts that need the session: commands, + * config, state, and questions to the user. Its resources close with the + * session. What it returns is added to the shared contributions. + */ + session?: ( + ctx: PluginSetupContext, + ) => PluginContributions | void | Promise +} + +/** Resources a plugin owns, with rollback and reverse-order cleanup. */ +export interface PluginResources { + /** + * Open a resource. `close` runs when the plugin is cleaned up, newest + * first. It is registered only after `open` resolves. + */ + acquire: ( + open: () => T | Promise, + close: (resource: T) => unknown, + ) => Promise + /** Aborts when the plugin is cleaned up. */ + signal: AbortSignal +} + +/** + * What the `setup` of a plugin with `lifetime: 'host'` receives. The setup + * runs once for each harness on a host, so it has no session, no user, no + * credentials, no keys, and no commands. Open the shared resource here (a + * pool, a child process, a watcher). A tool reads the user of its call from + * `ctx.context.run`. + */ +export interface HostPluginContext { + /** The name of the harness that the plugin is set up for. */ + harness: string + /** Resources that live as long as the shared setup. */ + resources: PluginResources +} + /** Plugin state that survives restarts, stored in the metadata store. */ export interface PluginState { get: () => Promise @@ -229,13 +281,7 @@ export interface PluginCommands { */ export interface PluginSetupContext> { /** Resources this plugin owns, with rollback and reverse-order cleanup. */ - resources: { - acquire: ( - open: () => T | Promise, - close: (resource: T) => unknown, - ) => Promise - signal: AbortSignal - } + resources: PluginResources /** Read a capability another plugin or middleware provided. Throws if absent. */ get: (capability: Capability) => T /** Read a capability, or `undefined` when nobody provided it. */ @@ -313,11 +359,8 @@ export interface PluginSetupContext> { turn?: TurnInfo } -/** - * A plugin. `TMetadata` is the type of the session index `metadata` that the - * plugin reads and writes, from its `metadataSchema`. - */ -export interface PluginDefinition> { +/** The fields of every plugin definition. */ +interface PluginDefinitionBase { /** A stable, unique name, for example `'acme/todos'`. */ name: string /** @@ -328,7 +371,6 @@ export interface PluginDefinition> { * metadata type does not fit is a type error. */ metadataSchema?: SchemaInput - lifetime?: PluginLifetime /** Capabilities that an earlier plugin or harness middleware must provide. */ requires?: ReadonlyArray /** Capabilities this plugin provides in `setup`. */ @@ -337,11 +379,42 @@ export interface PluginDefinition> { optionalRequires?: ReadonlyArray /** Agent outputs this plugin needs the session to have. */ needs?: { produces?: ReadonlyArray } +} + +/** A plugin that is set up for each session (default) or for each turn. */ +export interface SessionPluginDefinition< + TMetadata = Record, +> extends PluginDefinitionBase { + lifetime?: 'session' | 'turn' setup?: ( ctx: PluginSetupContext, ) => PluginContributions | void | Promise } +/** + * A plugin that is set up once for each harness on a host, and shared by + * every session of that harness. See {@link HostPluginContext}. + */ +export interface HostPluginDefinition< + TMetadata = Record, +> extends PluginDefinitionBase { + lifetime: 'host' + setup: ( + ctx: HostPluginContext, + ) => + | HostPluginContributions + | void + | Promise | void> +} + +/** + * A plugin. `TMetadata` is the type of the session index `metadata` that the + * plugin reads and writes, from its `metadataSchema`. + */ +export type PluginDefinition> = + | SessionPluginDefinition + | HostPluginDefinition + const PLUGIN_BRAND = Symbol.for('tanstack.ai.harnessPlugin') /** @@ -349,7 +422,9 @@ const PLUGIN_BRAND = Symbol.for('tanstack.ai.harnessPlugin') * be ignored. Other extra keys are the plugin's own API, such as the `diff` * of `snapshots()`. */ -const SETUP_KEYS: ReadonlySet = new Set([ +const SETUP_KEYS: ReadonlySet = new Set< + keyof PluginContributions | keyof HostPluginContributions +>([ 'adapter', 'tools', 'prompts', @@ -363,11 +438,15 @@ const SETUP_KEYS: ReadonlySet = new Set([ 'contribute', 'discoverTools', 'prepareTools', + 'session', ]) /** No `setup()` contribution at the top level of a definition. */ type NoSetupKeys = { - [K in Exclude]?: never + [K in Exclude< + keyof HostPluginContributions | keyof PluginContributions, + keyof PluginDefinitionBase | 'lifetime' | 'setup' + >]?: never } /** @@ -426,7 +505,10 @@ export function definePlugin< `definePlugin "${definition.name}"`, definition.metadataSchema, ) - return Object.freeze({ ...definition, [PLUGIN_BRAND]: true as const }) + // `assign`, not a spread: a spread of a union type parameter loses it. + return Object.freeze( + Object.assign({ [PLUGIN_BRAND]: true as const }, definition), + ) } /** A tool, prompt, or middleware with the plugin that contributed it. */ @@ -509,6 +591,14 @@ export interface MountEnvironment { services?: PluginServices /** The turn of a run mount. Plugins read it as `ctx.turn`. */ turn?: TurnInfo + /** + * The setup of a plugin with `lifetime: 'host'` that the host shares. The + * mount gives it back when it is disposed. Without it, a host plugin is + * set up for this mount only. + */ + hostPlugin?: (plugin: HostPluginDefinition) => Promise + /** The harness name a host plugin gets without `hostPlugin`. */ + harnessName?: string } const unavailable = (what: string) => () => { @@ -708,6 +798,23 @@ function promptOf( : prompt } +/** + * The shared setup of a host plugin. Without a host, the plugin is set up + * in `scope`, for this mount only. + */ +async function sharedSetup( + plugin: HostPluginDefinition, + scope: ResourceScope, + env: MountEnvironment, +): Promise { + if (env.hostPlugin) return env.hostPlugin(plugin) + const context = hostPluginContext(env.harnessName ?? '', scope) + return { + contributions: await plugin.setup(context), + release: () => Promise.resolve(), + } +} + /** A turn plugin added a command or a config option. */ const turnOnly = (plugin: string) => new Error( @@ -776,7 +883,7 @@ export async function mountPlugins( ...(plugin.requires ?? []), ...(plugin.optionalRequires ?? []), ]) - const contributions = await plugin.setup?.({ + const context: PluginSetupContext = { resources: { acquire: (open, close) => scope.acquire(open, close), signal: scope.signal, @@ -883,7 +990,19 @@ export async function mountPlugins( ready, }, ...(env.turn ? { turn: env.turn } : {}), - }) + } + let parts: Array + if (plugin.lifetime === 'host') { + // The shared setup goes back when this mount is disposed. + const shared = await scope.acquire( + () => sharedSetup(plugin, scope, env), + (setup) => setup.release(), + ) + const own = await shared.contributions?.session?.(context) + parts = [shared.contributions, own] + } else { + parts = [await plugin.setup?.(context)] + } for (const handle of plugin.provides ?? []) { if (!provided.has(handle)) { throw new Error( @@ -891,88 +1010,92 @@ export async function mountPlugins( ) } } - if (!contributions) continue - for (const tool of contributions.tools ?? []) { - const clash = tools.find((entry) => entry.value.name === tool.name) - if (clash) { - throw new Error( - `Duplicate tool "${tool.name}": first owner ${clash.owner}, second owner ${plugin.name}. No run started.`, - ) + for (const contributions of parts) { + if (!contributions) continue + for (const tool of contributions.tools ?? []) { + const clash = tools.find((entry) => entry.value.name === tool.name) + if (clash) { + throw new Error( + `Duplicate tool "${tool.name}": first owner ${clash.owner}, second owner ${plugin.name}. No run started.`, + ) + } + tools.push({ value: tool, owner: plugin.name }) } - tools.push({ value: tool, owner: plugin.name }) - } - ;(contributions.prompts ?? []).forEach((prompt, index) => { - const section = promptOf(prompt, plugin.name, index) - const clash = prompts.find((entry) => entry.value.id === section.id) - if (clash) { - throw new Error( - `Duplicate prompt id "${section.id}": first owner ${clash.owner}, second owner ${plugin.name}.`, - ) + ;(contributions.prompts ?? []).forEach((prompt, index) => { + const section = promptOf(prompt, plugin.name, index) + const clash = prompts.find((entry) => entry.value.id === section.id) + if (clash) { + throw new Error( + `Duplicate prompt id "${section.id}": first owner ${clash.owner}, second owner ${plugin.name}.`, + ) + } + prompts.push({ value: section, owner: plugin.name }) + }) + for (const item of contributions.middleware ?? []) { + middleware.push(item) + middlewareOwners.push(plugin.name) } - 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 ?? [])) - // The session reads commands and config only from its session plugins. - const isSessionOnly = - Object.keys(contributions.commands ?? {}).length > 0 || - Object.keys(contributions.config ?? {}).length > 0 - if (env.turn && isSessionOnly) throw turnOnly(plugin.name) - for (const [name, command] of Object.entries( - contributions.commands ?? {}, - )) { - const clash = commands.get(name) ?? env.takenCommands?.get(name) - if (clash) { - throw new Error( - `Duplicate command "${name}": first owner ${clash.owner}, second owner ${plugin.name}.`, - ) + agentMiddleware.push(...(contributions.agentMiddleware ?? [])) + generationMiddleware.push(...(contributions.generationMiddleware ?? [])) + // The session reads commands and config only from its session plugins. + const isSessionOnly = + Object.keys(contributions.commands ?? {}).length > 0 || + Object.keys(contributions.config ?? {}).length > 0 + if (env.turn && isSessionOnly) throw turnOnly(plugin.name) + for (const [name, command] of Object.entries( + contributions.commands ?? {}, + )) { + const clash = commands.get(name) ?? env.takenCommands?.get(name) + if (clash) { + throw new Error( + `Duplicate command "${name}": first owner ${clash.owner}, second owner ${plugin.name}.`, + ) + } + commands.set(name, { command, owner: plugin.name }) } - commands.set(name, { command, owner: plugin.name }) - } - for (const [key, option] of Object.entries(contributions.config ?? {})) { - const clash = config.get(key) ?? env.takenConfig?.get(key) - if (clash) { - throw new Error( - `Duplicate config key "${key}": first owner ${clash.owner}, second owner ${plugin.name}.`, - ) + for (const [key, option] of Object.entries( + contributions.config ?? {}, + )) { + const clash = config.get(key) ?? env.takenConfig?.get(key) + if (clash) { + throw new Error( + `Duplicate config key "${key}": first owner ${clash.owner}, second owner ${plugin.name}.`, + ) + } + config.set(key, { option, owner: plugin.name }) } - config.set(key, { option, owner: plugin.name }) - } - for (const item of contributions.contribute ?? []) { - let items = extensions.get(item.point) - if (!items) { - items = [] - extensions.set(item.point, items) + for (const item of contributions.contribute ?? []) { + let items = extensions.get(item.point) + if (!items) { + items = [] + extensions.set(item.point, items) + } + items.push({ value: item.value, owner: plugin.name }) + } + if (contributions.adapter) { + adapters.push(contributions.adapter) + adapterOwners.push(plugin.name) + } + if (contributions.discoverTools) { + discoverers.push({ + discover: contributions.discoverTools, + owner: plugin.name, + }) + } + if (contributions.prepareTools) { + preparers.push({ + prepare: contributions.prepareTools, + owner: plugin.name, + }) + } + for (const agent of contributions.agents ?? []) { + env.registry.add(agent, plugin.name) + agents.push(agent) + } + for (const agent of contributions.subagents ?? []) { + env.registry.add(agent, plugin.name) + subagents.push(agent) } - items.push({ value: item.value, owner: plugin.name }) - } - if (contributions.adapter) { - adapters.push(contributions.adapter) - adapterOwners.push(plugin.name) - } - if (contributions.discoverTools) { - discoverers.push({ - discover: contributions.discoverTools, - owner: plugin.name, - }) - } - if (contributions.prepareTools) { - preparers.push({ - prepare: contributions.prepareTools, - owner: plugin.name, - }) - } - for (const agent of contributions.agents ?? []) { - env.registry.add(agent, plugin.name) - agents.push(agent) - } - for (const agent of contributions.subagents ?? []) { - env.registry.add(agent, plugin.name) - subagents.push(agent) } } for (const plugin of plugins) { diff --git a/packages/ai-harness/src/purge.ts b/packages/ai-harness/src/purge.ts new file mode 100644 index 0000000000..d0ad2e76cf --- /dev/null +++ b/packages/ai-harness/src/purge.ts @@ -0,0 +1,62 @@ +import { resolveArtifactBlobKey } from '@tanstack/ai-persistence' +import { CHECKPOINT_NAMESPACE } from './log' +import { uploadBlobPrefix } from './media' +import { INTERRUPTED, SETTINGS, USAGE } from './session/helpers' +import type { ArtifactStore, BlobStore } from '@tanstack/ai-persistence' +import type { HarnessPersistence } from './host' + +/** The metadata rows the harness keeps for a thread, keyed by its store key. */ +const THREAD_METADATA = [ + SETTINGS, + 'harness:config', + INTERRUPTED, + USAGE, + 'harness:revert', +] + +/** + * Delete the stored data of one thread: its log (or its messages on a host + * with no log), its media, and the harness metadata rows. `storeId` is the + * store key of the thread, from `logIdOf`. It is also the store id of the + * thread's own log. A store the persistence does not have is skipped. + */ +// ponytail: plugin state, runs, interrupts, the inbox, and work claims stay. +// PR 5 moves them into the log or removes them, so `deleteLog` covers them. +export async function purgeThreadData({ + storeId, + stores, + media, +}: { + storeId: string + stores: HarnessPersistence['stores'] + /** The media stores of the host, which can be memory stores. */ + media: { artifacts: ArtifactStore; blobs: BlobStore } +}) { + const { log, messages, metadata } = stores + if (log) await log.deleteLog({ logId: storeId }) + else await messages?.saveThread(storeId, []) + + // Uploads and generated media are both artifacts of the thread. + const records = await media.artifacts.listForThread(storeId) + for (const record of records) { + await media.blobs.delete(resolveArtifactBlobKey(record)) + await media.artifacts.delete(record.artifactId) + } + // Upload bytes with no record: a put that failed after the bytes. + const prefix = uploadBlobPrefix(storeId) + let cursor: string | undefined + do { + const page = await media.blobs.list({ + prefix, + ...(cursor !== undefined ? { cursor } : {}), + }) + for (const object of page.objects) await media.blobs.delete(object.key) + cursor = page.truncated ? page.cursor : undefined + } while (cursor !== undefined) + + if (!metadata) return + for (const namespace of THREAD_METADATA) { + await metadata.delete(namespace, storeId) + } + await metadata.delete(CHECKPOINT_NAMESPACE, storeId) +} diff --git a/packages/ai-harness/src/session/base.ts b/packages/ai-harness/src/session/base.ts index 96454b2a2c..f5fb6d9a53 100644 --- a/packages/ai-harness/src/session/base.ts +++ b/packages/ai-harness/src/session/base.ts @@ -289,6 +289,7 @@ export abstract class SessionBase { protected readonly log: SessionDependencies['log'] /** How much this session keeps in memory. */ protected readonly retention: SessionDependencies['retention'] + protected readonly hostPlugins: SessionDependencies['hostPlugins'] /** What this session keeps for its inputs. */ protected readonly inputRecords: InputRecords protected readonly agentRegistry = new AgentRegistry() @@ -487,6 +488,7 @@ export abstract class SessionBase { }) this.principal = deps.principal this.retention = deps.retention + this.hostPlugins = deps.hostPlugins this.inputRecords = new InputRecords(deps.retention.settledOperations) this.notifyHost = deps.onEvent this.feed = tapFeed( diff --git a/packages/ai-harness/src/session/harness-session.ts b/packages/ai-harness/src/session/harness-session.ts index a26150d8b9..4690d46930 100644 --- a/packages/ai-harness/src/session/harness-session.ts +++ b/packages/ai-harness/src/session/harness-session.ts @@ -160,11 +160,12 @@ export class SessionImpl< protected async mountSessionPlugins(): Promise { this.plugins = this.harness.plugins?.() ?? [] this.sessionPlugins = await mountPlugins( - this.plugins.filter( - (plugin) => (plugin.lifetime ?? 'session') === 'session', - ), + // The session plugins and the host plugins. A turn mounts the rest. + this.plugins.filter((plugin) => plugin.lifetime !== 'turn'), { threadId: this.threadId, + hostPlugin: (plugin) => + this.hostPlugins.acquire(this.harness.name, plugin), registry: this.agentRegistry, harnessTools: this.harness.tools ?? [], harnessProvides: (this.harness.middleware ?? []).flatMap( diff --git a/packages/ai-harness/src/session/turn.ts b/packages/ai-harness/src/session/turn.ts index 581a290772..a40996b2c3 100644 --- a/packages/ai-harness/src/session/turn.ts +++ b/packages/ai-harness/src/session/turn.ts @@ -410,9 +410,7 @@ export abstract class SessionTurn< const removed = new Set(stored.plugins?.remove ?? []) // The 'turn' plugins. A removed turn plugin is not set up. const perRun = this.plugins.filter( - (plugin) => - (plugin.lifetime ?? 'session') !== 'session' && - !removed.has(plugin.name), + (plugin) => plugin.lifetime === 'turn' && !removed.has(plugin.name), ) if (perRun.length > 0) { runPlugins = await mountPlugins(perRun, { diff --git a/packages/ai-harness/src/session/types.ts b/packages/ai-harness/src/session/types.ts index 9fca9982a5..21adb6f562 100644 --- a/packages/ai-harness/src/session/types.ts +++ b/packages/ai-harness/src/session/types.ts @@ -26,6 +26,7 @@ import type { AgentInputOf, AgentResultOf, AnyAgent } from '../agents' import type { AnyHarness, HarnessAgentsOf } from '../define' import type { HarnessPersistence, Retention, SessionIndexWriter } from '../host' import type { PluginLifetime } from '../plugins' +import type { HostPlugins } from '../host-plugins' import type { AgentRun, ChatTurnResult, @@ -223,6 +224,8 @@ export interface SessionDependencies { index: SessionIndexWriter /** How much the session keeps in memory. */ retention: Retention + /** The plugins with `lifetime: 'host'` that the host shares. */ + hostPlugins: HostPlugins /** The session log of a durable host. */ log?: { store: LogStore diff --git a/packages/ai-harness/tests/host-plugins.test.ts b/packages/ai-harness/tests/host-plugins.test.ts new file mode 100644 index 0000000000..819fe82037 --- /dev/null +++ b/packages/ai-harness/tests/host-plugins.test.ts @@ -0,0 +1,227 @@ +import { describe, expect, it } from 'vitest' +import { toolDefinition } from '@tanstack/ai' +import { memoryPersistence } from '@tanstack/ai-persistence' +import { + createHarnessHost, + defineCommand, + defineHarness, + definePlugin, +} from '../src' +import { mockAdapter, text, toolCall } from './helpers' +import type { HarnessPlugin, HarnessRunContext } from '../src' +import type { Reply } from './helpers' + +/** Calls `who` once, then answers after the tool result. */ +const callWho: Reply = (options) => + options.messages.at(-1)?.role === 'tool' ? text('done') : toolCall('who', {}) + +/** A host plugin that counts its setups and its cleanups. */ +function countedPlugin() { + const counts = { setups: 0, disposed: 0 } + const seen: Array = [] + const plugin = definePlugin({ + name: 'test/shared', + lifetime: 'host', + setup: async (ctx) => { + counts.setups += 1 + await ctx.resources.acquire( + () => 'pool', + () => { + counts.disposed += 1 + }, + ) + return { + tools: [ + toolDefinition({ + name: 'who', + description: 'Name the caller', + }).server((_input, { context }) => { + seen.push(context.run.principal?.id) + return 'ok' + }), + ], + } + }, + }) + return { plugin, counts, seen } +} + +function setup( + plugins: () => ReadonlyArray, + idleCloseMs = 0, + name = 'test/host-plugins', +) { + const host = createHarnessHost({ + persistence: memoryPersistence(), + lifecycle: { idleCloseMs }, + }) + const harness = defineHarness({ + name, + adapter: mockAdapter(callWho).adapter, + plugins, + }) + return { host, harness } +} + +describe('lifetime host', () => { + it('sets up once for two sessions, and disposes after the last session closes', async () => { + const { plugin, counts } = countedPlugin() + const { host, harness } = setup(() => [plugin]) + + const first = await host.open(harness, { threadId: 'a' }) + const second = await host.open(harness, { threadId: 'b' }) + expect(counts).toEqual({ setups: 1, disposed: 0 }) + + await first.close() + expect(counts).toEqual({ setups: 1, disposed: 0 }) + await second.close() + expect(counts).toEqual({ setups: 1, disposed: 1 }) + + // The next session sets it up again. + const third = await host.open(harness, { threadId: 'c' }) + expect(counts).toEqual({ setups: 2, disposed: 1 }) + await third.close() + await host.close() + }) + + it('disposes at host.close() while sessions are open', async () => { + const { plugin, counts } = countedPlugin() + const { host, harness } = setup(() => [plugin], 60_000) + await host.open(harness, { threadId: 'a' }) + + await host.close() + + expect(counts).toEqual({ setups: 1, disposed: 1 }) + }) + + it('sets up once for each harness name', async () => { + const { plugin, counts } = countedPlugin() + const one = setup(() => [plugin], 60_000, 'test/one') + const other = defineHarness({ + name: 'test/other', + adapter: mockAdapter(callWho).adapter, + plugins: () => [plugin], + }) + + await one.host.open(one.harness, { threadId: 'a' }) + await one.host.open(other, { threadId: 'a' }) + + expect(counts.setups).toBe(2) + await one.host.close() + expect(counts.disposed).toBe(2) + }) + + it('gives its tool the user of the calling run', async () => { + const { plugin, seen } = countedPlugin() + const { host, harness } = setup(() => [plugin], 60_000) + const alice = await host.open(harness, { + threadId: 'a', + principal: { id: 'alice' }, + }) + const bob = await host.open(harness, { + threadId: 'b', + principal: { id: 'bob' }, + }) + + await alice.send('who am I?') + await bob.send('who am I?') + + expect(seen).toEqual(['alice', 'bob']) + await host.close() + }) + + it('still sets up a turn plugin for each turn, and a session plugin for each session', async () => { + const counts = { host: 0, session: 0, turn: 0 } + const shared = definePlugin({ + name: 'test/host', + lifetime: 'host', + setup: () => { + counts.host += 1 + }, + }) + const perSession = definePlugin({ + name: 'test/session', + setup: () => { + counts.session += 1 + }, + }) + const perTurn = definePlugin({ + name: 'test/turn', + lifetime: 'turn', + setup: () => { + counts.turn += 1 + }, + }) + const host = createHarnessHost({ persistence: memoryPersistence() }) + const harness = defineHarness({ + name: 'test/lifetimes', + adapter: mockAdapter(() => text('hi')).adapter, + plugins: () => [shared, perSession, perTurn], + }) + + for (const threadId of ['a', 'b']) { + const session = await host.open(harness, { threadId }) + await session.send('one') + await session.send('two') + } + + expect(counts).toEqual({ host: 1, session: 2, turn: 4 }) + await host.close() + }) + + it('runs its session part for each session, with that session', async () => { + const closed: Array = [] + const plugin = definePlugin({ + name: 'test/layered', + lifetime: 'host', + setup: () => ({ + session: async (ctx) => { + await ctx.resources.acquire( + () => ctx.session.threadId, + (threadId) => closed.push(threadId), + ) + return { + commands: { + thread: defineCommand({ + description: 'Name the thread', + run: () => ctx.session.threadId, + }), + }, + } + }, + }), + }) + const { host, harness } = setup(() => [plugin]) + const first = await host.open(harness, { threadId: 'a' }) + const second = await host.open(harness, { threadId: 'b' }) + + await expect(first.command({ name: 'thread' })).resolves.toBe('a') + await expect(second.command({ name: 'thread' })).resolves.toBe('b') + await first.close() + + expect(closed).toEqual(['a']) + await host.close() + }) + + it('sets up again after a setup that failed', async () => { + let attempts = 0 + const plugin = definePlugin({ + name: 'test/flaky', + lifetime: 'host', + setup: () => { + attempts += 1 + if (attempts === 1) throw new Error('no pool') + }, + }) + const { host, harness } = setup(() => [plugin]) + + await expect(host.open(harness, { threadId: 'a' })).rejects.toThrow( + 'no pool', + ) + const session = await host.open(harness, { threadId: 'a' }) + + expect(attempts).toBe(2) + await session.close() + await host.close() + }) +}) diff --git a/packages/ai-harness/tests/purge.test.ts b/packages/ai-harness/tests/purge.test.ts new file mode 100644 index 0000000000..b0fb354c1a --- /dev/null +++ b/packages/ai-harness/tests/purge.test.ts @@ -0,0 +1,198 @@ +import { describe, expect, it } from 'vitest' +import { memoryLogStore, memoryPersistence } from '@tanstack/ai-persistence' +import { + configOption, + createHarnessHost, + defineHarness, + definePlugin, + logIdOf, +} from '../src' +import { mockAdapter, text } from './helpers' +import type { HarnessPersistence } from '../src' + +const encoder = new TextEncoder() + +/** The metadata rows of a thread that a delete removes. */ +const ROWS = [ + 'harness:settings', + 'harness:config', + 'harness:interrupted', + 'harness:usage', + 'harness:revert', + 'harness:log-checkpoint', +] + +const level = definePlugin({ + name: 'test/level', + setup: () => ({ + config: { + level: configOption.select({ options: ['low', 'high'], default: 'low' }), + }, + }), +}) + +/** The stores of a host with no log, or of a durable host. */ +function persistenceOf(shape: 'no log' | 'durable', withIndex = true) { + const { messages, runs, metadata, sessions, artifacts, blobs } = + memoryPersistence().stores + const shared = { + runs, + metadata, + artifacts, + blobs, + ...(withIndex ? { sessions } : {}), + } + const persistence: HarnessPersistence = + shape === 'durable' + ? { stores: { log: memoryLogStore(), ...shared } } + : { stores: { messages, ...shared } } + return { persistence, metadata, artifacts, blobs, sessions } +} + +function setup(shape: 'no log' | 'durable', withIndex = true) { + const stores = persistenceOf(shape, withIndex) + const host = createHarnessHost({ + persistence: stores.persistence, + lifecycle: { idleCloseMs: 0 }, + }) + const harness = defineHarness({ + name: 'test/purge', + adapter: mockAdapter(() => text('answer')).adapter, + plugins: () => [level], + }) + return { ...stores, host, harness } +} + +/** The transcript of `threadId` as the store keeps it: log entries or messages. */ +async function storedTranscript( + persistence: HarnessPersistence, + threadId: string, + tenantId?: string, +) { + const key = logIdOf({ tenantId, logId: threadId }) + const { log, messages } = persistence.stores + if (log) return log.read({ logId: key }) + return (await messages?.loadThread(key)) ?? [] +} + +describe('host.sessions.delete', () => { + it.each(['no log', 'durable'] as const)( + 'purges the thread and its child on a host with %s', + async (shape) => { + const { + host, + harness, + persistence, + metadata, + artifacts, + blobs, + sessions, + } = setup(shape) + const key = logIdOf({ logId: 't1' }) + + // The child: a thread with its own transcript, then an entry under t1. + const child = await host.open(harness, { threadId: 'child-1' }) + await child.send('child work') + await child.close() + const childEntry = await sessions.get({ threadId: 'child-1' }) + await sessions.upsertIf({ + entry: { + threadId: 'child-1', + members: [], + createdAt: 1, + updatedAt: 1, + ...childEntry, + parentThreadId: 't1', + }, + revision: + (await sessions.getVersioned({ threadId: 'child-1' }))?.revision ?? + null, + }) + + const session = await host.open(harness, { threadId: 't1' }) + await session.send('hi') + await session.configure({ settings: { instructions: 'Be brief.' } }) + await session.setConfig({ key: 'level', value: 'high' }) + const upload = await session.putMedia({ + body: encoder.encode('cat'), + mimeType: 'image/png', + name: 'cat.png', + }) + // A generated file of a turn, as `withGenerationPersistence` keeps it. + await blobs.put('artifacts/op-1:image-1/a-1', encoder.encode('art')) + await artifacts.save({ + artifactId: 'a-1', + runId: 'op-1:image-1', + threadId: key, + blobKey: 'artifacts/op-1:image-1/a-1', + name: 'art.png', + mimeType: 'image/png', + size: 3, + createdAt: 1, + }) + await session.close() + // Rows that a turn writes only in some cases: seed them. + const seeded = [ + 'harness:usage', + 'harness:interrupted', + 'harness:revert', + 'harness:log-checkpoint', + ] + for (const namespace of seeded) { + await metadata.set(namespace, key, { seeded: true }) + } + expect(await storedTranscript(persistence, 't1')).not.toEqual([]) + expect(await storedTranscript(persistence, 'child-1')).not.toEqual([]) + expect(await metadata.get('harness:settings', key)).not.toBeNull() + expect(await metadata.get('harness:config', key)).not.toBeNull() + + await host.sessions.delete({ threadId: 't1' }) + + expect(await sessions.get({ threadId: 't1' })).toBeUndefined() + expect(await sessions.get({ threadId: 'child-1' })).toBeUndefined() + expect(await storedTranscript(persistence, 't1')).toEqual([]) + expect(await storedTranscript(persistence, 'child-1')).toEqual([]) + expect(await artifacts.listForThread(key)).toEqual([]) + expect((await blobs.list({ prefix: 'artifacts/' })).objects).toEqual([]) + expect(await artifacts.get(upload.id)).toBeNull() + for (const namespace of ROWS) { + expect(await metadata.get(namespace, key)).toBeNull() + } + await host.close() + }, + ) + + it('closes the open session of the thread on a host with no session index', async () => { + const { host, harness, persistence } = setup('no log', false) + const session = await host.open(harness, { threadId: 't1' }) + await session.send('hi') + + await host.sessions.delete({ threadId: 't1' }) + + expect(host.isOpen(harness, { threadId: 't1' })).toBe(false) + expect(await storedTranscript(persistence, 't1')).toEqual([]) + await host.close() + }) + + it('leaves the thread with the same id in another tenant', async () => { + const { host, harness, persistence, metadata } = setup('no log') + const other = { id: 'u2', tenantId: 'acme' } + const kept = await host.open(harness, { threadId: 't1', principal: other }) + await kept.send('keep me') + await kept.configure({ settings: { instructions: 'Stay.' } }) + await kept.close() + const mine = await host.open(harness, { threadId: 't1' }) + await mine.send('delete me') + await mine.close() + + await host.sessions.delete({ threadId: 't1' }) + + const otherKey = logIdOf({ tenantId: 'acme', logId: 't1' }) + expect(await storedTranscript(persistence, 't1', 'acme')).not.toEqual([]) + expect(await metadata.get('harness:settings', otherKey)).not.toBeNull() + expect( + await host.sessions.get({ threadId: 't1', principal: other }), + ).toMatchObject({ threadId: 't1', tenantId: 'acme' }) + await host.close() + }) +}) diff --git a/packages/ai-harness/tests/types.test-d.ts b/packages/ai-harness/tests/types.test-d.ts index 0d33227389..fa44834b0a 100644 --- a/packages/ai-harness/tests/types.test-d.ts +++ b/packages/ai-harness/tests/types.test-d.ts @@ -99,7 +99,7 @@ it('takes a harness without an adapter, and turn plugins', () => { expectTypeOf(picked.adapter).toEqualTypeOf< AnyTextAdapter | KeyedAdapter | undefined >() - expectTypeOf().toEqualTypeOf<'session' | 'turn'>() + expectTypeOf().toEqualTypeOf<'session' | 'turn' | 'host'>() }) it('takes attempt and time limits, and the log folds, on the harness', () => { diff --git a/packages/ai-mcp/src/harness-plugin.ts b/packages/ai-mcp/src/harness-plugin.ts index a22a739f2a..ef924ed90c 100644 --- a/packages/ai-mcp/src/harness-plugin.ts +++ b/packages/ai-mcp/src/harness-plugin.ts @@ -5,7 +5,9 @@ import type { AnyTool, JSONSchema } from '@tanstack/ai' import type { AnyCommand, PluginPrompt, + PluginResources, PluginSessionApi, + PluginSetupContext, TurnInfo, } from '@tanstack/ai-harness' import type { OnElicit } from './client' @@ -62,8 +64,10 @@ export type McpServerStatus = { * - The servers connect in the background. The first turn waits for them. * - `/mcp` lists the servers and their status. UIs read the same list from * the plugin state. - * - A `stdio` server runs as a child process (Node only). It stops when the - * session closes. + * - A `stdio` server runs as one child process (Node only) for every + * session of the harness on the host. It stops when the last of those + * sessions closes, or at `host.close()`. The other servers connect for + * each session. * - An `http` server with `oauth: true` signs in through `/connect `. * * @param options.servers - The servers, keyed by name. The name is the tool @@ -90,134 +94,206 @@ export function mcp(options: { const servers = Object.entries(options.servers) return definePlugin({ name: 'tanstack/mcp', - setup: async (ctx) => { - const statuses = new Map() - for (const [name] of servers) { - statuses.set(name, { name, status: 'connecting', toolCount: 0 }) + // The stdio servers connect once for the host. The rest is per session. + lifetime: 'host', + setup: (host) => { + // The sessions whose tool call of a shared server runs now, newest last. + const calling: Array = [] + // ponytail: a spec 2025 request of a shared server goes to the session + // of the newest running call. Two sessions that call one stdio server + // at the same time can get the question of the other one. + const askCaller: OnElicit = async (params) => { + const session = calling.at(-1) + if (!session) return { action: 'decline' } + return elicitFor(session)(params) } - const state = ctx.state({ servers: [...statuses.values()] }) - const setStatus = async (status: McpServerStatus) => { - const old = JSON.stringify(statuses.get(status.name)) - statuses.set(status.name, status) - if (old === JSON.stringify(status)) return - await state.update(() => ({ servers: [...statuses.values()] })) + const shared = new Map() + for (const [name, server] of servers) { + if (server.type !== 'stdio') continue + shared.set(name, openServer(host.resources, name, server, askCaller)) } - // State from an earlier run of this session is out of date. - await state.update(() => ({ servers: [...statuses.values()] })) + return { + session: (ctx) => sessionPart(ctx, servers, shared, calling), + } + }, + }) +} - /** Connects one server, in the background. Never rejects. */ - const connect = async (name: string, server: McpServerConfig) => { - const opened = await ctx.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, +/** The part of {@link mcp} that each session sets up. */ +async function sessionPart( + ctx: PluginSetupContext, + servers: ReadonlyArray<[string, McpServerConfig]>, + shared: ReadonlyMap, + calling: Array, +) { + const statuses = new Map() + for (const [name] of servers) { + statuses.set(name, { name, status: 'connecting', toolCount: 0 }) + } + const state = ctx.state({ servers: [...statuses.values()] }) + const setStatus = async (status: McpServerStatus) => { + const old = JSON.stringify(statuses.get(status.name)) + statuses.set(status.name, status) + if (old === JSON.stringify(status)) return + await state.update(() => ({ servers: [...statuses.values()] })) + } + // State from an earlier run of this session is out of date. + 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. + */ + const connect = async (name: string, server: McpServerConfig) => { + const sharedServer = shared.get(name) + const opened = await (sharedServer ?? + openServer(ctx.resources, name, server, elicitFor(ctx.session))) + const failed = opened instanceof Error + const status: McpServerStatus = failed + ? { name, status: 'failed', error: opened.message, toolCount: 0 } + : { name, status: 'connected', toolCount: opened.tools.length } + // 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 = sharedServer + ? trackCaller(opened.tools, ctx.session, calling) + : opened.tools + return askForInput(tools, ctx.session) + } + + const commands: Record = { + mcp: defineCommand({ + description: 'Show the MCP servers and their status', + run: () => [...statuses.values()].map(statusLine).join('\n'), + }), + } + const prompts: Array = [] + // One tool list for each server, in the order of `servers`. + const lists = await Promise.all( + 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 + } + // 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, + requestOptions: requestOptionsOf(server), + ...(server.fetch ? { fetch: server.fetch } : {}), + }).setup(ctx) + Object.assign(commands, connector.commands) + 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) => { + try { + const tools = await connector.discoverTools(turn) + const signedIn = + tools.length > 0 || (await turn.run.credentials.get(name)) !== null + await setStatus( + signedIn + ? { name, status: 'connected', toolCount: tools.length } + : { + name, + status: 'failed', + error: `Not signed in. Run /connect ${name}.`, + toolCount: 0, }, - elicitFor(ctx.session), - ) - try { - const tools = askForInput( - forCodeMode(server, await client.tools()), - ctx.session, - ) - return { client, tools } - } catch (error) { - // Connected, but no tool list: the server failed. - await client.close().catch(() => undefined) - throw error - } - }, - ({ client }) => client.close(), ) - .catch((error: unknown) => new Error(errorText(error))) - const failed = opened instanceof Error - const status: McpServerStatus = failed - ? { name, status: 'failed', error: opened.message, toolCount: 0 } - : { name, status: 'connected', toolCount: opened.tools.length } - // After the session closed, the state store can be gone. `/mcp` - // reads `statuses`, which is set first. - await setStatus(status).catch(() => undefined) - return failed ? [] : opened.tools + return forCodeMode(server, tools) + } catch (error) { + // One server that fails must not hide the tools of the others. + await setStatus({ + name, + status: 'failed', + error: errorText(error), + toolCount: 0, + }) + return [] + } } + }), + ) - const commands: Record = { - mcp: defineCommand({ - description: 'Show the MCP servers and their status', - run: () => [...statuses.values()].map(statusLine).join('\n'), - }), - } - const prompts: Array = [] - // One tool list for each server, in the order of `servers`. - const lists = await Promise.all( - 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 - } - // 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, + return { + commands, + prompts, + discoverTools: async (turn: TurnInfo) => { + const found = await Promise.all(lists.map((list) => list(turn))) + return found.flat() + }, + } +} + +/** A server that the sessions of a host share: its tools, or its error. */ +type SharedServer = ReturnType + +/** + * Connect one server in `resources`. Resolves to the error of a server that + * failed, so it never rejects. + */ +function openServer( + resources: PluginResources, + name: string, + server: McpServerConfig, + onElicit: OnElicit, +) { + 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), - ...(server.fetch ? { fetch: server.fetch } : {}), - }).setup(ctx) - Object.assign(commands, connector.commands) - 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) => { - try { - const tools = await connector.discoverTools(turn) - const signedIn = - tools.length > 0 || - (await turn.run.credentials.get(name)) !== null - await setStatus( - signedIn - ? { name, status: 'connected', toolCount: tools.length } - : { - name, - status: 'failed', - error: `Not signed in. Run /connect ${name}.`, - toolCount: 0, - }, - ) - return forCodeMode(server, tools) - } catch (error) { - // One server that fails must not hide the tools of the others. - await setStatus({ - name, - status: 'failed', - error: errorText(error), - toolCount: 0, - }) - return [] - } - } - }), - ) + clientOptions: harnessClientOptions, + }, + onElicit, + ) + try { + return { client, tools: forCodeMode(server, await client.tools()) } + } catch (error) { + // Connected, but no tool list: the server failed. + await client.close().catch(() => undefined) + throw error + } + }, + ({ client }) => client.close(), + ) + .catch((error: unknown) => new Error(errorText(error))) +} - return { - commands, - prompts, - discoverTools: async (turn) => { - const found = await Promise.all(lists.map((list) => list(turn))) - return found.flat() - }, - } - }, +/** Marks `session` as the caller of a shared server while a tool call runs. */ +function trackCaller( + tools: ReadonlyArray, + session: PluginSessionApi, + calling: Array, +) { + return tools.map((tool): AnyTool => { + const execute = tool.execute + if (!execute) return tool + return { + ...tool, + execute: async (args, context) => { + calling.push(session) + try { + return await execute(args, context) + } finally { + calling.splice(calling.lastIndexOf(session), 1) + } + }, + } }) } diff --git a/packages/ai-mcp/tests/harness-plugin.test.ts b/packages/ai-mcp/tests/harness-plugin.test.ts index 9ef7cdef56..7d66078e2b 100644 --- a/packages/ai-mcp/tests/harness-plugin.test.ts +++ b/packages/ai-mcp/tests/harness-plugin.test.ts @@ -67,6 +67,23 @@ async function openWith( return { host, session, model } } +/** A stdio server config that runs this file as an SDK server. */ +function stdioServer(): McpServerConfig { + return { + type: 'stdio', + command: process.execPath, + args: ['--import', 'jiti/register', fileURLToPath(import.meta.url)], + cwd: fileURLToPath(new URL('..', import.meta.url)), + env: { [childEnv]: '1' }, + } +} + +/** The pid that the `pid` tool gave in model call `call` of `model`. */ +function pidOf(model: ReturnType | undefined, call: number) { + const seen = JSON.stringify(model?.calls[call]?.messages) + return Number(/pid:(\d+)/.exec(seen)?.[1]) +} + function isRunning(pid: number) { try { process.kill(pid, 0) @@ -381,23 +398,11 @@ if (process.env[childEnv] === '1') { ) }) - it('runs a stdio server as a child process and stops it when the session closes', async () => { + it('runs a stdio server as a child process and stops it when the host closes', async () => { const model = recorder('local_pid') - const { host, session } = await openWith( - { - local: { - type: 'stdio', - command: process.execPath, - args: ['--import', 'jiti/register', fileURLToPath(import.meta.url)], - cwd: fileURLToPath(new URL('..', import.meta.url)), - env: { [childEnv]: '1' }, - }, - }, - model, - ) + const { host, session } = await openWith({ local: stdioServer() }, model) await session.send('which process runs the server?') - const seen = JSON.stringify(model.calls[1]?.messages) - const pid = Number(/pid:(\d+)/.exec(seen)?.[1]) + const pid = pidOf(model, 1) expect(isRunning(pid)).toBe(true) await host.close() @@ -406,6 +411,41 @@ if (process.env[childEnv] === '1') { }) }, 120_000) + it('runs one stdio child for two sessions, and stops it when the last one closes', async () => { + const host = createHarnessHost({ + persistence: memoryPersistence(), + lifecycle: { idleCloseMs: 0 }, + }) + cleanups.push(() => host.close()) + // One harness name, so one shared setup. Each model calls the tool once. + const models = [recorder('local_pid'), recorder('local_pid')] + const [first, second] = await Promise.all( + models.map((model, at) => + host.open( + defineHarness({ + name: 'test/mcp-plugin', + adapter: model.adapter, + plugins: () => [mcp({ servers: { local: stdioServer() } })], + }), + { threadId: `t-${at}` }, + ), + ), + ) + + await first?.send('which process runs the server?') + await second?.send('which process runs the server?') + const pid = pidOf(models[0], 1) + expect(pid).toBeGreaterThan(0) + expect(pidOf(models[1], 1)).toBe(pid) + + await first?.close() + expect(isRunning(pid)).toBe(true) + await second?.close() + await vi.waitFor(() => expect(isRunning(pid)).toBe(false), { + timeout: 10_000, + }) + }, 120_000) + it('signs in to an oauth server with /connect, then gives its tools', async () => { const protectedServer = await startProtectedServer() cleanups.push(() => protectedServer.close()) From 6510a12b6c20edfb6a79bf07fd142dbd6fb5c193 Mon Sep 17 00:00:00 2001 From: Alem Tuzlak Date: Sat, 17 Oct 2026 14:22:48 +0200 Subject: [PATCH 09/12] fix(ai-harness): definePlugin infers its definition only from the argument --- packages/ai-harness/src/plugins.ts | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/packages/ai-harness/src/plugins.ts b/packages/ai-harness/src/plugins.ts index 7f5cf19f68..75139b7693 100644 --- a/packages/ai-harness/src/plugins.ts +++ b/packages/ai-harness/src/plugins.ts @@ -488,7 +488,9 @@ export function definePlugin< definition: TDefinition & { metadataSchema?: TMetadataSchema } & NoSetupKeys, -): HarnessPlugin> { + // `NoInfer`: the type of a plugins array around the call must not pick + // `TDefinition`. Only the definition does. +): HarnessPlugin, MetadataOf> { if (definition.name.trim() === '') { throw new Error('definePlugin requires a non-empty name') } From 974b61ba694b187fe657edaecbf6bfb9afc4288f Mon Sep 17 00:00:00 2001 From: Alem Tuzlak Date: Sat, 17 Oct 2026 14:37:27 +0200 Subject: [PATCH 10/12] fix(ai-harness): key host plugins by the harness definition, not its name --- packages/ai-harness/src/host-plugins.ts | 40 +++++++++++++------ .../ai-harness/src/session/harness-session.ts | 3 +- .../ai-harness/tests/host-plugins.test.ts | 26 ++++++++++++ 3 files changed, 54 insertions(+), 15 deletions(-) diff --git a/packages/ai-harness/src/host-plugins.ts b/packages/ai-harness/src/host-plugins.ts index e837fb759b..71690742fe 100644 --- a/packages/ai-harness/src/host-plugins.ts +++ b/packages/ai-harness/src/host-plugins.ts @@ -1,4 +1,5 @@ import { ResourceScope, disposeAll } from './resources' +import type { AnyHarness } from './define' import type { HostPluginContext, HostPluginContributions, @@ -35,35 +36,48 @@ export function hostPluginContext( /** * @internal The plugins with `lifetime: 'host'` of one host. Each one is - * set up once for each harness name and plugin name, and shared by a + * set up once for each harness definition and plugin name, and shared by a * reference count. The last release disposes it, and so does `close()`. + * + * The key is the harness object, not its name. An app that builds a new + * definition (with other plugin options) under the same name gets its own + * setup, never the one of the older definition. */ export class HostPlugins { - private readonly entries = new Map() + private readonly entries = new Map>() private closed = false - /** The shared setup of `plugin` for the harness `harness`. */ + /** The shared setup of `plugin` for `harness`. */ async acquire( - harness: string, + harness: AnyHarness, plugin: HostPluginDefinition, ): Promise { if (this.closed) { throw new Error('The host is closed. Create a new host to open sessions.') } - const key = `${harness}\u0000${plugin.name}` - let entry = this.entries.get(key) + const plugins = this.entries.get(harness) ?? new Map() + this.entries.set(harness, plugins) + const key = plugin.name + /** Drop `entry` when it is still the one of `key`. */ + const forget = (entry: Entry) => { + if (plugins.get(key) !== entry) return false + plugins.delete(key) + if (plugins.size === 0) this.entries.delete(harness) + return true + } + let entry = plugins.get(key) if (!entry) { const scope = new ResourceScope() const ready = Promise.resolve().then(() => - plugin.setup(hostPluginContext(harness, scope)), + plugin.setup(hostPluginContext(harness.name, scope)), ) const created: Entry = { refs: 0, scope, ready } // A failed setup is not kept: the next session sets it up again. ready.catch(() => { - if (this.entries.get(key) === created) this.entries.delete(key) + forget(created) void scope.dispose().catch(() => {}) }) - this.entries.set(key, created) + plugins.set(key, created) entry = created } const held = entry @@ -82,9 +96,7 @@ export class HostPlugins { if (released) return Promise.resolve() released = true held.refs -= 1 - const isLast = held.refs <= 0 && this.entries.get(key) === held - if (!isLast) return Promise.resolve() - this.entries.delete(key) + if (held.refs > 0 || !forget(held)) return Promise.resolve() return held.scope.dispose() }, } @@ -93,7 +105,9 @@ export class HostPlugins { /** Dispose every shared setup. Later `acquire` calls throw. */ close(): Promise { this.closed = true - const scopes = [...this.entries.values()].map((entry) => entry.scope) + const scopes = [...this.entries.values()].flatMap((plugins) => + [...plugins.values()].map((entry) => entry.scope), + ) this.entries.clear() return disposeAll(scopes) } diff --git a/packages/ai-harness/src/session/harness-session.ts b/packages/ai-harness/src/session/harness-session.ts index 4690d46930..8418068684 100644 --- a/packages/ai-harness/src/session/harness-session.ts +++ b/packages/ai-harness/src/session/harness-session.ts @@ -164,8 +164,7 @@ export class SessionImpl< this.plugins.filter((plugin) => plugin.lifetime !== 'turn'), { threadId: this.threadId, - hostPlugin: (plugin) => - this.hostPlugins.acquire(this.harness.name, plugin), + hostPlugin: (plugin) => this.hostPlugins.acquire(this.harness, plugin), registry: this.agentRegistry, harnessTools: this.harness.tools ?? [], harnessProvides: (this.harness.middleware ?? []).flatMap( diff --git a/packages/ai-harness/tests/host-plugins.test.ts b/packages/ai-harness/tests/host-plugins.test.ts index 819fe82037..8395aaed61 100644 --- a/packages/ai-harness/tests/host-plugins.test.ts +++ b/packages/ai-harness/tests/host-plugins.test.ts @@ -64,6 +64,32 @@ function setup( } describe('lifetime host', () => { + it('gives a new harness definition with the same name its own setup', async () => { + const first = countedPlugin() + const second = countedPlugin() + const host = createHarnessHost({ + persistence: memoryPersistence(), + lifecycle: { idleCloseMs: 0 }, + }) + // Two definitions under one name, with other plugin options, as an app + // that builds its harness for each request does. + const define = (plugin: HarnessPlugin) => + defineHarness({ + name: 'test/same-name', + adapter: mockAdapter(callWho).adapter, + plugins: () => [plugin], + }) + + const a = await host.open(define(first.plugin), { threadId: 'a' }) + const b = await host.open(define(second.plugin), { threadId: 'b' }) + + expect(first.counts.setups).toBe(1) + expect(second.counts.setups).toBe(1) + await a.close() + await b.close() + await host.close() + }) + it('sets up once for two sessions, and disposes after the last session closes', async () => { const { plugin, counts } = countedPlugin() const { host, harness } = setup(() => [plugin]) From 657ee91d4b5b3facb122ab3bac12d807991abb76 Mon Sep 17 00:00:00 2001 From: Alem Tuzlak Date: Sat, 17 Oct 2026 15:16:47 +0200 Subject: [PATCH 11/12] test(ai-mcp): open two sessions of one harness definition to share a stdio server --- packages/ai-mcp/tests/connector-helpers.ts | 12 +++++-- packages/ai-mcp/tests/harness-plugin.test.ts | 38 ++++++++++---------- 2 files changed, 27 insertions(+), 23 deletions(-) diff --git a/packages/ai-mcp/tests/connector-helpers.ts b/packages/ai-mcp/tests/connector-helpers.ts index d7ce558e12..b21d961ae2 100644 --- a/packages/ai-mcp/tests/connector-helpers.ts +++ b/packages/ai-mcp/tests/connector-helpers.ts @@ -41,13 +41,19 @@ export function mockTextAdapter(chatStream: AnyTextAdapter['chatStream']) { return adapter } -/** A model that records the calls it gets, calls `tool` once, then stops. */ -export function recorder(tool?: string) { +/** + * A model that records the calls it gets, calls `tool` on the calls that + * `callsTool` picks (by default only the first one), and else stops. + */ +export function recorder( + tool?: string, + callsTool: (call: number) => boolean = (call) => call === 1, +) { const calls: Array = [] const now = () => Date.now() const adapter = mockTextAdapter((options) => { calls.push(options) - const first = calls.length === 1 && tool !== undefined + const first = tool !== undefined && callsTool(calls.length) return (async function* (): AsyncGenerator { yield { type: EventType.RUN_STARTED, diff --git a/packages/ai-mcp/tests/harness-plugin.test.ts b/packages/ai-mcp/tests/harness-plugin.test.ts index 7d66078e2b..445193a64c 100644 --- a/packages/ai-mcp/tests/harness-plugin.test.ts +++ b/packages/ai-mcp/tests/harness-plugin.test.ts @@ -417,30 +417,28 @@ if (process.env[childEnv] === '1') { lifecycle: { idleCloseMs: 0 }, }) cleanups.push(() => host.close()) - // One harness name, so one shared setup. Each model calls the tool once. - const models = [recorder('local_pid'), recorder('local_pid')] - const [first, second] = await Promise.all( - models.map((model, at) => - host.open( - defineHarness({ - name: 'test/mcp-plugin', - adapter: model.adapter, - plugins: () => [mcp({ servers: { local: stdioServer() } })], - }), - { threadId: `t-${at}` }, - ), - ), - ) + // One harness definition, so one shared setup. The model calls the + // tool on the first call of each turn: calls 1 and 3. + const model = recorder('local_pid', (call) => call % 2 === 1) + const harness = defineHarness({ + name: 'test/mcp-plugin', + adapter: model.adapter, + plugins: () => [mcp({ servers: { local: stdioServer() } })], + }) + const [first, second] = await Promise.all([ + host.open(harness, { threadId: 't-0' }), + host.open(harness, { threadId: 't-1' }), + ]) - await first?.send('which process runs the server?') - await second?.send('which process runs the server?') - const pid = pidOf(models[0], 1) + await first.send('which process runs the server?') + await second.send('which process runs the server?') + const pid = pidOf(model, 1) expect(pid).toBeGreaterThan(0) - expect(pidOf(models[1], 1)).toBe(pid) + expect(pidOf(model, 3)).toBe(pid) - await first?.close() + await first.close() expect(isRunning(pid)).toBe(true) - await second?.close() + await second.close() await vi.waitFor(() => expect(isRunning(pid)).toBe(false), { timeout: 10_000, }) From ce1293c6513e0aeed522d103971446886b389113 Mon Sep 17 00:00:00 2001 From: Alem Tuzlak Date: Sat, 17 Oct 2026 15:17:26 +0200 Subject: [PATCH 12/12] test(ai-harness, ai-harness-cli): give polled checks 5 s, for slower CI runs --- packages/ai-harness-cli/tests/signals.test.ts | 31 ++++++++++++------- packages/ai-harness/tests/handles.test.ts | 13 ++++++-- packages/ai-harness/tests/observe.test.ts | 17 ++++++---- 3 files changed, 40 insertions(+), 21 deletions(-) diff --git a/packages/ai-harness-cli/tests/signals.test.ts b/packages/ai-harness-cli/tests/signals.test.ts index b6c79d7621..7895f24165 100644 --- a/packages/ai-harness-cli/tests/signals.test.ts +++ b/packages/ai-harness-cli/tests/signals.test.ts @@ -16,6 +16,13 @@ import type { FakeResponse } from '@tanstack/ai/testing' import type { AnyHarness } from '@tanstack/ai-harness' import type { RunCliOptions } from '../src' +/** + * `vi.waitFor` with more time: CI runs these tests slower, for example under + * coverage, and a poll costs nothing when the check passes early. + */ +const waitFor = (check: () => T | Promise) => + vi.waitFor(check, { timeout: 5_000 }) + /** Resolves when `signal` aborts. */ function untilAborted(signal: AbortSignal | undefined) { return new Promise((resolve) => { @@ -88,12 +95,12 @@ describe('signals in line mode', () => { { argv: [], stdin: asStdin(input, { isTTY: false }) }, ) input.write('hello\n') - await vi.waitFor(() => expect(calls).toHaveLength(1)) + await waitFor(() => expect(calls).toHaveLength(1)) cli.signals.emit('SIGINT') - await vi.waitFor(() => expect(calls[0]?.aborted).toBe(true)) + await waitFor(() => expect(calls[0]?.aborted).toBe(true)) input.write('again\n') - await vi.waitFor(() => expect(cli.stdout.text).toContain('Second answer.')) + await waitFor(() => expect(cli.stdout.text).toContain('Second answer.')) cli.signals.emit('SIGINT') expect(await cli.running).toBe(EXIT.cancelled) @@ -114,7 +121,7 @@ describe('signals in line mode', () => { { argv: [], stdin }, ) input.write('/connect acme\n') - await vi.waitFor(() => expect(rawModes).toEqual([true])) + await waitFor(() => expect(rawModes).toEqual([true])) // Raw mode makes Ctrl+C a typed character. input.write('\u0003') @@ -131,7 +138,7 @@ describe('signals in print mode', () => { defineHarness({ name: 'test/signal-print', adapter: model }), { argv: ['-p', 'hi'] }, ) - await vi.waitFor(() => expect(calls).toHaveLength(1)) + await waitFor(() => expect(calls).toHaveLength(1)) cli.signals.emit('SIGTERM') @@ -149,16 +156,16 @@ describe('signals in print mode', () => { }) const close = vi.spyOn(host, 'close') const cli = startCli(harness, { argv: ['-p', 'hi'], host }) - await vi.waitFor(() => expect(calls).toHaveLength(1)) + await waitFor(() => expect(calls).toHaveLength(1)) cli.signals.emit('SIGINT') - await vi.waitFor(() => expect(calls[0]?.aborted).toBe(true)) + await waitFor(() => expect(calls[0]?.aborted).toBe(true)) cli.signals.emit('SIGINT') release() expect(await cli.running).toBe(EXIT.cancelled) expect(close).not.toHaveBeenCalled() - await vi.waitFor(() => + await waitFor(() => expect(host.isOpen(harness, { threadId: 'main' })).toBe(false), ) await host.close() @@ -187,7 +194,7 @@ describe('signals in a custom ui', () => { return new Promise(() => {}) }, }) - await vi.waitFor(() => expect(shown).toBe(true)) + await waitFor(() => expect(shown).toBe(true)) cli.signals.emit('SIGTERM') @@ -207,7 +214,7 @@ describe('signals in --serve', () => { defineHarness({ name: 'test/signal-serve', adapter: model }), { argv: ['--serve', '--port', '0', '--token', 'secret'], persistence }, ) - const url = await vi.waitFor(() => { + const url = await waitFor(() => { const [, served] = /at (\S+)\n/.exec(cli.stderr.text) ?? [] expect(served).toBeDefined() return served @@ -224,7 +231,7 @@ describe('signals in --serve', () => { }), }) const { inputId } = await receipt.json() - await vi.waitFor(() => expect(calls).toHaveLength(1)) + await waitFor(() => expect(calls).toHaveLength(1)) cli.signals.emit('SIGTERM') @@ -271,7 +278,7 @@ describe('signals in the stdio modes', () => { adapter: fakeText(), }) const cli = startCli(harness, { argv, ...options }) - await vi.waitFor(() => expect(listened).toContain('end')) + await waitFor(() => expect(listened).toContain('end')) return cli } diff --git a/packages/ai-harness/tests/handles.test.ts b/packages/ai-harness/tests/handles.test.ts index f616f3230f..09f40fd10b 100644 --- a/packages/ai-harness/tests/handles.test.ts +++ b/packages/ai-harness/tests/handles.test.ts @@ -7,6 +7,13 @@ import { finalizeHandle } from '../src/session/handle' import { mockAdapter, text } from './helpers' import type { HostLifecycleOptions } from '../src' +/** + * `vi.waitFor` with more time: CI runs these tests slower, for example under + * coverage, and a poll costs nothing when the check passes early. + */ +const waitFor = (check: () => T | Promise) => + vi.waitFor(check, { timeout: 5_000 }) + const BASE = 'http://x/api/harness' /** A memory host with `lifecycle`, and a harness that answers each turn. */ @@ -161,7 +168,7 @@ describe('a call after the session closed', () => { 'Hello back.', ]) // The released handle held the session only for its calls. - await vi.waitFor(() => expect(isOpen('t1')).toBe(false)) + await waitFor(() => expect(isOpen('t1')).toBe(false)) await host.close() }) }) @@ -226,7 +233,7 @@ describe('the HTTP handler', () => { } expect([...statuses]).toEqual([200]) - await vi.waitFor(() => + await waitFor(() => expect(hostCountsOf(host)).toEqual({ sessions: 0, handles: 0, @@ -274,7 +281,7 @@ describe('the HTTP handler', () => { ) expect(await response.text()).toContain('Hello back.') - await vi.waitFor(() => expect(isOpen('t1')).toBe(false)) + await waitFor(() => expect(isOpen('t1')).toBe(false)) await host.close() }) }) diff --git a/packages/ai-harness/tests/observe.test.ts b/packages/ai-harness/tests/observe.test.ts index 26e23e1d1f..da1198976a 100644 --- a/packages/ai-harness/tests/observe.test.ts +++ b/packages/ai-harness/tests/observe.test.ts @@ -13,6 +13,13 @@ import { hostCountsOf } from '../src/host' import { fakeSocket, mockAdapter, text } from './helpers' import type { HarnessHostOptions, Principal, SessionEvent } from '../src' +/** + * `vi.waitFor` with more time: CI runs these tests slower, for example under + * coverage, and a poll costs nothing when the check passes early. + */ +const waitFor = (check: () => T | Promise) => + vi.waitFor(check, { timeout: 5_000 }) + const BASE = 'http://x/api/harness' const THREAD = 't1' @@ -171,9 +178,7 @@ describe('host.observe on a durable host', () => { // A new session of the thread writes to the same log. await seed() - await vi.waitFor(() => - expect(replies).toEqual(['Hello back.', 'Hello back.']), - ) + await waitFor(() => expect(replies).toEqual(['Hello back.', 'Hello back.'])) const transcript = await second.transcript() expect(transcript.map((message) => message.content)).toEqual([ 'Hello.', @@ -201,7 +206,7 @@ describe('host.observe on a durable host', () => { for await (const entry of view.events()) seen.push(entry) })() // The read has the log open, and it waits for new events. - await vi.waitFor(() => expect(seen.length).toBeGreaterThan(0)) + await waitFor(() => expect(seen.length).toBeGreaterThan(0)) await host.close() @@ -308,7 +313,7 @@ describe('the read routes', () => { receive({ type: 'harness.subscribe', threadId: THREAD }) receive({ type: 'harness.snapshot' }) - await vi.waitFor(() => + await waitFor(() => expect(sent.map((frame) => frame.type)).toEqual([ 'harness.hello', 'harness.snapshot', @@ -323,7 +328,7 @@ describe('the read routes', () => { }) // The stream goes on with the session that the input opened. - await vi.waitFor(() => + await waitFor(() => expect( sent.some( (frame) =>