From 7e387421ffa3a394d78674ff09d262836ebae5d1 Mon Sep 17 00:00:00 2001 From: Alem Tuzlak Date: Tue, 13 Oct 2026 22:10:58 +0200 Subject: [PATCH 001/131] feat(ai-persistence): session stores for inputs, credentials, a log, leases, work claims, and a session index --- .changeset/persistence-session-stores.md | 16 + docs/config.json | 2 +- docs/persistence/build-your-own-adapter.md | 5 + docs/persistence/store-reference.md | 471 ++++++++++- packages/ai-persistence/src/index.ts | 25 +- packages/ai-persistence/src/memory.ts | 346 +++++++- packages/ai-persistence/src/middleware.ts | 3 + packages/ai-persistence/src/subagent-runs.ts | 46 +- .../ai-persistence/src/testkit/conformance.ts | 770 +++++++++++++++++- packages/ai-persistence/src/types.ts | 424 ++++++++++ packages/ai-persistence/tests/inbox.test.ts | 37 + .../tests/memory.conformance.test.ts | 15 +- packages/ai-persistence/tests/memory.test.ts | 4 + .../tests/persistence-types.test-d.ts | 8 + .../tests/persistence-validation.test.ts | 21 +- .../tests/subagent-persistence-cards.test.ts | 94 +++ 16 files changed, 2264 insertions(+), 23 deletions(-) create mode 100644 .changeset/persistence-session-stores.md create mode 100644 packages/ai-persistence/tests/inbox.test.ts diff --git a/.changeset/persistence-session-stores.md b/.changeset/persistence-session-stores.md new file mode 100644 index 0000000000..1bc7360477 --- /dev/null +++ b/.changeset/persistence-session-stores.md @@ -0,0 +1,16 @@ +--- +'@tanstack/ai-persistence': minor +--- + +Add opt-in session stores. The chat and generation middleware work as before without them. + +- `inbox` (`InboxStore`): the inputs a host accepted, with an idempotent `append` by `inputId`, `listPending`, `markApplied`, `markRejected`, and `get`. +- `credentials` (`CredentialStore`): API keys and OAuth tokens by scope and id. An `oauth` credential can keep the OAuth `client` it was issued to. `list` never returns secret values. +- `log` (`LogStore`): an append-only log per thread. `append` writes a batch at a position, all or nothing, and rejects with `LogConflictError` when another writer took that position. `read` and `subscribe` complete it. `memoryLogStore()` is the reference store. It is not part of `memoryPersistence()`. +- `leases` (`LeaseStore`, `TurnLease`, `TurnLeaseKey`): turn leases from a system that already has them. +- `workClaims` (`WorkClaimStore`): claims on threads that have pending work, with `claim`, `release`, and `listExpired`. +- `sessions` (`SessionIndexStore`): one index entry per thread, with `upsert`, `get`, `list`, and `delete`. `list` filters by parent thread, owner, title `search`, `harness` name, and exact `metadata` values, and pages with a cursor. The optional `getVersioned` and `upsertIf` give compare-and-set writes. + +`defineInboxStore`, `defineCredentialStore`, `defineLogStore`, `defineWorkClaimStore`, and `defineSessionIndexStore` type an adapter. `memoryPersistence()` now also has `inbox`, `credentials`, `workClaims`, and `sessions`. `runPersistenceConformance` runs the cases of each session store when the store is present, so an existing backend needs no `skip` entry. + +With `stores.sessions`, `withPersistence` writes one index entry for each subagent child (`subagent:`), with the parent thread, the tool call, and the owner and `harness` name of the parent entry. diff --git a/docs/config.json b/docs/config.json index 98cc0f1dbd..4301a613c1 100644 --- a/docs/config.json +++ b/docs/config.json @@ -712,7 +712,7 @@ "label": "Build Your Own Adapter", "to": "persistence/build-your-own-adapter", "addedAt": "2026-08-04", - "updatedAt": "2026-09-29" + "updatedAt": "2026-10-13" }, { "label": "Migrations", diff --git a/docs/persistence/build-your-own-adapter.md b/docs/persistence/build-your-own-adapter.md index 02c5587640..3e78841c9c 100644 --- a/docs/persistence/build-your-own-adapter.md +++ b/docs/persistence/build-your-own-adapter.md @@ -189,6 +189,11 @@ The optional run methods are `listByThread`, `listByParentRun`, and `skipMethods`. An omitted `listByParentRun` needs no entry: the subagent checks skip on their own. +The session stores (`inbox`, `credentials`, `log`, `workClaims`, `sessions`) are +opt-in too. Their cases run only when your persistence has the store, so a backend +without them needs no `skip` entry. See the +[session stores](./store-reference#session-stores) for their rules. + ### Turn on the newer checks Some checks were added after the suite shipped. They are off by default, so an diff --git a/docs/persistence/store-reference.md b/docs/persistence/store-reference.md index 84d927932b..b52780f39d 100644 --- a/docs/persistence/store-reference.md +++ b/docs/persistence/store-reference.md @@ -16,6 +16,12 @@ there is no separate enable list. | `generationRuns` | Generation run status and result metadata, keyed by its own `runId`. | `withGenerationPersistence`, required | | `artifacts` | File metadata. Needs `blobs`. | `withGenerationPersistence`, portable snapshots | | `blobs` | File bytes. Needs `artifacts`. | `withGenerationPersistence`, portable snapshots | +| `inbox` | Inputs a host accepted, one entry per input id. | Your host code | +| `credentials` | API keys and OAuth tokens by scope and id. | Your host code | +| `log` | Append-only session log per thread. | Your host code | +| `leases` | Turn leases from a system that already has them. | Your host code | +| `workClaims` | Claims on threads that have pending work. | Your host code | +| `sessions` | One index entry per thread, for a list of sessions. | `withPersistence` (subagent children), your host code | Named groupings of the chat stores (`ChatTranscriptStores`, `ChatPersistenceStores`, `ChatWithInterruptsStores`) are covered in [Controls](./controls). @@ -588,6 +594,457 @@ asserts it. Ignoring `range` and returning the whole file is what makes entire artifact for every seek. A reference-only backend that stores no bytes skips `blobs` altogether instead. +## Session stores + +A host that keeps long-lived sessions needs more than a transcript. It must remember +the inputs it accepted, the keys a user saved, which host works on which thread, and a +list of sessions to show. The session stores hold that data. + +All of them are opt-in. `withPersistence` reads only `sessions`, to index subagent +children. The other stores are for your host code. + +## InboxStore + +A client sends an input, and the host must not lose it in a crash or run it twice. +Write the input to the inbox before you answer the client. `InboxStore` keeps one +entry per `inputId`. + +```ts +type InboxStatus = 'pending' | 'applied' | 'rejected' | 'expired' + +interface InboxEntry { + inputId: string // the idempotency key + threadId: string + principal?: { id: string; tenantId?: string } // who sent it + input: unknown // stored as-is + status: InboxStatus + createdAt: number // epoch ms + expiresAt?: number + operationId?: string // the operation that applied it + reason?: string // why it was rejected +} + +interface InboxStore { + append(entry: Omit): Promise + listPending(threadId: string): Promise> + markApplied(inputId: string, operationId: string): Promise + markRejected(inputId: string, reason: string): Promise + get(inputId: string): Promise +} +``` + +Rules for an implementation: + +- `inputId` is unique across all threads. Let the client send a UUID. +- `append` stores a new entry as `'pending'`. For an `inputId` that is already + stored, it returns the stored entry and ignores the new payload. +- `listPending` returns the pending entries of one thread, oldest `createdAt` first. +- `markApplied` and `markRejected` do nothing for an unknown `inputId`. +- `get` returns `null` for an unknown `inputId`. + +In SQL, make `input_id` the primary key and write with `INSERT ... ON CONFLICT DO +NOTHING`, then read the row back. + +## CredentialStore + +Your users connect their own accounts: an API key, or the OAuth tokens of a sign-in. +`CredentialStore` keeps each secret by scope and by an id that you choose, for example +`'github'`. + +```ts +import type { Scope } from '@tanstack/ai-persistence' + +type Credential = + | { type: 'api_key'; value: string } + | { + type: 'oauth' + accessToken: string + refreshToken?: string + expiresAt?: number // epoch ms + scopes?: Array + client?: { + clientId: string + clientSecret?: string + redirectUri?: string + issuer?: string + } + } + +interface CredentialStore { + get(scope: Scope, id: string): Promise + set(scope: Scope, id: string, credential: Credential): Promise + delete(scope: Scope, id: string): Promise + list( + scope: Scope, + ): Promise> +} +``` + +Rules for an implementation: + +- Key each credential by `scope.tenantId`, `scope.userId`, and `id`. Do not use + `scope.threadId`. A credential saved without `scope.userId` belongs to the tenant. +- A user credential is not visible at the tenant scope, and a tenant credential is + not visible at the user scope. +- `list` returns ids and types only. It never returns a secret value. +- `client` is the OAuth client that the tokens belong to. A token refresh needs it. +- Encrypt the secrets at rest. + +## LogStore + +Some hosts keep each thread's events, transcript, inputs, and tool steps in one +append-only log, so that a turn survives a crash. `LogStore` stores that log. + +```ts +interface LogRecord { + type: string + [key: string]: unknown +} + +interface LogEntry { + seq: number + record: LogRecord +} + +interface LogStore { + append( + threadId: string, + seq: number, + records: ReadonlyArray, + ): Promise + read( + threadId: string, + options?: { after?: number; limit?: number }, + ): Promise> + subscribe(threadId: string, listener: () => void): () => void +} +``` + +- `append` writes the batch at `seq`, `seq + 1`, and so on. The first record of a + thread is at `1`. +- The batch is atomic. Write all of it or none of it, in one transaction. +- `seq` must be the next free position. For a taken position or a gap, reject with + `LogConflictError` from `@tanstack/ai-persistence`, and write nothing. This is how + a writer finds a second writer on the same thread. +- An empty batch writes nothing. +- `read` returns the entries after `after`, in `seq` order, at most `limit`. A thread + with no records returns `[]`. Return copies, so a caller cannot change the log. +- `subscribe` calls `listener` after each append that the store can see. A store that + cannot see appends from other processes can poll, or call only for its own appends. + +In SQL, a primary key on `(thread_id, seq)` gives the conflict rule. Insert the batch +in one transaction, and map the key violation to `LogConflictError`: + +```sql +CREATE TABLE session_log ( + thread_id TEXT NOT NULL, + seq INTEGER NOT NULL, + record TEXT NOT NULL, + PRIMARY KEY (thread_id, seq) +); +``` + +A key violation covers a taken position. A gap needs one more check in the same +transaction: `seq` must equal `MAX(seq) + 1` for the thread, or `1`. + +Each record is JSON. Store it as text or as a JSON column. The store does not read +the records. + +`memoryLogStore()` is the reference store. It is not part of `memoryPersistence()`, +so add it yourself: + +```ts +import { + defineAIPersistence, + memoryLogStore, + memoryPersistence, +} from '@tanstack/ai-persistence' + +const persistence = defineAIPersistence({ + stores: { ...memoryPersistence().stores, log: memoryLogStore() }, +}) +``` + +## LeaseStore + +Your system can already have leases: a job queue, or a submission table that says +which worker owns a job. `LeaseStore` lets a host use them for each attempt of a turn. + +A host does four things with `leases`: + +- It takes a lease when a turn attempt starts. +- It renews the lease while the attempt runs. +- It releases the lease after the attempt ends. +- After a crash, it calls `isAlive` to find attempts that nobody runs any more. + +```ts +interface TurnLeaseKey { + threadId: string + inputId: string + operationId: string // the operation that runs the attempt, also the AG-UI run id + attempt: number // 1 for the first run of the input, then 1 more after each crash +} + +interface TurnLease extends TurnLeaseKey { + ownerId: string // the host that runs the attempt + expiresAt: number // epoch ms +} + +interface LeaseStore { + acquire(lease: TurnLease): Promise + renew(lease: TurnLease): Promise + release(lease: TurnLease): Promise + isAlive(key: TurnLeaseKey): Promise +} +``` + +Rules for an implementation: + +- `isAlive` returns `true` while another host still runs the attempt. That means a + lease exists and has not expired. +- `renew` and `release` for a lease you do not hold can do nothing. +- Key a lease by input id and attempt. Two attempts of one input are two leases. + +This store keeps leases in a `Map`: + +```ts +import type { + LeaseStore, + TurnLease, + TurnLeaseKey, +} from '@tanstack/ai-persistence' + +function createMemoryLeases(): LeaseStore { + const leases = new Map() + const keyOf = (key: TurnLeaseKey) => `${key.inputId}:${key.attempt}` + + return { + acquire: async (lease) => { + leases.set(keyOf(lease), lease) + }, + renew: async (lease) => { + leases.set(keyOf(lease), lease) + }, + release: async (lease) => { + leases.delete(keyOf(lease)) + }, + isAlive: async (key) => { + const lease = leases.get(keyOf(key)) + return lease !== undefined && lease.expiresAt >= Date.now() + }, + } +} +``` + +The conformance suite has no `leases` cases. The rules depend on the system that +issues the leases. + +## WorkClaimStore + +A host can stop while a thread still has work. Then another host must find that +thread and finish the work. A host claims a thread while the thread has work, and +renews the claim. When the host stops, the claim expires. Another host's sweep finds +the thread with `listExpired` and takes it over with `claim`. + +```ts +interface WorkClaimStore { + claim(entry: { + threadId: string + harness: string // the name of the setup that runs the thread + ownerId: string // the host that holds the claim + until: number // epoch ms + }): Promise + release(threadId: string, ownerId: string): Promise + listExpired(options: { + now: number // epoch ms + limit?: number + }): Promise> +} +``` + +Rules for an implementation: + +- `claim` is atomic. Two callers never both get `true` for one thread at one moment. +- `claim` returns `true` and sets `until` when nobody holds the thread, when + `ownerId` already holds it (a renewal), or when the claim expired (a takeover). +- `claim` returns `false` while another owner holds a claim that has not expired. +- `release` removes the claim only for its owner. For another `ownerId`, it does + nothing. +- `listExpired` returns the claims that expired before `now`, oldest first, at most + `limit`. Without `limit`, it returns all of them. + +In SQL, one row per thread and a conditional upsert give the atomic rule. The `WHERE` +clause refuses a live claim of another owner: + +```sql +CREATE TABLE work_claims ( + thread_id TEXT PRIMARY KEY, + harness TEXT NOT NULL, + owner_id TEXT NOT NULL, + until INTEGER NOT NULL +); + +INSERT INTO work_claims (thread_id, harness, owner_id, until) +VALUES (:thread_id, :harness, :owner_id, :until) +ON CONFLICT (thread_id) DO UPDATE + SET harness = excluded.harness, owner_id = excluded.owner_id, until = excluded.until + WHERE work_claims.owner_id = excluded.owner_id + OR work_claims.until <= :now; +``` + +The claim is `true` when the statement wrote one row. An index on `until` keeps +`listExpired` fast. + +`memoryPersistence()` includes a memory `workClaims` store. It does not survive a +restart. + +## SessionIndexStore + +Your app shows a list of sessions: the title, the owner, the time of the last change, +and the token totals. Reading every transcript to build that list is slow. +`SessionIndexStore` keeps one small entry per thread. The index holds no messages. +The data of each thread stays in the stores that hold it. + +```ts +interface SessionIndexEntry { + threadId: string + harness?: string // the name of the setup that runs the thread + title?: string + parentThreadId?: string // a child session, for example `subagent:` + parentToolCallId?: string // the tool call in the parent thread that started it + createdAt: number // epoch ms, the first open + updatedAt: number // epoch ms, the last change + pinned?: boolean + principal?: { id: string; tenantId?: string } // the owner + usage?: { + turns: number // model calls + promptTokens: number + completionTokens: number + totalTokens: number + cachedTokens: number + cacheWriteTokens: number + cost?: number // USD + } + metadata?: Record // app data, for example pinned files +} + +interface SessionIndexListOptions { + limit?: number + cursor?: string + parentThreadId?: string | null + principal?: { id: string; tenantId?: string } + search?: string + harness?: string + metadata?: Record +} + +interface SessionIndexPage { + entries: Array + cursor?: string // only when `truncated` + truncated?: boolean +} + +interface SessionIndexStore { + upsert(entry: SessionIndexEntry): Promise + get(threadId: string): Promise + list(options?: SessionIndexListOptions): Promise + delete(threadId: string): Promise + // Optional: compare-and-set. Add both, or neither. + getVersioned?( + threadId: string, + ): Promise<{ entry: SessionIndexEntry; revision: string } | undefined> + upsertIf?( + entry: SessionIndexEntry, + expectedRevision: string | null, + ): Promise<{ ok: true; revision: string } | { ok: false; reason: 'conflict' }> +} +``` + +Rules for `upsert`, `get`, and `delete`: + +- `upsert` replaces the whole entry for `entry.threadId`. A field that the new entry + does not have is gone. Do not merge. To change one field, read the entry and write + a changed copy. +- `get` returns `undefined` for an unknown thread, not `null`. +- `delete` removes the entry only. The messages, the log, and the other data of the + thread stay. For an unknown thread, `delete` does nothing. + +Rules for `list`: + +- Sort by `updatedAt`, the newest first. Entries with the same `updatedAt` go in + `threadId` order. `pinned` gets no special place. +- `parentThreadId: 'x'` returns only the entries whose parent is `x`. `null` returns + only the entries with no parent. +- `principal` returns only the entries of that owner id. With a `tenantId`, the + tenant must match too. An entry without a `principal` never matches. +- `search` returns only the entries whose `title` contains the text, without case. + `harness` and `metadata` match exact strings. All filters must match. +- With `limit`, when more entries match, return `truncated: true` and a `cursor`. The + next call with that `cursor` returns the entries after the last entry of the page. + +With `stores.sessions`, `withPersistence` writes an entry for each subagent child. The +entry has the parent thread, the tool call that started the child, and the owner and +`harness` of the parent entry. + +### Keep the writes of several hosts + +Two hosts can write one entry at the same time. For example, one host writes the +usage of a turn while another host writes the pinned files of the app. Without +`getVersioned` and `upsertIf`, the second write can remove the change of the first. + +Rules for `getVersioned` and `upsertIf`: + +- `getVersioned` returns the entry and its revision, or `undefined` for an unknown + thread. The format of the revision is yours. +- Each write gives the entry a new revision: `upsert` and `upsertIf` both. After + `delete`, the next write gets a new revision. The deleted revision is never used + again. +- `upsertIf` writes only when the current revision is `expectedRevision`, and returns + `{ ok: true, revision }`. Pass `null` to write only when the thread has no entry. +- On any other revision, `upsertIf` writes nothing and returns + `{ ok: false, reason: 'conflict' }`. + +In SQL, keep a revision column. Then: + +- For a revision, write with `UPDATE ... WHERE thread_id = ? AND revision = ?`. + When no row changed, return a conflict. +- For `null`, write with `INSERT ... ON CONFLICT DO NOTHING`. When no row was + added, return a conflict. + +`memoryPersistence()` has both methods. Its cursor is `:` of the +last entry of a page. + +### Check the session stores with the conformance suite + +`runPersistenceConformance` runs the cases of `credentials`, `inbox`, `workClaims`, +`log`, and `sessions` when your persistence has that store. A persistence without one +needs no `skip` entry for it. The compare-and-set cases run only when `sessions` has +`getVersioned` and `upsertIf`. + +To check one session store alone, skip the stores that the suite reads by default: + +```ts +import { runPersistenceConformance } from '@tanstack/ai-persistence/testkit' +import { mySessionIndex } from './session-index' + +runPersistenceConformance( + 'my session index', + () => ({ stores: { sessions: mySessionIndex() } }), + { + skip: [ + 'messages', + 'activities', + 'runs', + 'interrupts', + 'metadata', + 'generationRuns', + 'artifacts', + 'blobs', + ], + }, +) +``` + ## How the records relate The thread is not a table of its own: it exists as the `thread_id` the other records @@ -645,12 +1102,14 @@ erDiagram Each store has a `define*Store` helper (`defineMessageStore`, `defineRunStore`, `defineInterruptStore`, `defineMetadataStore`, `defineGenerationRunStore`, -`defineArtifactStore`, `defineBlobStore`). They compose into `defineAIPersistence`, -which tracks exact presence: the stores you passed autocomplete on -`persistence.stores`, and reading one you did not pass is a compile error. - -Those seven keys are the only ones `stores` accepts. Anything else throws -`Unknown AIPersistence store key` at construction. +`defineArtifactStore`, `defineBlobStore`, `defineInboxStore`, +`defineCredentialStore`, `defineLogStore`, `defineWorkClaimStore`, +`defineSessionIndexStore`). They compose into `defineAIPersistence`, which tracks +exact presence: the stores you passed autocomplete on `persistence.stores`, and +reading one you did not pass is a compile error. + +The keys in the table at the top, plus `activities`, are the only ones `stores` +accepts. Anything else throws `Unknown AIPersistence store key` at construction. To annotate the value instead, use a named shape: diff --git a/packages/ai-persistence/src/index.ts b/packages/ai-persistence/src/index.ts index 4440d726c3..74538c5710 100644 --- a/packages/ai-persistence/src/index.ts +++ b/packages/ai-persistence/src/index.ts @@ -10,6 +10,12 @@ export { defineGenerationRunStore, defineArtifactStore, defineBlobStore, + defineInboxStore, + defineCredentialStore, + defineLogStore, + defineWorkClaimStore, + defineSessionIndexStore, + LogConflictError, // Run lifecycle helpers owned by @tanstack/ai, transiting through ./types so // this package's public surface stays a single import for backend authors. isTerminalRunStatus, @@ -28,6 +34,23 @@ export type { InterruptStatus, InterruptStore, MetadataStore, + // Session store contracts + InboxEntry, + InboxStatus, + InboxStore, + Credential, + CredentialStore, + LogEntry, + LogRecord, + LogStore, + LeaseStore, + TurnLease, + TurnLeaseKey, + WorkClaimStore, + SessionIndexEntry, + SessionIndexListOptions, + SessionIndexPage, + SessionIndexStore, // Named product shapes (prefer these over a sparse bag) ChatTranscriptStores, ChatPersistenceStores, @@ -107,7 +130,7 @@ export { export { parseRangeHeader, resolveBlobRange } from './blob-range' // Reference in-memory implementation -export { memoryPersistence } from './memory' +export { memoryLogStore, memoryPersistence } from './memory' // Persistence-owned capabilities only. Locks: @tanstack/ai. export { diff --git a/packages/ai-persistence/src/memory.ts b/packages/ai-persistence/src/memory.ts index 153df924b1..03dcf2cd9f 100644 --- a/packages/ai-persistence/src/memory.ts +++ b/packages/ai-persistence/src/memory.ts @@ -1,4 +1,4 @@ -import { defineAIPersistence } from './types' +import { LogConflictError, defineAIPersistence } from './types' import { resolveBlobRange } from './blob-range' import type { ActivityRecord, ModelMessage } from '@tanstack/ai' import type { @@ -13,15 +13,28 @@ import type { BlobRange, BlobRecord, BlobStore, + Credential, + CredentialStore, GenerationRunRecord, GenerationRunStore, + InboxEntry, + InboxStore, InterruptCommitEntry, InterruptRecord, InterruptStore, + LogEntry, + LogRecord, + LogStore, MessageStore, MetadataStore, RunRecord, RunStore, + Scope, + SessionIndexEntry, + SessionIndexListOptions, + SessionIndexPage, + SessionIndexStore, + WorkClaimStore, } from './types' const compareUtf8Bytes = (left: string, right: string): number => { @@ -591,6 +604,321 @@ class MemoryBlobStore implements BlobStore { } } +class MemoryInboxStore implements InboxStore { + private readonly entries = new Map() + append(entry: Omit): Promise { + // Copies in and out, so a caller never changes a stored entry. + const existing = this.entries.get(entry.inputId) + if (existing) return Promise.resolve(structuredClone(existing)) + const stored: InboxEntry = { ...entry, status: 'pending' } + this.entries.set(entry.inputId, structuredClone(stored)) + return Promise.resolve(stored) + } + listPending(threadId: string): Promise> { + // Map iteration keeps insertion order. The stable sort keeps it for ties. + const pending = [...this.entries.values()] + .filter( + (entry) => entry.threadId === threadId && entry.status === 'pending', + ) + .sort((a, b) => a.createdAt - b.createdAt) + .map((entry) => structuredClone(entry)) + return Promise.resolve(pending) + } + markApplied(inputId: string, operationId: string): Promise { + const existing = this.entries.get(inputId) + if (existing) { + this.entries.set(inputId, { ...existing, status: 'applied', operationId }) + } + return Promise.resolve() + } + markRejected(inputId: string, reason: string): Promise { + const existing = this.entries.get(inputId) + if (existing) { + this.entries.set(inputId, { ...existing, status: 'rejected', reason }) + } + return Promise.resolve() + } + get(inputId: string): Promise { + const existing = this.entries.get(inputId) + return Promise.resolve(existing ? structuredClone(existing) : null) + } +} + +class MemoryCredentialStore implements CredentialStore { + private readonly values = new Map>() + // A credential without a userId belongs to the tenant. + private owner(scope: Scope): string { + return JSON.stringify([scope.tenantId ?? null, scope.userId ?? null]) + } + get(scope: Scope, id: string): Promise { + const stored = this.values.get(this.owner(scope))?.get(id) + return Promise.resolve(stored ? structuredClone(stored) : null) + } + set(scope: Scope, id: string, credential: Credential): Promise { + const key = this.owner(scope) + let bucket = this.values.get(key) + if (!bucket) { + bucket = new Map() + this.values.set(key, bucket) + } + bucket.set(id, structuredClone(credential)) + return Promise.resolve() + } + delete(scope: Scope, id: string): Promise { + this.values.get(this.owner(scope))?.delete(id) + return Promise.resolve() + } + list( + scope: Scope, + ): Promise< + Array<{ id: string; type: Credential['type']; expiresAt?: number }> + > { + const bucket = this.values.get(this.owner(scope)) + return Promise.resolve( + [...(bucket?.entries() ?? [])].map(([id, credential]) => ({ + id, + type: credential.type, + ...(credential.type === 'oauth' && credential.expiresAt !== undefined + ? { expiresAt: credential.expiresAt } + : {}), + })), + ) + } +} + +/** A JSON copy, so the log never shares an object with a caller. */ +function copyRecord(record: LogRecord): LogRecord { + return JSON.parse(JSON.stringify(record)) +} + +class MemoryLogStore implements LogStore { + private readonly threads = new Map>() + private readonly listeners = new Map void>>() + + append( + threadId: string, + seq: number, + records: ReadonlyArray, + ): Promise { + if (records.length === 0) return Promise.resolve() + const log = this.threads.get(threadId) ?? [] + if (seq !== log.length + 1) { + return Promise.reject(new LogConflictError(threadId, seq)) + } + // Copy the whole batch before the push, so a record that cannot become + // JSON rejects the batch with nothing written. + let copies: Array + try { + copies = records.map(copyRecord) + } catch (error) { + return Promise.reject(error) + } + log.push(...copies) + this.threads.set(threadId, log) + for (const listener of [...(this.listeners.get(threadId) ?? [])]) { + // The batch is written. A listener error must not fail the append. + try { + listener() + } catch (error) { + console.error('[ai-persistence] log listener failed', error) + } + } + return Promise.resolve() + } + + read( + threadId: string, + options?: { after?: number; limit?: number }, + ): Promise> { + const log = this.threads.get(threadId) ?? [] + const after = options?.after ?? 0 + const end = + options?.limit === undefined ? log.length : after + options.limit + const entries = log.slice(after, end).map((record, index) => ({ + seq: after + index + 1, + record: copyRecord(record), + })) + return Promise.resolve(entries) + } + + subscribe(threadId: string, listener: () => void): () => void { + let set = this.listeners.get(threadId) + if (!set) { + set = new Set() + this.listeners.set(threadId, set) + } + set.add(listener) + return () => { + set.delete(listener) + } + } +} + +/** + * In-process reference {@link LogStore}, for tests and one-process hosts. The + * log is lost when the process stops. Records are JSON-copied in and out. + * + * Not part of `memoryPersistence()`: add it yourself. + * + * @example + * ```ts + * const persistence = defineAIPersistence({ + * stores: { ...memoryPersistence().stores, log: memoryLogStore() }, + * }) + * ``` + */ +export function memoryLogStore(): LogStore { + return new MemoryLogStore() +} + +type WorkClaim = Parameters[0] + +class MemoryWorkClaimStore implements WorkClaimStore { + private readonly claims = new Map() + claim(entry: WorkClaim): Promise { + // Synchronous, so two claims in one process never both win. + const held = this.claims.get(entry.threadId) + const isTaken = + held !== undefined && + held.ownerId !== entry.ownerId && + held.until > Date.now() + if (isTaken) return Promise.resolve(false) + this.claims.set(entry.threadId, { ...entry }) + return Promise.resolve(true) + } + release(threadId: string, ownerId: string): Promise { + if (this.claims.get(threadId)?.ownerId === ownerId) { + this.claims.delete(threadId) + } + return Promise.resolve() + } + listExpired(options: { + now: number + limit?: number + }): Promise> { + const expired = [...this.claims.values()] + .filter((claim) => claim.until <= options.now) + .sort((a, b) => a.until - b.until) + .slice(0, options.limit) + .map(({ threadId, harness }) => ({ threadId, harness })) + return Promise.resolve(expired) + } +} + +// Newest `updatedAt` first, then `threadId` order. The cursor uses it too. +const bySessionOrder = ( + a: Pick, + b: Pick, +) => b.updatedAt - a.updatedAt || compareUtf8Bytes(a.threadId, b.threadId) + +// The cursor is `:` of the last entry of a page. +function parseSessionCursor(cursor: string) { + const split = cursor.indexOf(':') + const updatedAt = Number(cursor.slice(0, split)) + if (split <= 0 || !Number.isFinite(updatedAt)) { + throw new Error(`Invalid session index cursor: ${cursor}`) + } + return { updatedAt, threadId: cursor.slice(split + 1) } +} + +class MemorySessionIndexStore implements SessionIndexStore { + private readonly entries = new Map< + string, + { entry: SessionIndexEntry; revision: string } + >() + // One counter for the whole store, so an entry that is deleted and written + // again never gets a revision that an earlier reader saw. + private lastRevision = 0 + upsert(entry: SessionIndexEntry): Promise { + this.write(entry) + return Promise.resolve() + } + get(threadId: string): Promise { + const stored = this.entries.get(threadId) + return Promise.resolve(stored && structuredClone(stored.entry)) + } + getVersioned( + threadId: string, + ): Promise<{ entry: SessionIndexEntry; revision: string } | undefined> { + const stored = this.entries.get(threadId) + return Promise.resolve(stored && structuredClone(stored)) + } + upsertIf( + entry: SessionIndexEntry, + expectedRevision: string | null, + ): Promise< + { ok: true; revision: string } | { ok: false; reason: 'conflict' } + > { + const current = this.entries.get(entry.threadId)?.revision ?? null + if (current !== expectedRevision) { + return Promise.resolve({ ok: false, reason: 'conflict' }) + } + return Promise.resolve({ ok: true, revision: this.write(entry) }) + } + list(options: SessionIndexListOptions = {}): Promise { + const { limit, cursor, parentThreadId, principal, search, harness } = + options + const after = cursor === undefined ? undefined : parseSessionCursor(cursor) + const needle = search?.toLowerCase() + const metadata = Object.entries(options.metadata ?? {}) + const matching = [...this.entries.values()] + .map((stored) => stored.entry) + .filter((entry) => { + // `null` asks for the entries with no parent. + const isChild = + parentThreadId === undefined || + (entry.parentThreadId ?? null) === parentThreadId + const isSameTenant = + principal?.tenantId === undefined || + entry.principal?.tenantId === principal.tenantId + const isOwned = + principal === undefined || + (entry.principal?.id === principal.id && isSameTenant) + const isAfterCursor = + after === undefined || bySessionOrder(after, entry) < 0 + const isFound = + needle === undefined || + (entry.title?.toLowerCase().includes(needle) ?? false) + const isHarness = harness === undefined || entry.harness === harness + const hasMetadata = metadata.every( + ([key, value]) => entry.metadata?.[key] === value, + ) + return ( + isChild && + isOwned && + isFound && + isHarness && + hasMetadata && + isAfterCursor + ) + }) + .sort(bySessionOrder) + const page = limit === undefined ? matching : matching.slice(0, limit) + const last = page.at(-1) + const hasMore = + limit !== undefined && matching.length > limit && last !== undefined + return Promise.resolve({ + entries: page.map((entry) => structuredClone(entry)), + ...(hasMore + ? { cursor: `${last.updatedAt}:${last.threadId}`, truncated: true } + : {}), + }) + } + delete(threadId: string): Promise { + this.entries.delete(threadId) + return Promise.resolve() + } + // Copies in and out, so a caller never changes a stored entry. + private write(entry: SessionIndexEntry): string { + const revision = String(++this.lastRevision) + this.entries.set(entry.threadId, { + entry: structuredClone(entry), + revision, + }) + return revision + } +} + interface MemoryPersistenceStores { messages: MessageStore activities: ActivityStore @@ -600,15 +928,21 @@ interface MemoryPersistenceStores { metadata: MetadataStore artifacts: ArtifactStore blobs: BlobStore + inbox: InboxStore + credentials: CredentialStore + workClaims: WorkClaimStore + sessions: SessionIndexStore } /** * In-process reference backend for the full state + generation store set. * * Returns messages + activities + runs + generationRuns + interrupts + - * metadata + artifacts + blobs. Locks are not included — use - * `InMemoryLockStore` + `withLocks` from `@tanstack/ai` when a test or - * single-process app needs coordination. + * metadata + artifacts + blobs + inbox + credentials + workClaims + sessions. + * The session log is not included: a host can turn on more behavior when + * `stores.log` is present, so add `memoryLogStore()` yourself. Locks are not + * included either. Use `InMemoryLockStore` + `withLocks` from `@tanstack/ai` + * when a test or single-process app needs coordination. */ export function memoryPersistence() { const stores: MemoryPersistenceStores = { @@ -620,6 +954,10 @@ export function memoryPersistence() { metadata: new MemoryMetadataStore(), artifacts: new MemoryArtifactStore(), blobs: new MemoryBlobStore(), + inbox: new MemoryInboxStore(), + credentials: new MemoryCredentialStore(), + workClaims: new MemoryWorkClaimStore(), + sessions: new MemorySessionIndexStore(), } return defineAIPersistence({ stores }) } diff --git a/packages/ai-persistence/src/middleware.ts b/packages/ai-persistence/src/middleware.ts index d71982be7b..1c8d654d1b 100644 --- a/packages/ai-persistence/src/middleware.ts +++ b/packages/ai-persistence/src/middleware.ts @@ -2100,6 +2100,9 @@ export function withPersistence( ...(wantsInterrupts && persistence.stores.interrupts ? { interrupts: persistence.stores.interrupts } : {}), + ...(persistence.stores.sessions + ? { sessions: persistence.stores.sessions } + : {}), }) function resetStreaming(state: RunStateEntry) { diff --git a/packages/ai-persistence/src/subagent-runs.ts b/packages/ai-persistence/src/subagent-runs.ts index e6f3d5be51..c968fdbac4 100644 --- a/packages/ai-persistence/src/subagent-runs.ts +++ b/packages/ai-persistence/src/subagent-runs.ts @@ -17,7 +17,7 @@ import type { UIMessage, } from '@tanstack/ai' import { mergeStoredMessages } from './merge-stored' -import type { InterruptStore, MessageStore } from './types' +import type { InterruptStore, MessageStore, SessionIndexStore } from './types' function withSubagentInfo( metadata: ModelMessage['metadata'], @@ -255,6 +255,8 @@ export function createSubagentRunRecorder(stores: { messages: MessageStore runs?: RunStore interrupts?: InterruptStore + /** When present, each child gets a session index entry. */ + sessions?: SessionIndexStore /** Minimum milliseconds between writes while a child streams. */ intervalMs?: number }) { @@ -478,6 +480,48 @@ export function createSubagentRunRecorder(stores: { children.get(chunk.parentSubagentRunId)?.processor.processChunk(chunk) } await saveChild(id) + const note = children.get(id) + if (stores.sessions && note) await indexChild(stores.sessions, id, note) + } + + // One index entry per child. `upsert` replaces the whole entry, so keep the + // fields that other writers set, like a title. The child has the owner and + // the `harness` name of its parent thread, so a list filtered by them shows + // it. ponytail: get-then-upsert, use `upsertIf` if two hosts write one child. + // Best-effort: the index is a side store, so a failed write never fails the + // run. + async function indexChild( + sessions: SessionIndexStore, + subagentRunId: string, + note: ChildNote, + ) { + const childThreadId = childStoreId(subagentRunId) + try { + const [current, parent] = await Promise.all([ + sessions.get(childThreadId), + sessions.get(note.parentThreadId), + ]) + const now = Date.now() + await sessions.upsert({ + ...current, + threadId: childThreadId, + parentThreadId: note.parentThreadId, + ...(note.parentToolCallId !== undefined && { + parentToolCallId: note.parentToolCallId, + }), + ...(parent?.principal !== undefined && { + principal: parent.principal, + }), + ...(parent?.harness !== undefined && { harness: parent.harness }), + createdAt: current?.createdAt ?? now, + updatedAt: now, + }) + } catch (error) { + console.warn( + `[ai-persistence] session index write failed for '${childThreadId}'`, + error, + ) + } } /** diff --git a/packages/ai-persistence/src/testkit/conformance.ts b/packages/ai-persistence/src/testkit/conformance.ts index 7b958783d7..fae36d9de0 100644 --- a/packages/ai-persistence/src/testkit/conformance.ts +++ b/packages/ai-persistence/src/testkit/conformance.ts @@ -7,11 +7,14 @@ * store the persistence exposes and is the authoritative compatibility gate for * the store interfaces in `../types.ts`. * - * Covers all eight stores: the five chat state stores (`messages`, - * `activities`, `runs`, `interrupts`, `metadata`) and the three generation - * stores (`generationRuns`, `artifacts`, `blobs`). Locks are not part of this - * suite — they are a separate coordination concern (`LockStore` + `withLocks`), - * not a store. + * Covers the five chat state stores (`messages`, `activities`, `runs`, + * `interrupts`, `metadata`), the three generation stores (`generationRuns`, + * `artifacts`, `blobs`), and the opt-in session stores (`credentials`, `inbox`, + * `workClaims`, `log`, `sessions`). The session store cases run only when the + * store is present, so they need no `skip` entry. `leases` has no cases: its + * rules depend on the system that issues the leases. Locks are not part of + * this suite — they are a separate coordination concern (`LockStore` + + * `withLocks`), not a store. * * SKIPPING (declare or fail): a backend that deliberately omits a store must * declare it in `options.skip`, and one that omits an OPTIONAL store method @@ -39,12 +42,16 @@ * starting with `rc-`, or it silently changes that expected set. */ import { beforeAll, describe, expect, it } from 'vitest' +import { LogConflictError } from '../types' import type { ModelMessage } from '@tanstack/ai' import type { AIPersistence, AIPersistenceStores, ArtifactRecord, + LogStore, RunStore, + SessionIndexEntry, + SessionIndexPage, } from '../types' type MakePersistence = () => Promise | AIPersistence @@ -1573,6 +1580,759 @@ export function runPersistenceConformance( }) }) + // The session stores below are opt-in. A backend without one skips its + // cases and needs no `skip` entry. + describe('credentials', () => { + it('keys by user and tenant, lists without values, and deletes', async (ctx) => { + const store = persistence.stores.credentials + if (!store) return ctx.skip('credentials store not provided') + + const user = { threadId: 't', userId: 'u1', tenantId: 'org' } + const tenant = { threadId: 't', tenantId: 'org' } + await store.set(user, 'github', { + type: 'oauth', + accessToken: 'secret-access', + refreshToken: 'secret-refresh', + expiresAt: 100, + }) + await store.set(tenant, 'notion', { + type: 'api_key', + value: 'secret-key', + }) + + expect(await store.get(user, 'github')).toMatchObject({ + accessToken: 'secret-access', + }) + // A user credential is not visible at tenant scope, and the other way around. + expect(await store.get(tenant, 'github')).toBeNull() + expect(await store.get(user, 'notion')).toBeNull() + + const listed = await store.list(user) + expect(listed).toEqual([ + { id: 'github', type: 'oauth', expiresAt: 100 }, + ]) + // `list` never returns secret values. + expect(JSON.stringify(listed)).not.toContain('secret') + + await store.delete(user, 'github') + expect(await store.get(user, 'github')).toBeNull() + expect(await store.list(user)).toEqual([]) + }) + }) + + describe('inbox', () => { + it('appends idempotently, lists pending oldest first, and settles entries', async (ctx) => { + const store = persistence.stores.inbox + if (!store) return ctx.skip('inbox store not provided') + // Own ids, because the suite shares one persistence. + const id = (name: string) => `${name}-${crypto.randomUUID()}` + const thread = id('inbox-thread') + const [in1, in2, other, missing] = [ + 'in-1', + 'in-2', + 'in-other', + 'missing', + ].map(id) + + const first = await store.append({ + inputId: in1, + threadId: thread, + input: { op: 'prompt', message: 'hi' }, + createdAt: 1, + }) + expect(first.status).toBe('pending') + // Same id again: the stored entry wins, the new payload is ignored. + const again = await store.append({ + inputId: in1, + threadId: thread, + input: { op: 'prompt', message: 'changed' }, + createdAt: 5, + }) + expect(again.input).toEqual({ op: 'prompt', message: 'hi' }) + + await store.append({ + inputId: in2, + threadId: thread, + principal: { id: 'user-1' }, + input: { op: 'steer', message: 'shorter' }, + createdAt: 2, + }) + await store.append({ + inputId: other, + threadId: id('other-thread'), + input: { op: 'prompt', message: 'x' }, + createdAt: 0, + }) + + expect( + (await store.listPending(thread)).map((entry) => entry.inputId), + ).toEqual([in1, in2]) + + await store.markApplied(in1, 'op-1') + await store.markRejected(in2, 'busy') + expect(await store.listPending(thread)).toEqual([]) + expect(await store.get(in1)).toMatchObject({ + status: 'applied', + operationId: 'op-1', + }) + expect(await store.get(in2)).toMatchObject({ + status: 'rejected', + reason: 'busy', + principal: { id: 'user-1' }, + }) + + // Unknown ids are no-ops. + await store.markApplied(missing, 'op-x') + await store.markRejected(missing, 'nope') + expect(await store.get(missing)).toBeNull() + }) + }) + + // Each case uses its own threads, because the suite shares one persistence. + describe('workClaims', () => { + const newThread = () => `claims-${crypto.randomUUID()}` + const later = () => Date.now() + 60 * 60_000 + + it('claims, renews, refuses a second owner, and releases by the owner only', async (ctx) => { + const store = persistence.stores.workClaims + if (!store) return ctx.skip('workClaims store not provided') + const threadId = newThread() + const claim = (ownerId: string) => + store.claim({ threadId, harness: 'h', ownerId, until: later() }) + + expect(await claim('a')).toBe(true) + // The owner renews. + expect(await claim('a')).toBe(true) + // Another owner is refused while the claim is live. + expect(await claim('b')).toBe(false) + // Only the owner releases. + await store.release(threadId, 'b') + expect(await claim('b')).toBe(false) + await store.release(threadId, 'a') + expect(await claim('b')).toBe(true) + await store.release(threadId, 'b') + }) + + it('gives one of two concurrent claims', async (ctx) => { + const store = persistence.stores.workClaims + if (!store) return ctx.skip('workClaims store not provided') + const threadId = newThread() + const results = await Promise.all( + ['a', 'b'].map((ownerId) => + store.claim({ threadId, harness: 'h', ownerId, until: later() }), + ), + ) + expect(results.filter(Boolean)).toHaveLength(1) + await Promise.all( + ['a', 'b'].map((owner) => store.release(threadId, owner)), + ) + }) + + it('lists expired claims oldest first, with a limit, and lets another owner take one over', async (ctx) => { + const store = persistence.stores.workClaims + if (!store) return ctx.skip('workClaims store not provided') + const now = Date.now() + const older = newThread() + const newer = newThread() + const live = newThread() + await store.claim({ + threadId: newer, + harness: 'h', + ownerId: 'gone', + until: now - 1_000, + }) + await store.claim({ + threadId: older, + harness: 'h2', + ownerId: 'gone', + until: now - 2_000, + }) + await store.claim({ + threadId: live, + harness: 'h', + ownerId: 'alive', + until: later(), + }) + + const ours = (await store.listExpired({ now })).filter((entry) => + [older, newer, live].includes(entry.threadId), + ) + expect(ours).toEqual([ + { threadId: older, harness: 'h2' }, + { threadId: newer, harness: 'h' }, + ]) + expect(await store.listExpired({ now, limit: 1 })).toHaveLength(1) + + // A sweep takes an expired claim over. + expect( + await store.claim({ + threadId: older, + harness: 'h2', + ownerId: 'sweeper', + until: later(), + }), + ).toBe(true) + expect( + (await store.listExpired({ now })).map((entry) => entry.threadId), + ).not.toContain(older) + + await store.release(older, 'sweeper') + await store.release(newer, 'gone') + await store.release(live, 'alive') + }) + }) + + // Each case uses its own thread, because the suite shares one persistence. + describe('log', () => { + const newThread = () => `log-${crypto.randomUUID()}` + const types = async (store: LogStore, threadId: string) => + (await store.read(threadId)).map((entry) => [ + entry.seq, + entry.record.type, + ]) + + it('writes a batch at consecutive positions and reads it in order', async (ctx) => { + const store = persistence.stores.log + if (!store) return ctx.skip('log store not provided') + const threadId = newThread() + + await store.append(threadId, 1, [ + { type: 'a', n: 1 }, + { type: 'b', n: 2 }, + ]) + await store.append(threadId, 3, [{ type: 'c', nested: { deep: [1] } }]) + + expect(await store.read(threadId)).toEqual([ + { seq: 1, record: { type: 'a', n: 1 } }, + { seq: 2, record: { type: 'b', n: 2 } }, + { seq: 3, record: { type: 'c', nested: { deep: [1] } } }, + ]) + }) + + it('rejects a taken position or a gap with LogConflictError and writes nothing', async (ctx) => { + const store = persistence.stores.log + if (!store) return ctx.skip('log store not provided') + const threadId = newThread() + await store.append(threadId, 1, [{ type: 'first' }]) + + const taken = store.append(threadId, 1, [ + { type: 'late-1' }, + { type: 'late-2' }, + ]) + await expect(taken).rejects.toBeInstanceOf(LogConflictError) + await expect(taken).rejects.toMatchObject({ threadId, seq: 1 }) + + // A gap (position 3 while 2 is free) is a conflict too. + await expect( + store.append(threadId, 3, [{ type: 'gap' }]), + ).rejects.toBeInstanceOf(LogConflictError) + + expect(await types(store, threadId)).toEqual([[1, 'first']]) + }) + + it('writes nothing for an empty batch', async (ctx) => { + const store = persistence.stores.log + if (!store) return ctx.skip('log store not provided') + const threadId = newThread() + + await store.append(threadId, 1, []) + + expect(await store.read(threadId)).toEqual([]) + await store.append(threadId, 1, [{ type: 'still-first' }]) + expect(await types(store, threadId)).toEqual([[1, 'still-first']]) + }) + + it('reads after a position, with a limit', async (ctx) => { + const store = persistence.stores.log + if (!store) return ctx.skip('log store not provided') + const threadId = newThread() + await store.append(threadId, 1, [ + { type: 'r1' }, + { type: 'r2' }, + { type: 'r3' }, + { type: 'r4' }, + ]) + + const page = await store.read(threadId, { after: 1, limit: 2 }) + expect(page.map((entry) => entry.seq)).toEqual([2, 3]) + expect( + (await store.read(threadId, { after: 3 })).map((entry) => entry.seq), + ).toEqual([4]) + expect(await store.read(threadId, { after: 4 })).toEqual([]) + expect(await store.read(threadId, { limit: 0 })).toEqual([]) + expect(await store.read(newThread())).toEqual([]) + }) + + it('keeps threads apart', async (ctx) => { + const store = persistence.stores.log + if (!store) return ctx.skip('log store not provided') + const left = newThread() + const right = newThread() + + await store.append(left, 1, [{ type: 'left' }]) + // Position 1 is free on the other thread. + await store.append(right, 1, [{ type: 'right' }]) + + expect(await types(store, left)).toEqual([[1, 'left']]) + expect(await types(store, right)).toEqual([[1, 'right']]) + }) + + it('calls a listener after each append until it unsubscribes', async (ctx) => { + const store = persistence.stores.log + if (!store) return ctx.skip('log store not provided') + const threadId = newThread() + let calls = 0 + const stop = store.subscribe(threadId, () => { + calls += 1 + }) + + await store.append(threadId, 1, [{ type: 'one' }]) + await expect.poll(() => calls).toBe(1) + stop() + await store.append(threadId, 2, [{ type: 'two' }]) + + // Give a store that notifies later (a poller) the time to call again. + await new Promise((resolve) => setTimeout(resolve, 50)) + expect(calls).toBe(1) + }) + + it('does not share record objects with the caller', async (ctx) => { + const store = persistence.stores.log + if (!store) return ctx.skip('log store not provided') + const threadId = newThread() + const written = { type: 'shared', list: [1] } + await store.append(threadId, 1, [written]) + written.list.push(2) + + const [first] = await store.read(threadId) + if (first) first.record.type = 'changed' + + expect(await store.read(threadId)).toEqual([ + { seq: 1, record: { type: 'shared', list: [1] } }, + ]) + }) + }) + + // Each case uses its own ids, because the suite shares one persistence. + describe('sessions', () => { + const newId = () => `sessions-${crypto.randomUUID()}` + const sessionEntry = ( + threadId: string, + updatedAt: number, + extra: Partial = {}, + ) => ({ threadId, createdAt: 1, updatedAt, ...extra }) + const threadIds = (page: SessionIndexPage) => + page.entries.map((entry) => entry.threadId) + + it('gets every field of an upserted entry, and undefined for an unknown thread', async (ctx) => { + const store = persistence.stores.sessions + if (!store) return ctx.skip('sessions store not provided') + const threadId = newId() + const written: SessionIndexEntry = { + threadId, + harness: 'coder', + title: 'Fix the build', + parentThreadId: `${threadId}-parent`, + parentToolCallId: 'call-1', + createdAt: 1, + updatedAt: 2, + pinned: true, + principal: { id: 'user-1', tenantId: 'org-1' }, + usage: { + turns: 2, + promptTokens: 10, + completionTokens: 5, + totalTokens: 15, + cachedTokens: 4, + cacheWriteTokens: 1, + cost: 0.25, + }, + metadata: { model: 'm-1' }, + } + + await store.upsert(written) + + expect(await store.get(threadId)).toEqual(written) + expect(await store.get(newId())).toBeUndefined() + }) + + it('replaces the whole entry on a second upsert', async (ctx) => { + const store = persistence.stores.sessions + if (!store) return ctx.skip('sessions store not provided') + const threadId = newId() + await store.upsert( + sessionEntry(threadId, 2, { title: 'old', pinned: true }), + ) + + await store.upsert(sessionEntry(threadId, 3, { title: 'new' })) + + // `pinned` is gone: upsert does not merge. + expect(await store.get(threadId)).toEqual({ + threadId, + createdAt: 1, + updatedAt: 3, + title: 'new', + }) + }) + + it('lists newest first and pages with a cursor', async (ctx) => { + const store = persistence.stores.sessions + if (!store) return ctx.skip('sessions store not provided') + const owner = { id: newId() } + const a = `${owner.id}-a` + const b = `${owner.id}-b` + const c = `${owner.id}-c` + await store.upsert(sessionEntry(c, 20, { principal: owner })) + await store.upsert(sessionEntry(a, 30, { principal: owner })) + await store.upsert(sessionEntry(b, 20, { principal: owner })) + + // Newest first. The same `updatedAt` goes in `threadId` order. + expect(threadIds(await store.list({ principal: owner }))).toEqual([ + a, + b, + c, + ]) + + const first = await store.list({ principal: owner, limit: 2 }) + expect(threadIds(first)).toEqual([a, b]) + expect(first.truncated).toBe(true) + + const second = await store.list({ + principal: owner, + limit: 2, + cursor: required(first.cursor, 'the cursor of a truncated page'), + }) + expect(threadIds(second)).toEqual([c]) + expect(second.truncated).toBeFalsy() + expect(second.cursor).toBeUndefined() + }) + + it('filters by parent thread and by owner', async (ctx) => { + const store = persistence.stores.sessions + if (!store) return ctx.skip('sessions store not provided') + const prefix = newId() + const owner = { id: `${prefix}-user`, tenantId: 'org-1' } + const parent = `${prefix}-parent` + const child = `${prefix}-child` + const unowned = `${prefix}-unowned` + const otherTenant = `${prefix}-other-tenant` + await store.upsert(sessionEntry(parent, 10, { principal: owner })) + await store.upsert( + sessionEntry(child, 20, { + principal: owner, + parentThreadId: parent, + parentToolCallId: 'call-1', + }), + ) + await store.upsert( + sessionEntry(unowned, 30, { parentThreadId: parent }), + ) + await store.upsert( + sessionEntry(otherTenant, 40, { + principal: { id: owner.id, tenantId: 'org-2' }, + }), + ) + + expect(threadIds(await store.list({ parentThreadId: parent }))).toEqual( + [unowned, child], + ) + // With a tenant, the tenant must match. An entry without an owner + // never matches. + expect(threadIds(await store.list({ principal: owner }))).toEqual([ + child, + parent, + ]) + // Without a tenant, the id alone decides. + expect( + threadIds(await store.list({ principal: { id: owner.id } })), + ).toEqual([otherTenant, child, parent]) + expect( + threadIds( + await store.list({ + principal: { id: `${prefix}-nobody` }, + parentThreadId: parent, + }), + ), + ).toEqual([]) + }) + + it('lists only the entries with no parent for parentThreadId null', async (ctx) => { + const store = persistence.stores.sessions + if (!store) return ctx.skip('sessions store not provided') + const owner = { id: newId() } + const parent = `${owner.id}-parent` + const child = `${owner.id}-child` + await store.upsert(sessionEntry(parent, 10, { principal: owner })) + await store.upsert( + sessionEntry(child, 20, { principal: owner, parentThreadId: parent }), + ) + + expect( + threadIds( + await store.list({ principal: owner, parentThreadId: null }), + ), + ).toEqual([parent]) + }) + + it('filters by title search, without case', async (ctx) => { + const store = persistence.stores.sessions + if (!store) return ctx.skip('sessions store not provided') + const owner = { id: newId() } + const build = `${owner.id}-build` + const docs = `${owner.id}-docs` + const untitled = `${owner.id}-untitled` + await store.upsert( + sessionEntry(build, 10, { principal: owner, title: 'Fix the Build' }), + ) + await store.upsert( + sessionEntry(docs, 20, { principal: owner, title: 'Write docs' }), + ) + await store.upsert(sessionEntry(untitled, 30, { principal: owner })) + + expect( + threadIds(await store.list({ principal: owner, search: 'BUILD' })), + ).toEqual([build]) + // An entry with no title never matches. + expect( + threadIds(await store.list({ principal: owner, search: '' })), + ).toEqual([docs, build]) + }) + + it('filters by harness name', async (ctx) => { + const store = persistence.stores.sessions + if (!store) return ctx.skip('sessions store not provided') + const owner = { id: newId() } + const coder = `${owner.id}-coder` + const writer = `${owner.id}-writer` + const none = `${owner.id}-none` + await store.upsert( + sessionEntry(coder, 10, { principal: owner, harness: 'coder' }), + ) + await store.upsert( + sessionEntry(writer, 20, { principal: owner, harness: 'writer' }), + ) + await store.upsert(sessionEntry(none, 30, { principal: owner })) + + expect( + threadIds(await store.list({ principal: owner, harness: 'coder' })), + ).toEqual([coder]) + }) + + it('filters by exact metadata values', async (ctx) => { + const store = persistence.stores.sessions + if (!store) return ctx.skip('sessions store not provided') + const owner = { id: newId() } + const both = `${owner.id}-both` + const projectOnly = `${owner.id}-project-only` + const other = `${owner.id}-other` + const none = `${owner.id}-none` + await store.upsert( + sessionEntry(both, 10, { + principal: owner, + metadata: { project: 'p-1', branch: 'main' }, + }), + ) + await store.upsert( + sessionEntry(projectOnly, 20, { + principal: owner, + metadata: { project: 'p-1' }, + }), + ) + await store.upsert( + sessionEntry(other, 30, { + principal: owner, + metadata: { project: 'p-2', branch: 'main' }, + }), + ) + await store.upsert(sessionEntry(none, 40, { principal: owner })) + + expect( + threadIds( + await store.list({ + principal: owner, + metadata: { project: 'p-1' }, + }), + ), + ).toEqual([projectOnly, both]) + expect( + threadIds( + await store.list({ + principal: owner, + metadata: { project: 'p-1', branch: 'main' }, + }), + ), + ).toEqual([both]) + }) + + it('combines search, harness, and metadata with AND', async (ctx) => { + const store = persistence.stores.sessions + if (!store) return ctx.skip('sessions store not provided') + const owner = { id: newId() } + const match = `${owner.id}-match` + const wrongTitle = `${owner.id}-wrong-title` + const wrongHarness = `${owner.id}-wrong-harness` + const wrongProject = `${owner.id}-wrong-project` + const base = { + principal: owner, + title: 'Fix the build', + harness: 'coder', + metadata: { project: 'p-1' }, + } + await store.upsert(sessionEntry(match, 10, base)) + await store.upsert( + sessionEntry(wrongTitle, 20, { ...base, title: 'Write docs' }), + ) + await store.upsert( + sessionEntry(wrongHarness, 30, { ...base, harness: 'writer' }), + ) + await store.upsert( + sessionEntry(wrongProject, 40, { + ...base, + metadata: { project: 'p-2' }, + }), + ) + + expect( + threadIds( + await store.list({ + principal: owner, + search: 'build', + harness: 'coder', + metadata: { project: 'p-1' }, + }), + ), + ).toEqual([match]) + }) + + it('pages the filtered list with a cursor', async (ctx) => { + const store = persistence.stores.sessions + if (!store) return ctx.skip('sessions store not provided') + const owner = { id: newId() } + const a = `${owner.id}-a` + const b = `${owner.id}-b` + const c = `${owner.id}-c` + await store.upsert( + sessionEntry(a, 50, { principal: owner, title: 'build a' }), + ) + await store.upsert( + sessionEntry(`${owner.id}-skip`, 40, { + principal: owner, + title: 'docs', + }), + ) + await store.upsert( + sessionEntry(b, 30, { principal: owner, title: 'build b' }), + ) + await store.upsert( + sessionEntry(c, 20, { principal: owner, title: 'build c' }), + ) + + const first = await store.list({ + principal: owner, + search: 'build', + limit: 2, + }) + expect(threadIds(first)).toEqual([a, b]) + expect(first.truncated).toBe(true) + + const second = await store.list({ + principal: owner, + search: 'build', + limit: 2, + cursor: required(first.cursor, 'the cursor of a truncated page'), + }) + expect(threadIds(second)).toEqual([c]) + expect(second.truncated).toBeFalsy() + }) + + it('deletes an entry, and ignores an unknown thread', async (ctx) => { + const store = persistence.stores.sessions + if (!store) return ctx.skip('sessions store not provided') + const owner = { id: newId() } + const kept = `${owner.id}-kept` + const removed = `${owner.id}-removed` + await store.upsert(sessionEntry(kept, 1, { principal: owner })) + await store.upsert(sessionEntry(removed, 2, { principal: owner })) + + await store.delete(removed) + await store.delete(newId()) + + expect(await store.get(removed)).toBeUndefined() + expect(threadIds(await store.list({ principal: owner }))).toEqual([ + kept, + ]) + }) + + // Compare-and-set is optional. A store without it skips these cases. + it('writes with upsertIf only at the expected revision', async (ctx) => { + const store = persistence.stores.sessions + if (!store?.getVersioned || !store.upsertIf) + return ctx.skip('sessions store has no compare-and-set') + const threadId = newId() + expect(await store.getVersioned(threadId)).toBeUndefined() + + // `null`: only when absent. + const created = await store.upsertIf(sessionEntry(threadId, 1), null) + if (!created.ok) throw new Error('the first upsertIf conflicted') + expect(await store.upsertIf(sessionEntry(threadId, 9), null)).toEqual({ + ok: false, + reason: 'conflict', + }) + expect(await store.getVersioned(threadId)).toEqual({ + entry: sessionEntry(threadId, 1), + revision: created.revision, + }) + + const changed = await store.upsertIf( + sessionEntry(threadId, 2, { title: 'new' }), + created.revision, + ) + if (!changed.ok) + throw new Error('upsertIf at the read revision conflicted') + expect(changed.revision).not.toBe(created.revision) + + // A stale revision conflicts and leaves the entry as it was. + expect( + await store.upsertIf(sessionEntry(threadId, 3), created.revision), + ).toEqual({ ok: false, reason: 'conflict' }) + expect(await store.get(threadId)).toEqual( + sessionEntry(threadId, 2, { title: 'new' }), + ) + }) + + it('changes the revision on a plain upsert, and never reuses one after delete', async (ctx) => { + const store = persistence.stores.sessions + if (!store?.getVersioned || !store.upsertIf) + return ctx.skip('sessions store has no compare-and-set') + const threadId = newId() + await store.upsert(sessionEntry(threadId, 1)) + const read = await store.getVersioned(threadId) + if (!read) throw new Error('getVersioned found no entry') + + await store.upsert(sessionEntry(threadId, 2)) + + expect( + await store.upsertIf(sessionEntry(threadId, 3), read.revision), + ).toEqual({ ok: false, reason: 'conflict' }) + expect((await store.get(threadId))?.updatedAt).toBe(2) + + // A deleted revision does not match the entry written after it. + const last = required( + await store.getVersioned(threadId), + 'the versioned entry', + ) + await store.delete(threadId) + expect(await store.getVersioned(threadId)).toBeUndefined() + expect( + await store.upsertIf(sessionEntry(threadId, 4), last.revision), + ).toEqual({ ok: false, reason: 'conflict' }) + expect(await store.get(threadId)).toBeUndefined() + const recreated = await store.upsertIf(sessionEntry(threadId, 5), null) + if (!recreated.ok) throw new Error('upsertIf after delete conflicted') + expect(recreated.revision).not.toBe(last.revision) + }) + }) + describe('metadata', () => { it('sets, gets, namespaces, and deletes without composite-key collisions', async (ctx) => { const store = resolveStore('metadata') diff --git a/packages/ai-persistence/src/types.ts b/packages/ai-persistence/src/types.ts index 18eae8b56f..17f28716ce 100644 --- a/packages/ai-persistence/src/types.ts +++ b/packages/ai-persistence/src/types.ts @@ -637,6 +637,412 @@ export interface BlobStore { list: (options?: BlobListOptions) => Promise } +// =========================================================================== +// Session stores +// =========================================================================== +// +// Opt-in stores for a host that keeps long-lived sessions: the inputs it +// accepted, user credentials, an append-only log, turn leases, work claims, +// and a list of sessions. `withPersistence` writes only `sessions`, an entry +// for each subagent child. + +/** Lifecycle of one inbox entry. */ +export type InboxStatus = 'pending' | 'applied' | 'rejected' | 'expired' + +/** + * One input a client sent to a session (a prompt, a steer message, a + * follow-up, an interrupt answer, a command). The host writes it before it + * answers with a receipt, so an accepted input survives a crash. + */ +export interface InboxEntry { + /** + * Idempotency key, unique across all threads (for example a UUID). A second + * append with the same id is a no-op. + */ + inputId: string + threadId: string + /** Who sent it, from the host's authorization. */ + principal?: { id: string; tenantId?: string } + /** The input itself. The store keeps it as-is. The host validates it. */ + input: unknown + status: InboxStatus + createdAt: number + expiresAt?: number + /** The operation that applied the input. */ + operationId?: string + /** Why the input was rejected. */ + reason?: string +} + +/** Durable store for the inputs of a session. */ +export interface InboxStore { + /** + * Store a new entry as `'pending'`, or return the existing entry unchanged + * when `inputId` is already present. + */ + append: (entry: Omit) => Promise + /** Pending entries of a thread, oldest first. */ + listPending: (threadId: string) => Promise> + /** Mark an entry applied by `operationId`. A no-op for an unknown id. */ + markApplied: (inputId: string, operationId: string) => Promise + /** Mark an entry rejected with `reason`. A no-op for an unknown id. */ + markRejected: (inputId: string, reason: string) => Promise + /** The entry for `inputId`, or `null`. */ + get: (inputId: string) => Promise +} + +/** Type an {@link InboxStore} implementation inline. */ +export function defineInboxStore(store: InboxStore): InboxStore { + return store +} + +/** A secret a user or an organization saved: an API key or OAuth tokens. */ +export type Credential = + | { type: 'api_key'; value: string } + | { + type: 'oauth' + accessToken: string + refreshToken?: string + expiresAt?: number + scopes?: Array + /** + * The OAuth client these tokens belong to (for example one made by + * dynamic client registration). A refresh needs it. `issuer` is the + * authorization server that issued the client and the tokens. + */ + client?: { + clientId: string + clientSecret?: string + redirectUri?: string + issuer?: string + } + } + +/** + * Durable store for credentials, keyed by scope and credential id (for + * example `'github'`). The key is `scope.tenantId` and `scope.userId`, not + * `scope.threadId`. A credential saved without `scope.userId` belongs to the + * tenant. A scope with neither is one shared bucket. Encrypt at rest in your + * implementation. + */ +export interface CredentialStore { + get: (scope: Scope, id: string) => Promise + set: (scope: Scope, id: string, credential: Credential) => Promise + delete: (scope: Scope, id: string) => Promise + /** Ids and types only. `list` never returns secret values. */ + list: ( + scope: Scope, + ) => Promise< + Array<{ id: string; type: Credential['type']; expiresAt?: number }> + > +} + +/** Type a {@link CredentialStore} implementation inline. */ +export function defineCredentialStore(store: CredentialStore): CredentialStore { + return store +} + +/** + * One record of a session log. `type` names the kind of record. The other + * fields are JSON. The store does not read records: the writer owns the + * `type` names. + */ +export interface LogRecord { + type: string + [key: string]: unknown +} + +/** A {@link LogRecord} and its position in the log. */ +export interface LogEntry { + /** The position of the record. The first record of a thread is at 1. */ + seq: number + record: LogRecord +} + +/** + * The error {@link LogStore.append} rejects with when `seq` is not the next + * free position of the thread. It means that another writer appended first. + * Nothing of the rejected batch is written. + */ +export class LogConflictError extends Error { + readonly threadId: string + readonly seq: number + + constructor(threadId: string, seq: number) { + super( + `Log conflict on thread ${JSON.stringify(threadId)}: position ${seq} is not the next free position.`, + ) + this.name = 'LogConflictError' + this.threadId = threadId + this.seq = seq + } +} + +/** + * Durable append-only log of a session, one log per thread. A host keeps the + * events, the transcript, the inputs, and the tool steps of a thread in it. + * + * The log has one writer at a time. `append` is a compare-and-append: the + * writer names the position of the first record, and the store refuses the + * batch when another writer took that position first. In SQL, a primary key + * on `(thread_id, seq)` and one transaction per batch give this. + */ +export interface LogStore { + /** + * Write `records` at positions `seq`, `seq + 1`, and so on. + * + * INVARIANT (atomic): the whole batch is written, or none of it is. A crash + * or an error never leaves a part of the batch. + * + * INVARIANT (compare-and-append): `seq` must be the next free position (the + * last position of the thread plus 1, or 1 for a new thread). For any other + * `seq`, a position that is taken or a gap, reject with + * {@link LogConflictError} and write nothing. + * + * An empty batch writes nothing and resolves. + */ + append: ( + threadId: string, + seq: number, + records: ReadonlyArray, + ) => Promise + /** + * The entries of `threadId` with `seq` greater than `after` (default 0), in + * ascending order, at most `limit` of them. `limit: 0` gives `[]`. + * + * INVARIANT: a thread that has no records gives `[]`. The returned records + * are copies: a change to them does not change the log. + */ + read: ( + threadId: string, + options?: { after?: number; limit?: number }, + ) => Promise> + /** + * Call `listener` after each append to `threadId` that this store can see. + * A store that cannot see appends from other processes can poll. Returns a + * function that stops the calls. + */ + subscribe: (threadId: string, listener: () => void) => () => void +} + +/** Type a {@link LogStore} implementation inline. */ +export function defineLogStore(store: LogStore): LogStore { + return store +} + +/** The id of one attempt of a turn. */ +export interface TurnLeaseKey { + threadId: string + inputId: string + /** The operation that runs the attempt. Also the AG-UI run id. */ + operationId: string + /** 1 for the first run of the input, then 1 more after each crash. */ + attempt: number +} + +/** A lease on one attempt of a turn. */ +export interface TurnLease extends TurnLeaseKey { + /** The host that runs the attempt. */ + ownerId: string + /** Epoch milliseconds. */ + expiresAt: number +} + +/** + * Leases for turns, from a system that already has them (a job queue, a + * submission table). A host takes a lease when an attempt starts, renews it + * while the attempt runs, and releases it when the attempt ends. After a + * crash, a host asks `isAlive` to find the attempts that nobody runs any more. + */ +export interface LeaseStore { + acquire: (lease: TurnLease) => Promise + renew: (lease: TurnLease) => Promise + release: (lease: TurnLease) => Promise + /** True while another host still runs the attempt. */ + isAlive: (key: TurnLeaseKey) => Promise +} + +/** + * Claims on threads that have pending work. A host claims a thread while it + * works on it and renews the claim. When the host stops, the claim expires, + * and another host's sweep finds the thread with `listExpired` and takes it + * over with `claim`. + */ +export interface WorkClaimStore { + /** + * Claim `threadId` for `ownerId` until `until` (epoch ms). Renews when the + * owner already holds it. Returns false when another owner holds a claim + * that has not expired. Atomic: two callers never both get true for one + * thread and one moment. + */ + claim: (entry: { + threadId: string + /** The name of the setup that runs the thread, so a sweep can open it. */ + harness: string + ownerId: string + until: number + }) => Promise + /** The thread is idle: remove the claim. Only its owner can. */ + release: (threadId: string, ownerId: string) => Promise + /** Claims that expired before `now`: their host stopped. Oldest first. */ + listExpired: (options: { + now: number + limit?: number + }) => Promise> +} + +/** Type a {@link WorkClaimStore} implementation inline. */ +export function defineWorkClaimStore(store: WorkClaimStore): WorkClaimStore { + return store +} + +/** + * One entry of the session index: what a list of sessions shows for one + * thread. The index holds no messages. The thread data stays in the stores + * that hold it. + * + * @property createdAt - Epoch ms when the thread was first opened. + * @property updatedAt - Epoch ms of the last change. {@link SessionIndexStore.list} + * sorts by it. + */ +export interface SessionIndexEntry> { + threadId: string + /** The name of the setup that runs the thread. */ + harness?: string + title?: string + /** + * The thread that started this one. A subagent child has the thread id + * `subagent:` and the thread of its parent here. + */ + parentThreadId?: string + /** The tool call in the parent thread that started this thread. */ + parentToolCallId?: string + createdAt: number + updatedAt: number + pinned?: boolean + /** Who owns the thread, from the host's authorization. */ + principal?: { id: string; tenantId?: string } + /** Token totals of the thread. */ + usage?: { + /** Model calls. */ + turns: number + promptTokens: number + completionTokens: number + totalTokens: number + /** Input tokens the provider read from its prompt cache. */ + cachedTokens: number + /** Input tokens the provider wrote to its prompt cache. */ + cacheWriteTokens: number + /** Cost in USD, when the host counts it. */ + cost?: number + } + /** Data the app owns, for example the context items a user pinned. */ + metadata?: TMetadata +} + +/** Options for {@link SessionIndexStore.list}. */ +export interface SessionIndexListOptions { + /** The most entries in one page. Without it, the page has all entries. */ + limit?: number + /** The `cursor` of the previous page. */ + cursor?: string + /** + * Only the entries whose `parentThreadId` is this thread. `null`: only the + * entries with no parent (the top-level sessions). + */ + parentThreadId?: string | null + /** + * Only the entries owned by this principal id. With `tenantId`, the + * tenant must match too. An entry without a principal never matches. + */ + principal?: { id: string; tenantId?: string } + /** + * Only the entries whose `title` contains this text, without case. An entry + * without a title never matches. + */ + search?: string + /** Only the entries whose `harness` is exactly this name. */ + harness?: string + /** + * Only the entries whose `metadata` has each key with exactly this string + * value. + */ + metadata?: Record +} + +/** + * One page of a {@link SessionIndexStore.list} scan. + * + * @property cursor - The token for the next page. Only when `truncated`. + * @property truncated - `true` when more entries match after this page. + */ +export interface SessionIndexPage { + entries: Array + cursor?: string + truncated?: boolean +} + +/** Durable index of sessions, one entry per thread. */ +export interface SessionIndexStore { + /** + * Write the entry for `entry.threadId`. + * + * INVARIANT (full replace): a second upsert for the same thread replaces + * the whole entry. A field that the new entry does not have is gone. To + * change one field, `get` the entry and upsert a changed copy. + */ + upsert: (entry: SessionIndexEntry) => Promise + /** The entry for `threadId`, or `undefined`. */ + get: (threadId: string) => Promise + /** + * Optional. The entry with a revision, or `undefined` when it is absent. + * Pair it with `upsertIf` for compare-and-set writes: implement both or + * neither. Each write (`upsert` or `upsertIf`) gives the entry a new + * revision, and a revision does not come back after a `delete`. + */ + getVersioned?: ( + threadId: string, + ) => Promise<{ entry: SessionIndexEntry; revision: string } | undefined> + /** + * Optional. Write `entry` only when the current revision of its thread is + * `expectedRevision` (`null` means "only when absent"). Returns the new + * revision, or `{ ok: false, reason: 'conflict' }` when someone else wrote + * first. With both methods, hosts on several machines do not lose each + * other's changes to one entry. + */ + upsertIf?: ( + entry: SessionIndexEntry, + expectedRevision: string | null, + ) => Promise< + { ok: true; revision: string } | { ok: false; reason: 'conflict' } + > + /** + * The entries that match the filters, newest `updatedAt` first. Entries with + * the same `updatedAt` are in `threadId` order. Pinned entries get no + * special place. All filters must match. + * + * CURSOR SEMANTICS: when `limit` is given and more entries match, the page + * is `truncated: true` with a `cursor`. Pass that `cursor` back to get the + * entries that follow the last entry of the page. An entry that changes + * between two pages can move in the order. + */ + list: (options?: SessionIndexListOptions) => Promise + /** + * Remove the index entry for `threadId`. A no-op for an unknown id. + * + * This removes the entry only. The messages, the log, and the other data of + * the thread stay. The stores that hold that data remove it. + */ + delete: (threadId: string) => Promise +} + +/** Type a {@link SessionIndexStore} implementation inline. */ +export function defineSessionIndexStore( + store: SessionIndexStore, +): SessionIndexStore { + return store +} + /** * Sparse bag of **state** store keys — composition / validation only. * @@ -657,6 +1063,18 @@ export interface AIPersistenceStores { generationRuns?: GenerationRunStore artifacts?: ArtifactStore blobs?: BlobStore + /** Session inputs. */ + inbox?: InboxStore + /** User and tenant credentials. */ + credentials?: CredentialStore + /** Append-only session log. */ + log?: LogStore + /** Turn leases from a system that already has them. */ + leases?: LeaseStore + /** Claims on threads that have pending work. */ + workClaims?: WorkClaimStore + /** The session index. */ + sessions?: SessionIndexStore } /** @@ -835,6 +1253,12 @@ const storeKeys = [ 'metadata', 'artifacts', 'blobs', + 'inbox', + 'credentials', + 'log', + 'leases', + 'workClaims', + 'sessions', ] satisfies Array const storeKeySet = new Set(storeKeys) diff --git a/packages/ai-persistence/tests/inbox.test.ts b/packages/ai-persistence/tests/inbox.test.ts new file mode 100644 index 0000000000..b88414483e --- /dev/null +++ b/packages/ai-persistence/tests/inbox.test.ts @@ -0,0 +1,37 @@ +import { describe, expect, it } from 'vitest' +import { + defineAIPersistence, + defineInboxStore, + memoryPersistence, +} from '../src' + +describe('inbox store', () => { + it('ships in memoryPersistence', () => { + expect(memoryPersistence().stores.inbox).toBeDefined() + }) + + it('returns a copy, so callers cannot change stored entries', async () => { + const inbox = memoryPersistence().stores.inbox + const entry = await inbox.append({ + inputId: 'in-1', + threadId: 't', + input: { op: 'prompt' }, + createdAt: 1, + }) + entry.status = 'applied' + + expect((await inbox.get('in-1'))?.status).toBe('pending') + }) + + it('accepts a user-implemented inbox in defineAIPersistence', () => { + const inbox = defineInboxStore({ + append: (entry) => Promise.resolve({ ...entry, status: 'pending' }), + listPending: () => Promise.resolve([]), + markApplied: () => Promise.resolve(), + markRejected: () => Promise.resolve(), + get: () => Promise.resolve(null), + }) + + expect(defineAIPersistence({ stores: { inbox } }).stores.inbox).toBe(inbox) + }) +}) diff --git a/packages/ai-persistence/tests/memory.conformance.test.ts b/packages/ai-persistence/tests/memory.conformance.test.ts index a3a693f674..ff6aec26bd 100644 --- a/packages/ai-persistence/tests/memory.conformance.test.ts +++ b/packages/ai-persistence/tests/memory.conformance.test.ts @@ -1,6 +1,13 @@ import { runPersistenceConformance } from '../src/testkit/conformance' -import { memoryPersistence } from '../src/memory' +import { memoryLogStore, memoryPersistence } from '../src/memory' +import { defineAIPersistence } from '../src/types' -runPersistenceConformance('memory', () => memoryPersistence(), { - checks: ['messages.metadata', 'runs.listByThread.state'], -}) +// memoryPersistence() has no log, so add one to run the log cases too. +runPersistenceConformance( + 'memory', + () => + defineAIPersistence({ + stores: { ...memoryPersistence().stores, log: memoryLogStore() }, + }), + { checks: ['messages.metadata', 'runs.listByThread.state'] }, +) diff --git a/packages/ai-persistence/tests/memory.test.ts b/packages/ai-persistence/tests/memory.test.ts index e89ff1cd72..95f0628be1 100644 --- a/packages/ai-persistence/tests/memory.test.ts +++ b/packages/ai-persistence/tests/memory.test.ts @@ -17,11 +17,15 @@ describe('memoryPersistence', () => { 'activities', 'artifacts', 'blobs', + 'credentials', 'generationRuns', + 'inbox', 'interrupts', 'messages', 'metadata', 'runs', + 'sessions', + 'workClaims', ]) }) diff --git a/packages/ai-persistence/tests/persistence-types.test-d.ts b/packages/ai-persistence/tests/persistence-types.test-d.ts index 39e0520190..06a89922db 100644 --- a/packages/ai-persistence/tests/persistence-types.test-d.ts +++ b/packages/ai-persistence/tests/persistence-types.test-d.ts @@ -22,12 +22,16 @@ import type { ChatPersistence, ChatTranscriptPersistence, ChatTranscriptStores, + CredentialStore, GenerationRunStore, + InboxStore, InterruptStore, MessagePage, MessageStore, MetadataStore, RunStore, + SessionIndexStore, + WorkClaimStore, } from '../src' declare const messages: MessageStore @@ -143,6 +147,10 @@ expectTypeOf(memoryPersistence()).toEqualTypeOf< metadata: MetadataStore artifacts: ArtifactStore blobs: BlobStore + inbox: InboxStore + credentials: CredentialStore + workClaims: WorkClaimStore + sessions: SessionIndexStore }> >() const transcript: ChatTranscriptPersistence = messagesOnly diff --git a/packages/ai-persistence/tests/persistence-validation.test.ts b/packages/ai-persistence/tests/persistence-validation.test.ts index 7c669775dd..6359c89c1d 100644 --- a/packages/ai-persistence/tests/persistence-validation.test.ts +++ b/packages/ai-persistence/tests/persistence-validation.test.ts @@ -1,8 +1,9 @@ import { describe, expect, it } from 'vitest' +import { memoryLogStore, memoryPersistence } from '../src/memory' import { withGenerationPersistence, withPersistence } from '../src/middleware' import { reconstructChat } from '../src/reconstruct' import { defineAIPersistence } from '../src/types' -import type { ChatTranscriptPersistence } from '../src/types' +import type { ChatTranscriptPersistence, LeaseStore } from '../src/types' import { createGenerationRunStore, createInterruptStore, @@ -96,4 +97,22 @@ describe('persistence store dependency validation', () => { ), ).rejects.toThrow(/requires stores\.messages/i) }) + + it('accepts every session store key', () => { + const leases: LeaseStore = { + acquire: async () => {}, + renew: async () => {}, + release: async () => {}, + isAlive: async () => false, + } + const { inbox, credentials, workClaims, sessions } = + memoryPersistence().stores + const log = memoryLogStore() + + expect(() => + defineAIPersistence({ + stores: { inbox, credentials, log, leases, workClaims, sessions }, + }), + ).not.toThrow() + }) }) diff --git a/packages/ai-persistence/tests/subagent-persistence-cards.test.ts b/packages/ai-persistence/tests/subagent-persistence-cards.test.ts index 0671294ef5..750edced4f 100644 --- a/packages/ai-persistence/tests/subagent-persistence-cards.test.ts +++ b/packages/ai-persistence/tests/subagent-persistence-cards.test.ts @@ -651,6 +651,7 @@ describe('subagent run recorder', () => { const { stores } = memoryPersistence() const recorder = createSubagentRunRecorder({ messages: stores.messages, + sessions: stores.sessions, intervalMs: 0, }) await feed(recorder, 'run-1', [ @@ -662,6 +663,10 @@ describe('subagent run recorder', () => { ...textMessage('m1', 'Inner notes', { subagentRunId: 'inner' }), ]) + expect(await stores.sessions.get('subagent:inner')).toMatchObject({ + parentThreadId: 'subagent:outer', + parentToolCallId: 'call_n', + }) // The outer child's chat can continue it. The top thread cannot. const fromOuter = { threadId: 'desk:writer', subagentRunId: 'outer' } expect((await recorder.loadChild('inner', fromOuter))?.messages).toEqual([ @@ -672,6 +677,35 @@ describe('subagent run recorder', () => { ).toBeUndefined() }) + it('gives each child the owner of its parent thread', async () => { + const { stores } = memoryPersistence() + await stores.sessions.upsert({ + threadId: 'desk', + principal: { id: 'alice' }, + createdAt: t, + updatedAt: t, + }) + const recorder = createSubagentRunRecorder({ + messages: stores.messages, + sessions: stores.sessions, + intervalMs: 0, + }) + await feed(recorder, 'run-1', [ + childStarted('outer', { parentToolCallId: 'call_1' }), + childStarted('inner', { + parentSubagentRunId: 'outer', + parentToolCallId: 'call_n', + }), + ]) + + const owners = await Promise.all( + ['subagent:outer', 'subagent:inner'].map( + async (threadId) => (await stores.sessions.get(threadId))?.principal, + ), + ) + expect(owners).toEqual([{ id: 'alice' }, { id: 'alice' }]) + }) + it('loads a waiting continued child without the messages of the call it waits in', async () => { const { stores } = memoryPersistence() const recorder = createSubagentRunRecorder({ @@ -800,6 +834,66 @@ describe('subagent run recorder', () => { }) describe('agent tools with persistence', () => { + it('writes one index entry per child and keeps fields others wrote', async () => { + const persistence = memoryPersistence() + const middleware = withPersistence(persistence) + const sessions = persistence.stores.sessions + const { agent } = noteTaker() + const first = await subagentTurn(middleware, agent, { + runId: 'run-1', + toolCallId: 'call_1', + args: { prompt: 'Take notes' }, + }) + const threadId = `subagent:${startedId(first)}` + const entry = await sessions.get(threadId) + expect(entry).toEqual({ + threadId, + parentThreadId: 'desk', + parentToolCallId: 'call_1', + createdAt: expect.any(Number), + updatedAt: expect.any(Number), + }) + if (!entry) throw new Error('no index entry') + + // Another writer names the session. A continue keeps the name. + await sessions.upsert({ ...entry, title: 'Tide notes' }) + await subagentTurn(middleware, agent, { + runId: 'run-2', + toolCallId: 'call_2', + args: { sessionId: startedId(first), prompt: 'More' }, + }) + expect(await sessions.get(threadId)).toEqual({ + threadId, + title: 'Tide notes', + parentThreadId: 'desk', + parentToolCallId: 'call_2', + createdAt: entry.createdAt, + updatedAt: expect.any(Number), + }) + }) + + it('gives a child entry the harness and owner of its parent entry', async () => { + const persistence = memoryPersistence() + const sessions = persistence.stores.sessions + await sessions.upsert({ + threadId: 'desk', + harness: 'notes/desk', + principal: { id: 'alice' }, + createdAt: 1, + updatedAt: 1, + }) + const { agent } = noteTaker() + const chunks = await subagentTurn(withPersistence(persistence), agent, { + runId: 'run-1', + toolCallId: 'call_1', + args: { prompt: 'Take notes' }, + }) + + const entry = await sessions.get(`subagent:${startedId(chunks)}`) + expect(entry?.harness).toBe('notes/desk') + expect(entry?.principal).toEqual({ id: 'alice' }) + }) + it('continues a stored child by sessionId', async () => { const middleware = withPersistence(memoryPersistence()) const { agent, seen } = noteTaker() From a0d68a375bbeb351d5a7f0eae298d541aa6a8da7 Mon Sep 17 00:00:00 2001 From: Alem Tuzlak Date: Tue, 13 Oct 2026 22:12:10 +0200 Subject: [PATCH 002/131] test(ai-persistence): name each inbox id in the conformance case --- packages/ai-persistence/src/testkit/conformance.ts | 10 ++++------ 1 file changed, 4 insertions(+), 6 deletions(-) diff --git a/packages/ai-persistence/src/testkit/conformance.ts b/packages/ai-persistence/src/testkit/conformance.ts index fae36d9de0..b76a86f21e 100644 --- a/packages/ai-persistence/src/testkit/conformance.ts +++ b/packages/ai-persistence/src/testkit/conformance.ts @@ -1627,12 +1627,10 @@ export function runPersistenceConformance( // Own ids, because the suite shares one persistence. const id = (name: string) => `${name}-${crypto.randomUUID()}` const thread = id('inbox-thread') - const [in1, in2, other, missing] = [ - 'in-1', - 'in-2', - 'in-other', - 'missing', - ].map(id) + const in1 = id('in-1') + const in2 = id('in-2') + const other = id('in-other') + const missing = id('missing') const first = await store.append({ inputId: in1, From ebc8bc80b819c3a8d1d766334b60ffd3bffeb46b Mon Sep 17 00:00:00 2001 From: Alem Tuzlak Date: Tue, 13 Oct 2026 22:22:52 +0200 Subject: [PATCH 003/131] fix(ai-persistence): read the sessions store through the persistence plan --- packages/ai-persistence/src/middleware.ts | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/packages/ai-persistence/src/middleware.ts b/packages/ai-persistence/src/middleware.ts index 1c8d654d1b..7291857964 100644 --- a/packages/ai-persistence/src/middleware.ts +++ b/packages/ai-persistence/src/middleware.ts @@ -1727,6 +1727,7 @@ interface PersistencePlan { wantsInterrupts: boolean wantsArtifactPersistence: boolean runs: AIPersistence['stores']['runs'] + sessions: AIPersistence['stores']['sessions'] } function resolvePersistencePlan(persistence: AIPersistence): PersistencePlan { @@ -1736,6 +1737,7 @@ function resolvePersistencePlan(persistence: AIPersistence): PersistencePlan { persistence.stores.artifacts !== undefined && persistence.stores.blobs !== undefined, runs: persistence.stores.runs, + sessions: persistence.stores.sessions, } } @@ -2061,7 +2063,7 @@ export function withPersistence( const snapshotStreaming = options.snapshotStreaming ?? false const snapshotIntervalMs = options.snapshotIntervalMs ?? 1000 const plan = resolvePersistencePlan(persistence) - const { wantsInterrupts, runs } = plan + const { wantsInterrupts, runs, sessions } = plan const messageStore = persistence.stores.messages if (!messageStore) { // validateChatPersistenceStores already throws; this narrows for TypeScript. @@ -2100,9 +2102,7 @@ export function withPersistence( ...(wantsInterrupts && persistence.stores.interrupts ? { interrupts: persistence.stores.interrupts } : {}), - ...(persistence.stores.sessions - ? { sessions: persistence.stores.sessions } - : {}), + ...(sessions ? { sessions } : {}), }) function resetStreaming(state: RunStateEntry) { From fd7612916f431c50e9e0fb6c47df93def01f4584 Mon Sep 17 00:00:00 2001 From: Alem Tuzlak Date: Tue, 13 Oct 2026 22:40:28 +0200 Subject: [PATCH 004/131] feat(ai-harness): new package with the harness contracts and plugin API --- packages/ai-harness/package.json | 77 ++ packages/ai-harness/src/agents.ts | 184 +++++ packages/ai-harness/src/auth.ts | 196 +++++ packages/ai-harness/src/commands.ts | 163 ++++ packages/ai-harness/src/config.ts | 92 +++ packages/ai-harness/src/define.ts | 380 +++++++++ packages/ai-harness/src/extensions.ts | 40 + packages/ai-harness/src/feed.ts | 117 +++ packages/ai-harness/src/index.ts | 170 ++++ packages/ai-harness/src/media-ref.ts | 166 ++++ packages/ai-harness/src/operation.ts | 124 +++ packages/ai-harness/src/plugins.ts | 967 +++++++++++++++++++++++ packages/ai-harness/src/resources.ts | 84 ++ packages/ai-harness/src/session-entry.ts | 109 +++ packages/ai-harness/src/types.ts | 421 ++++++++++ packages/ai-harness/src/usage.ts | 96 +++ packages/ai-harness/src/utils.ts | 3 + packages/ai-harness/tsconfig.json | 8 + packages/ai-harness/vite.config.ts | 35 + pnpm-lock.yaml | 40 +- 20 files changed, 3460 insertions(+), 12 deletions(-) create mode 100644 packages/ai-harness/package.json create mode 100644 packages/ai-harness/src/agents.ts create mode 100644 packages/ai-harness/src/auth.ts create mode 100644 packages/ai-harness/src/commands.ts create mode 100644 packages/ai-harness/src/config.ts create mode 100644 packages/ai-harness/src/define.ts create mode 100644 packages/ai-harness/src/extensions.ts create mode 100644 packages/ai-harness/src/feed.ts create mode 100644 packages/ai-harness/src/index.ts create mode 100644 packages/ai-harness/src/media-ref.ts create mode 100644 packages/ai-harness/src/operation.ts create mode 100644 packages/ai-harness/src/plugins.ts create mode 100644 packages/ai-harness/src/resources.ts create mode 100644 packages/ai-harness/src/session-entry.ts create mode 100644 packages/ai-harness/src/types.ts create mode 100644 packages/ai-harness/src/usage.ts create mode 100644 packages/ai-harness/src/utils.ts create mode 100644 packages/ai-harness/tsconfig.json create mode 100644 packages/ai-harness/vite.config.ts diff --git a/packages/ai-harness/package.json b/packages/ai-harness/package.json new file mode 100644 index 0000000000..5d45fa9b9d --- /dev/null +++ b/packages/ai-harness/package.json @@ -0,0 +1,77 @@ +{ + "name": "@tanstack/ai-harness", + "version": "0.0.0", + "description": "Build agent harnesses on TanStack AI: long-lived sessions, typed agents, and plugins", + "keywords": [ + "agent", + "ai", + "harness", + "plugins", + "session", + "tanstack" + ], + "license": "MIT", + "author": "", + "repository": { + "type": "git", + "url": "git+https://github.com/TanStack/ai.git", + "directory": "packages/ai-harness" + }, + "files": [ + "dist", + "src" + ], + "type": "module", + "sideEffects": false, + "module": "./dist/esm/index.js", + "types": "./dist/esm/index.d.ts", + "exports": { + ".": { + "types": "./dist/esm/index.d.ts", + "import": "./dist/esm/index.js" + } + }, + "scripts": { + "build": "vite build", + "clean": "premove ./build ./dist", + "lint:fix": "oxlint src --type-aware --fix", + "test:build": "publint --strict", + "test:oxlint": "oxlint src --type-aware", + "test:coverage": "vitest run --passWithNoTests --coverage --coverage.include='src/**' --coverage.reporter=text-summary --coverage.reporter=json-summary", + "test:lib": "vitest --passWithNoTests", + "test:lib:dev": "pnpm test:lib --watch", + "test:types": "tsc" + }, + "devDependencies": { + "@tanstack/ai": "workspace:*", + "@tanstack/ai-persistence": "workspace:*", + "@vitest/coverage-v8": "4.1.10", + "zod": "^4.2.0" + }, + "peerDependencies": { + "@tanstack/ai": "workspace:^", + "@tanstack/ai-persistence": "workspace:^" + }, + "nx": { + "targets": { + "test:lib": { + "dependsOn": [ + "build", + "^build" + ] + }, + "test:coverage": { + "dependsOn": [ + "build", + "^build" + ] + }, + "test:types": { + "dependsOn": [ + "build", + "^build" + ] + } + } + } +} diff --git a/packages/ai-harness/src/agents.ts b/packages/ai-harness/src/agents.ts new file mode 100644 index 0000000000..fa350e2772 --- /dev/null +++ b/packages/ai-harness/src/agents.ts @@ -0,0 +1,184 @@ +import type { + AgentProduces, + DefinedAgent, + InferSchemaType, + SchemaInput, +} from '@tanstack/ai' +import type { ToolStep } from './durable-tool' +import type { AgentStartOptions } from './session' +import type { AgentRun } from './types' + +/** Any agent made with `defineAgent`, with its types kept. */ +export type AnyAgent = DefinedAgent + +/** The input an agent takes: its `inputSchema` type, or `undefined`. */ +export type AgentInputOf = TAgent extends { + inputSchema?: infer TSchema +} + ? TSchema extends SchemaInput + ? InferSchemaType + : undefined + : undefined + +/** + * What running an agent resolves to: + * - its `outputSchema` type, when it has one; + * - else the value its promise `run` resolves to; + * - else the child's text. + */ +export type AgentResultOf = + TAgent extends DefinedAgent< + any, + any, + infer TSchema, + any, + any, + infer TResult, + any + > + ? TSchema extends SchemaInput + ? InferSchemaType + : unknown extends TResult + ? string + : TResult + : unknown + +/** A read-only view of the agents a session can run. */ +export interface AgentRegistryView { + /** Every agent, in registration order. */ + list: () => ReadonlyArray + /** The agent named `name`, or `undefined`. */ + get: (name: string) => AnyAgent | undefined + /** The first agent whose `produces` matches, or `undefined`. */ + find: (query: { produces: AgentProduces }) => AnyAgent | undefined +} + +/** + * The agents of one session, with the owner of each (the harness, or a + * plugin name). Adding a name twice with two different agents is an error that + * names both owners. Adding the same agent object twice is allowed, so an + * agent can sit in both `agents` and `subagents.agents`. + */ +export class AgentRegistry implements AgentRegistryView { + private readonly agents = new Map< + string, + { agent: AnyAgent; owner: string } + >() + + add(agent: AnyAgent, owner: string): void { + const existing = this.agents.get(agent.name) + if (existing && existing.agent !== agent) { + throw new Error( + `Duplicate agent "${agent.name}": first owner ${existing.owner}, second owner ${owner}.`, + ) + } + if (!existing) this.agents.set(agent.name, { agent, owner }) + } + + /** Add `owner`'s agent, or replace it. A name another owner has throws. */ + set(agent: AnyAgent, owner: string): void { + const existing = this.agents.get(agent.name) + if (existing && existing.owner !== owner) { + throw new Error( + `Agent "${agent.name}" belongs to ${existing.owner}. ${owner} cannot replace it.`, + ) + } + this.agents.set(agent.name, { agent, owner }) + } + + /** Remove `owner`'s agent `name`. Another owner's name does nothing. */ + delete(name: string, owner: string): void { + if (this.agents.get(name)?.owner === owner) this.agents.delete(name) + } + + /** Remove every agent that the harness itself did not add. */ + deletePluginAgents(): void { + for (const [name, entry] of this.agents) { + if (entry.owner !== 'the harness') this.agents.delete(name) + } + } + + list(): ReadonlyArray { + return [...this.agents.values()].map((entry) => entry.agent) + } + + get(name: string): AnyAgent | undefined { + return this.agents.get(name)?.agent + } + + find(query: { produces: AgentProduces }): AnyAgent | undefined { + return this.list().find((agent) => agent.produces === query.produces) + } + + /** A copy for one chat turn, so turn plugins can add agents for that turn only. */ + fork(): AgentRegistry { + const copy = new AgentRegistry() + for (const [name, entry] of this.agents) copy.agents.set(name, entry) + return copy + } +} + +/** `harnessContext(ctx).agents`: start agents in the background. */ +export interface AgentStarter { + /** + * Start `agent` (a definition, or a name the session knows) in the + * background, and return its run at once. The child counts against the + * subagent tree budget of the run that starts it. Over the budget, `start` + * throws and nothing starts. + */ + start: ( + agent: AnyAgent | string, + input?: unknown, + options?: AgentStartOptions, + ) => AgentRun +} + +/** What a harness session adds to an agent run that it starts from code. */ +export interface HarnessAgentContext { + /** + * The steps of the run. On a durable host, in a run started with + * `resume: true`, a finished step returns its stored value when the run + * runs again, and `fn` does not run again. Else `fn` runs each time. + */ + step: ToolStep + /** Start agents in the background, as children of this run. */ + agents: AgentStarter +} + +const NO_SESSION: HarnessAgentContext = { + step: { do: async (_name, fn) => fn() }, + agents: { + start: () => { + throw new Error( + 'harnessContext(ctx).agents needs a run that a harness session started.', + ) + }, + }, +} + +const contexts = new WeakMap() + +/** + * Read `step` and `agents` in the `run` of an agent: `harnessContext(ctx)`. + * A harness session gives them to each agent run it starts from code (a + * session agent, `harnessContext(ctx).agents.start`, a background child). + * In a tool-mode child or outside a harness, `step.do` runs `fn` each time + * and `agents.start` throws. + */ +export function harnessContext(ctx: object): HarnessAgentContext { + return contexts.get(ctx) ?? NO_SESSION +} + +/** `agent`, with `extras` as its {@link harnessContext} in every run. */ +export function withHarnessContext( + agent: AnyAgent, + extras: Partial, +): AnyAgent { + return { + ...agent, + run: (ctx) => { + contexts.set(ctx, { ...NO_SESSION, ...extras }) + return agent.run(ctx) + }, + } +} diff --git a/packages/ai-harness/src/auth.ts b/packages/ai-harness/src/auth.ts new file mode 100644 index 0000000000..31f6b2d384 --- /dev/null +++ b/packages/ai-harness/src/auth.ts @@ -0,0 +1,196 @@ +import { isKeyedAdapter } from '@tanstack/ai' +import type { ProviderKeys, Scope } from '@tanstack/ai' +import type { ByokProvider, ProviderId } from '@tanstack/ai/byok' +import type { Credential, CredentialStore } from '@tanstack/ai-persistence' + +/** + * Thrown when a tool or command needs a credential the user has not saved. + * The session publishes `harness.auth_required`, so a host can show the + * sign-in link or run `/connect`. + */ +export class AuthRequiredError extends Error { + readonly connector: string + readonly url: string | undefined + constructor(connector: string, url?: string) { + super( + url + ? `Sign in to ${connector} first: ${url}` + : `Sign in to ${connector} first. Run /connect ${connector}.`, + ) + this.name = 'AuthRequiredError' + this.connector = connector + this.url = url + } +} + +/** + * Credentials of the running turn's sender (outside a turn, of the session's + * principal), as plugins see them. + */ +export interface CredentialsAccess { + get: (id: string) => Promise + /** + * Like `get`, but throws {@link AuthRequiredError} when the credential is + * missing. With `wait: true`, a tool of a chat turn waits for the sign-in + * instead: the turn stops with an interrupt with the reason + * `'auth_required'`. When the credential is saved (for example by + * `connect:`), the turn goes on and the tool runs again. Cancel the + * interrupt to fail the tool. `wait` is for a tool of the running chat + * turn. Do not pass it elsewhere (a command, a background agent, a hook): + * the credentials of a command never wait, and outside a chat turn `wait` + * does nothing. + */ + require: (id: string, options?: { wait?: boolean }) => Promise + set: (id: string, credential: Credential) => Promise + delete: (id: string) => Promise + list: () => Promise< + Array<{ id: string; type: Credential['type']; expiresAt?: number }> + > +} + +/** + * What a session adds to `credentialsFor`: `canWait` says if a `require` + * with `wait` can stop the turn now, and `onSet` runs after each save. + */ +export interface CredentialHooks { + canWait?: (id: string) => boolean + onSet?: (id: string) => Promise +} + +/** + * The error that chat() reads as "stop the turn for input" (the shape of + * `MCPInputRequiredError`), with the reason `'auth_required'`. + */ +function signInRequired(error: AuthRequiredError) { + return Object.assign(new Error(error.message), { + name: 'MCPInputRequiredError', + kind: 'form' as const, + reason: 'auth_required', + request: { + connector: error.connector, + ...(error.url ? { url: error.url } : {}), + }, + }) +} + +/** + * The tenant scope of a user scope: the same tenant, no user. A credential + * saved there belongs to every user of the tenant. + */ +const tenantOf = ({ userId, ...tenant }: Scope): Scope | undefined => + userId === undefined ? undefined : tenant + +/** + * Scope credentials to one principal and thread. A function `scope` is read + * at each call, so the scope can follow the sender of the running turn. + * + * A read finds the user's own credential first, then the tenant's (saved + * without a `userId`). A save or a delete changes the user's own only. + */ +export function credentialsFor( + store: CredentialStore, + scope: Scope | (() => Scope), + onMissing: (error: AuthRequiredError) => void, + hooks?: CredentialHooks, +): CredentialsAccess { + const scopeNow = typeof scope === 'function' ? scope : () => scope + // ponytail: a refreshed tenant OAuth token is saved for the user, so each + // user keeps a copy. A provider that rotates refresh tokens would need a + // save for the whole tenant to share one token. + const get = async (id: string) => { + const user = scopeNow() + const own = await store.get(user, id) + if (own) return own + const tenant = tenantOf(user) + return tenant ? store.get(tenant, id) : null + } + return { + get, + require: async (id, options) => { + const credential = await get(id) + if (credential) return credential + const error = new AuthRequiredError(id) + onMissing(error) + if (options?.wait && hooks?.canWait?.(id)) throw signInRequired(error) + throw error + }, + set: async (id, credential) => { + await store.set(scopeNow(), id, credential) + await hooks?.onSet?.(id) + }, + delete: (id) => store.delete(scopeNow(), id), + list: async () => { + const user = scopeNow() + const tenant = tenantOf(user) + const own = await store.list(user) + if (!tenant) return own + const shared = await store.list(tenant) + return [ + ...own, + ...shared.filter((item) => !own.some((o) => o.id === item.id)), + ] + }, + } +} + +/** + * The first name in `provider.env` that is set, with its value. A provider + * id has no env names. Where `process` is missing (a browser, a worker) it + * returns `null`. + */ +export function envKeyOf(provider: ByokProvider | ProviderId) { + if (typeof provider === 'string') return null + const env = globalThis.process?.env + for (const name of provider.env ?? []) { + const value = env?.[name] + if (typeof value === 'string' && value.length > 0) return { name, value } + } + return null +} + +/** + * The provider keys of the session's principal. A key is the `api_key` + * credential saved under the provider id (`/connect ` saves it), else + * the first `provider.env` name that is set. A missing key calls `onMissing` + * and throws {@link AuthRequiredError}, so the user sees `/connect `. + */ +export function providerKeysFor( + credentials: CredentialsAccess, + onMissing: (error: AuthRequiredError) => void, +) { + const providerIdOf = (provider: ByokProvider | ProviderId) => + typeof provider === 'string' ? provider : provider.id + const get = async (provider: ByokProvider | ProviderId) => { + const saved = await credentials.get(providerIdOf(provider)) + if (saved?.type === 'api_key') return saved.value + return envKeyOf(provider)?.value ?? null + } + const requireKey = async (provider: ByokProvider | ProviderId) => { + const key = await get(provider) + if (key !== null) return key + const error = new AuthRequiredError(providerIdOf(provider)) + onMissing(error) + throw error + } + const keys: ProviderKeys = { + get, + require: requireKey, + adapter: async (adapter) => + isKeyedAdapter(adapter) + ? adapter.create(await requireKey(adapter.provider)) + : adapter, + } + return keys +} + +/** Remove secret-looking values from a message before it reaches a model or a log. */ +export function scrubSecrets( + text: string, + secrets: ReadonlyArray, +): string { + let scrubbed = text + for (const secret of secrets) { + if (secret.length >= 6) scrubbed = scrubbed.split(secret).join('[redacted]') + } + return scrubbed +} diff --git a/packages/ai-harness/src/commands.ts b/packages/ai-harness/src/commands.ts new file mode 100644 index 0000000000..a9fdac9cba --- /dev/null +++ b/packages/ai-harness/src/commands.ts @@ -0,0 +1,163 @@ +import type { InferSchemaType, ModelMessage, SchemaInput } from '@tanstack/ai' +import type { SessionIndexEntry } from '@tanstack/ai-persistence' +import type { CredentialsAccess } from './auth' +import type { SessionSnapshot } from './session' +import type { SessionEntryChange, SessionEntryUpdate } from './session-entry' +import type { + ChatTurnResult, + Operation, + Principal, + Receipt, + ThreadSettings, +} from './types' + +/** A question a command asks the user. The host renders it from the schema. */ +export interface Question { + message: string + /** The answer's shape. Without one, the answer is any text. */ + schema?: TSchema + /** + * The answer is a secret, for example an API key or a password. Hosts hide + * it while the user types, and the session does not keep it in the inbox. + */ + secret?: boolean + /** + * A page the user must open before they answer, for example a sign-in or a + * verify page. Hosts show the link. The session does not open it. + */ + url?: string +} + +export type AnswerOf = TSchema extends SchemaInput + ? InferSchemaType + : string + +/** + * The parts of the session a plugin can use. `TMetadata` is the type of the + * session index `metadata`, from the `metadataSchema` of the plugin. + */ +export interface PluginSessionApi> { + threadId: string + /** + * Who sent the input of the running chat turn. Outside a turn: the + * principal that opened the session. In a command's `run` context: who + * runs the command. Code that runs after its turn ends (a background + * agent's tool) gets the sender of the turn that runs at that time. + */ + principal: Principal | undefined + snapshot: () => SessionSnapshot + /** + * Start a chat turn, as if the user typed `text`. While a turn runs, the new + * turn waits in the queue. Cancel the returned turn to drop it before it starts. + */ + prompt: (text: string) => Operation + /** + * Add `text` to the transcript as an assistant note, not as a user prompt. + * When no turn runs, the note is written at once. Else it is written before + * the next turn starts. With `wake: true`, a new turn that sends `text` is + * also queued, so the model can act on the note. Resolves when the note is + * written or queued and the wake turn is queued, not when that turn ends. + */ + note: (text: string, options?: { wake?: boolean }) => Promise + /** The saved transcript. */ + transcript: () => Promise> + /** Replace the saved transcript (for example after a summary). */ + replaceTranscript: (messages: Array) => Promise + /** + * The session index entry of this thread. `undefined` when the host has + * no `stores.sessions`, or the thread has no entry. + */ + entry: () => Promise | undefined> + /** + * Change the session index entry of this thread. Each field you pass + * replaces the stored one. Pass a function of the current entry to change + * part of a field: it runs again when another host wrote first. The + * `metadataSchema` of the plugin and of the harness check `metadata`, and + * a refused value throws `SessionMetadataError`. Resolves to the changed + * entry. Without `stores.sessions`, or for a thread with no entry, it + * writes nothing and resolves to `undefined`. + * + * @example + * ```ts + * await ctx.session.updateEntry({ title: 'Fix the login bug' }) + * ``` + */ + updateEntry: ( + update: SessionEntryUpdate< + TMetadata, + Omit, 'pinned'> & + Pick + >, + ) => Promise | undefined> + /** + * Ask the user and wait for the answer. Hosts show it (the CLI prompts, + * ACP shows an elicitation). Rejects when the session closes. + */ + ask: ( + question: Question, + ) => Promise> + /** + * Tell clients that the user must sign in. Hosts show the link (the CLI + * opens the browser) or the device code. + */ + authRequired: (info: { + connector: string + url?: string + userCode?: string + }) => void + /** Change a session setting, the same as a client. Applies at the next turn. */ + setConfig: (key: string, value: unknown) => Promise + /** + * The stored settings of the thread (`session.configure`). Read them at + * each call, for example the working folder `cwd`: a change applies from + * the next turn. + */ + settings: () => ThreadSettings +} + +/** What a command handler receives besides its input. */ +export interface CommandContext { + signal: AbortSignal + /** The session, for the user who runs the command: `session.prompt` runs as that user. */ + session: PluginSessionApi + /** Who runs the command. Default: the principal that opened the session. */ + principal?: Principal + /** + * The credentials of the user who runs the command. Save a sign-in here, + * so it belongs to that user. + */ + credentials: CredentialsAccess +} + +/** A user action: a slash command, a button, a dashboard action. Not a model tool. */ +export interface CommandDefinition< + TSchema extends SchemaInput | undefined = SchemaInput | undefined, +> { + description: string + /** The input shape, checked before `run`. */ + input?: TSchema + run: ( + input: TSchema extends SchemaInput ? InferSchemaType : undefined, + ctx: CommandContext, + ) => unknown +} + +/** A command with any input type. */ +export type AnyCommand = CommandDefinition + +/** + * Define a command with a typed input. + * + * @example + * ```ts + * const clear = defineCommand({ + * description: 'Remove every todo', + * run: async () => ({ removed: await db.clear() }), + * }) + * ``` + */ +export function defineCommand< + const TSchema extends SchemaInput | undefined = undefined, +>(command: CommandDefinition): CommandDefinition { + return command +} diff --git a/packages/ai-harness/src/config.ts b/packages/ai-harness/src/config.ts new file mode 100644 index 0000000000..72df9a16d7 --- /dev/null +++ b/packages/ai-harness/src/config.ts @@ -0,0 +1,92 @@ +/** + * A session setting a plugin declares, for example the model or the + * thinking level. Hosts render it from this description, for example a + * picker in the CLI. `serveAcp` does not send config options yet. A change + * applies at the next turn. + */ +export type ConfigOption = + | { + type: 'select' + options: ReadonlyArray + default: string + description?: string + /** + * A kind for hosts that group settings, for example `'model'` or + * `'thought_level'`. `serveAcp` does not send it yet. + */ + category?: string + } + | { type: 'boolean'; default: boolean; description?: string } + | { type: 'text'; default: string; description?: string } + | { + type: 'number' + default: number + min?: number + max?: number + description?: string + } + +type Without = Omit +type Option = Extract< + ConfigOption, + { type: TType } +> + +/** Helpers to declare {@link ConfigOption}s. */ +export const configOption = { + select: (option: Without, 'type'>): Option<'select'> => { + if (!option.options.includes(option.default)) { + throw new Error( + `configOption.select: the default "${option.default}" is not an option.`, + ) + } + return { type: 'select', ...option } + }, + boolean: (option: Without, 'type'>): Option<'boolean'> => ({ + type: 'boolean', + ...option, + }), + text: (option: Without, 'type'>): Option<'text'> => ({ + type: 'text', + ...option, + }), + number: (option: Without, 'type'>): Option<'number'> => ({ + type: 'number', + ...option, + }), +} + +/** Check a value for an option. Returns the value, or throws with a reason. */ +export function checkConfigValue( + key: string, + option: ConfigOption, + value: unknown, +): unknown { + const fail = (expected: string): never => { + throw new Error( + `Config "${key}" expects ${expected}, got ${JSON.stringify(value)}.`, + ) + } + switch (option.type) { + case 'select': + if (typeof value !== 'string' || !option.options.includes(value)) { + fail(`one of ${option.options.join(', ')}`) + } + return value + case 'boolean': + if (typeof value !== 'boolean') fail('true or false') + return value + case 'text': + if (typeof value !== 'string') fail('text') + return value + case 'number': { + if (typeof value !== 'number' || !Number.isFinite(value)) + return fail('a number') + if (option.min !== undefined && value < option.min) + fail(`a number >= ${option.min}`) + if (option.max !== undefined && value > option.max) + fail(`a number <= ${option.max}`) + return value + } + } +} diff --git a/packages/ai-harness/src/define.ts b/packages/ai-harness/src/define.ts new file mode 100644 index 0000000000..d54f1f3375 --- /dev/null +++ b/packages/ai-harness/src/define.ts @@ -0,0 +1,380 @@ +import { assertMetadataSchema } from './session-entry' +import type { + AgentLoopStrategy, + AnyChatMiddleware, + AnyTextAdapter, + AnyTool, + InterruptDefinition, + KeyedAdapter, + ModelMessage, + PromptCacheOptions, + ReasoningOption, + SchemaInput, + SubagentRouterPick, + SubagentsBag, + SystemPrompt, +} from '@tanstack/ai' +import type { AnyAgent } from './agents' +import type { MediaOptions } from './media' +import type { HarnessPlugin, PluginMetadataCheck } from './plugins' +import type { HarnessSession } from './session' +import type { MetadataOf } from './session-entry' +import type { HarnessTurnOptions, RecoverHook } from './turn' +import type { BusyPolicy, Principal, ThreadSettings, UserInput } from './types' + +/** + * The `subagents` option: the same as `chat({ subagents })`, without + * `binding` and `keys`. The session binds its own provider keys. + */ +export type HarnessSubagents> = Omit< + SubagentsBag, + 'binding' | 'keys' +> + +/** + * What `routing.router` gets: the fields of the `subagents.router` context, + * plus the harness fields of the turn. + */ +export interface HarnessRouterContext { + /** The model messages of the turn: the history, then the new user message. */ + messages: Array + /** + * The root agents of this turn, in registration order: the harness + * `agents`, then the `agents` of session plugins, then of turn plugins. An + * agent that is only in `subagents` is not in this list. + */ + agents: ReadonlyArray + /** Aborted when the turn is cancelled. */ + abortSignal?: AbortSignal + session: SessionWithMetadata + /** + * The user input of the turn as sent: text, or content parts. `undefined` + * for a turn with no new user message. + */ + input?: UserInput + /** The id of the turn operation. */ + operationId: string + /** The id of the turn input, when the turn has one. */ + inputId?: string + /** Who sent the turn input. Default: the principal that opened the session. */ + principal?: Principal + /** The `context` of the turn input. Client data: do not trust it. */ + context?: unknown + /** The main text adapter of the turn, with the user's keys bound. */ + adapter: AnyTextAdapter +} + +/** + * The `routing` option. The other fields mean the same as in `subagents`. + * `limits` defaults to `subagents.limits`. With `strategy: 'handoff'`, the + * main model answers after the agents, in the same turn, and it keeps its + * `subagents`. + */ +export interface HarnessRouting extends Pick< + SubagentsBag, + 'order' | 'strategy' | 'limits' | 'sandbox' +> { + /** + * Pick who answers the turn, with the picks of `subagents.router`: + * `'main'`, an agent name, a list of names, `{ names, order }`, or + * `{ steps }`. Give an agent with `inputSchema` its input with + * `{ name, input }`. A throw fails the turn. + */ + router: ( + ctx: HarnessRouterContext, + ) => SubagentRouterPick | Promise +} + +/** + * Limits for one input on a durable host (a host with `stores.log`). A host + * without a log reads only `interruptedToolResult` and `truncatedToolResult`. + */ +export interface HarnessDurability { + /** + * Total attempts for one input, the first run included. Each recovery after + * a crash is one more attempt. When no attempt is left, the input fails + * with code `'attempts_exhausted'`. Default 10. + */ + maxAttempts?: number + /** + * Time from the first run of an input until it fails with code + * `'timeout'`, in milliseconds. Retries do not reset it. Default: no + * timeout. + */ + timeoutMs?: number + /** + * Decide how an input that a crashed host left recovers: an input that + * never ran, and a turn whose lease expired. The context has the decision + * the harness takes by default. Return `undefined` to keep it. Return + * `{ action: 'run', overrides }` to run the turn with `TurnOverrides`. + */ + recover?: RecoverHook + /** + * The tool result that recovery gives a `replay: 'never'` tool call that + * a crash cut: the content and the error of its tool message. Default: a + * JSON note that the tool may or may not have run. + */ + interruptedToolResult?: string + /** + * The tool result of each tool call of an answer that stopped at the output + * limit (finish reason `length`): the content and the error of its tool + * message. The function form gets the call. Such a call never runs, also + * with `replay: 'safe'`. When set, the answer keeps its calls, each cut + * call gets this text, and the turn calls the model again. Recovery gives + * the same text to a cut call with no result in the log. Not set: a cut + * answer ends the turn, and recovery uses the default. Default: `'The answer was cut off at the output limit before + * this tool call was complete. The call did not run.'` + */ + truncatedToolResult?: + | string + | ((call: { toolCallId: string; toolName: string }) => string) + /** + * Continue an answer that a crash cut. Before the recovered attempt runs, + * recovery adds the part of the answer that the log has as an assistant + * message, then one user message for each note, in order, in one append. + * The model goes on after the notes. The assistant message has the answer + * text and each signed thinking block (a finished block with its + * signature, as the transcript keeps thinking), so the provider can replay + * it. A thinking block with no signature and tool calls are not kept. An + * attempt that streamed no answer text and no signed thinking (for example + * only thinking with no signature) runs again with no partial and no note. + * Signed thinking alone is enough to continue. A string `note` is one + * note. Default `false`. The default note: `'The previous answer was cut + * off. Continue exactly where it stopped, without repeating it.'` + */ + continueCutOff?: boolean | { note?: string | ReadonlyArray } +} + +/** + * What `defineHarness` takes. Where it overlaps with `chat()`, the option + * names and types are the same. + */ +export interface HarnessConfig< + TAdapter extends AnyTextAdapter = AnyTextAdapter, + TAgents extends ReadonlyArray = ReadonlyArray, + TSubagents extends ReadonlyArray = ReadonlyArray, + TMetadataSchema = undefined, +> { + /** A stable name, for example `'acme/studio'`. */ + name: string + /** What the harness does. Shown when another agent can call it. */ + description?: string + /** + * The main agent-loop model, the same as `chat({ adapter })`. A + * `keyedAdapter(...)` is built for each turn with the user's own key: the + * key saved with `/connect `, else the provider's env var. + * Leave it out when a plugin `adapter()` picks the model of every turn. A + * turn that gets no model from any of them fails. + */ + adapter?: TAdapter | KeyedAdapter + /** + * Models a thread can pick by name with `session.configure({ model })`. + * The stored name picks the model of each turn, after `overrides.adapter` + * and before a plugin `adapter()`. + */ + models?: Record> + systemPrompts?: Array + tools?: ReadonlyArray + middleware?: ReadonlyArray + /** + * When a turn stops calling the model. Defaults to `maxIterations(50)`. + * Without `durability.truncatedToolResult`, the turn also stops after an + * answer cut at the output limit. + */ + agentLoopStrategy?: AgentLoopStrategy + /** + * How the server tools of one model call run, the same as + * `chat({ toolExecution })`. `'parallel'` starts them together. + * `'sequential'` runs them one at a time, in call order. Default + * `'parallel'`. + */ + toolExecution?: 'parallel' | 'sequential' + modelOptions?: TAdapter['~types']['providerOptions'] + /** + * Automatic prompt caching for every session, the same as + * `chat({ promptCache })`. Default `'short'`. `'none'` turns it off. A + * `key` here is shared by all sessions, so usually leave it out: then each + * session uses its threadId as the key. + */ + promptCache?: PromptCacheOptions + /** + * How hard the main model thinks on every turn, as the adapter gets it in + * `TextOptions`. Not set: no reasoning is sent, and the provider default + * applies. `overrides.reasoning` of a turn replaces it. + */ + reasoning?: ReasoningOption + interrupts?: ReadonlyArray> + /** + * Runtime context passed to middleware hooks and server tools. When this + * and the `context` of a turn input are both plain objects, a turn gets + * both merged, and this value wins for a key in both, so a client cannot + * replace a server value. Otherwise a turn gets this value when it is set, + * else the input's. + */ + context?: unknown + /** + * Typed agents the session can run from code, commands, plugins, and + * exposed clients: `session.agents..run(input)`. + */ + agents?: TAgents + /** + * Agents the main model can call as tools, the same as `chat({ subagents })`. + * They are also registered in `session.agents`. + */ + subagents?: HarnessSubagents + /** + * A router for each new turn. It sends the turn to the main model, or to + * root agents: the harness `agents` and the `agents` of plugins. A resolve + * continues the turn it answers, so it does not call the router. + */ + routing?: HarnessRouting + /** Called once per session, so every session gets fresh plugin instances. */ + plugins?: () => ReadonlyArray< + HarnessPlugin & PluginMetadataCheck> + > + /** + * The shape of the `metadata` of the session index entry of each thread, + * as a Standard Schema (for example a zod object). Each write checks it: + * `host.sessions.update`, the `update` op of `POST sessions`, and + * `updateEntry` of plugins and sessions. A write it refuses throws + * `SessionMetadataError` and writes nothing. It also types the metadata + * of those writes, of `session.entry()`, of the hooks, and of a client + * made with `createHarnessClient`. + */ + metadataSchema?: TMetadataSchema + /** What a `prompt` does while a chat turn runs. Default `'queue'`. */ + busy?: BusyPolicy + /** Attempt and time limits for each input on a durable host. */ + durability?: HarnessDurability + /** Hooks that control a chat turn: retries, joins, and the end of a turn. */ + turn?: HarnessTurnOptions + /** + * Files sent to a turn, and media agents make: the size limit, the kinds a + * user can send, the kinds the model reads, and an optional transcriber for + * audio the model cannot read. + */ + media?: MediaOptions + /** What clients may call. Nothing is exposed by default. */ + expose?: { + agents?: ReadonlyArray + /** + * The thread settings a client may change with a `configure` input. An + * input with any other field is refused with `not_exposed`. Server code + * that calls `session.configure()` can change every field. `plugins` and + * `instructions` let a client turn off a plugin and write into the + * system prompt, so expose them only to clients you trust. + */ + settings?: ReadonlyArray + /** + * The plugin config keys a client may set with a `config` input, for + * example `client.setConfig('model', 'fast')`. An input with any other + * key is refused with `not_exposed`. Server code that calls + * `session.setConfig()` can set every key. Warning: `mode` of + * `permissions()` lets a client turn approvals off with `bypass`. Keep + * `mode` on the server, or expose it only to clients you trust. + */ + config?: ReadonlyArray + /** + * The commands a client may run with a `command` input, for example + * `client.command('undo')`, or `/undo` in a session view on a client. + * An input with any other command is refused with `not_exposed`. Server + * code that calls `session.command()` can run every command. Warning: + * `/mode bypass` of `permissions()` turns approvals off, so expose `mode` + * only to clients you trust. Plugins add their commands when a session + * opens, so `defineHarness` does not check these names. + */ + commands?: ReadonlyArray + } +} + +const HARNESS_KIND = 'tanstack-ai-harness' as const + +/** An immutable harness definition. Importing it starts nothing. */ +export type HarnessDefinition< + TAdapter extends AnyTextAdapter = AnyTextAdapter, + TAgents extends ReadonlyArray = ReadonlyArray, + TSubagents extends ReadonlyArray = ReadonlyArray, + TMetadataSchema = undefined, +> = Readonly> & { + readonly kind: typeof HARNESS_KIND + readonly version: 1 +} + +/** A harness definition with any type parameters. */ +export type AnyHarness = HarnessDefinition + +/** A session whose harness has the metadata schema `TMetadataSchema`. */ +export type SessionWithMetadata = HarnessSession< + HarnessDefinition +> + +/** The type of the session index `metadata` of a harness. */ +export type SessionMetadataOf = + THarness extends HarnessDefinition + ? MetadataOf + : Record + +/** Every agent a harness registers, as a union. */ +export type HarnessAgentsOf = + THarness extends HarnessDefinition + ? TAgents[number] | TSubagents[number] + : never + +/** True for a value made with `defineHarness`. */ +export function isHarnessDefinition(value: unknown): value is AnyHarness { + return ( + typeof value === 'object' && + value !== null && + (value as { kind?: unknown }).kind === HARNESS_KIND && + (value as { version?: unknown }).version === 1 + ) +} + +/** + * Define a harness: a reusable, typed agent configuration. Use the same option + * names as `chat()`, plus `agents`, `plugins`, `busy`, `expose`, and `media`. + * + * @example + * ```ts + * const studio = defineHarness({ + * name: 'acme/studio', + * adapter: anthropicText('claude-sonnet-4-5'), + * agents: [heroImage], + * subagents: { agents: [researcher] }, + * }) + * ``` + */ +export function defineHarness< + TAdapter extends AnyTextAdapter = AnyTextAdapter, + const TAgents extends ReadonlyArray = readonly [], + const TSubagents extends ReadonlyArray = readonly [], + TMetadataSchema extends SchemaInput | undefined = undefined, +>( + config: HarnessConfig, +): HarnessDefinition { + if (config.name.trim() === '') { + throw new Error('defineHarness requires a non-empty name') + } + assertMetadataSchema(`defineHarness "${config.name}"`, config.metadataSchema) + const byName = new Map() + for (const agent of [ + ...(config.agents ?? []), + ...(config.subagents?.agents ?? []), + ]) { + const existing = byName.get(agent.name) + if (existing && existing !== agent) { + throw new Error( + `defineHarness "${config.name}": two different agents are named "${agent.name}".`, + ) + } + byName.set(agent.name, agent) + } + for (const name of config.expose?.agents ?? []) { + if (!byName.has(name)) { + throw new Error( + `defineHarness "${config.name}": expose.agents names "${name}", which is not a registered agent.`, + ) + } + } + return Object.freeze({ ...config, kind: HARNESS_KIND, version: 1 as const }) +} diff --git a/packages/ai-harness/src/extensions.ts b/packages/ai-harness/src/extensions.ts new file mode 100644 index 0000000000..3181822750 --- /dev/null +++ b/packages/ai-harness/src/extensions.ts @@ -0,0 +1,40 @@ +/** + * A list one plugin reads and other plugins add to, in any order. For + * example, a permissions plugin reads rules that tool plugins contribute. + * + * @example + * ```ts + * export const PermissionRules = createExtensionPoint('acme/permission-rules') + * // a tool plugin: contribute: [PermissionRules.item({ tool: 'write_file', decision: 'ask' })] + * // the permissions plugin, at run time: ctx.collect(PermissionRules) + * ``` + */ +export interface ExtensionPoint { + readonly name: string + item: (value: T) => ExtensionItem +} + +/** One contribution to an extension point. */ +export interface ExtensionItem { + readonly point: string + readonly value: T +} + +export function createExtensionPoint(name: string): ExtensionPoint { + return { name, item: (value) => ({ point: name, value }) } +} + +/** + * A typed plugin event. `ctx.emit(event, value)` sends it to every + * `ctx.on(event, handler)` in the session, and clients see it as a + * `harness.plugin.event` CUSTOM event. + */ +export interface PluginEvent { + readonly name: string + /** Type only. Never read at runtime. */ + readonly __type?: T +} + +export function createPluginEvent(name: string): PluginEvent { + return { name } +} diff --git a/packages/ai-harness/src/feed.ts b/packages/ai-harness/src/feed.ts new file mode 100644 index 0000000000..76280a4411 --- /dev/null +++ b/packages/ai-harness/src/feed.ts @@ -0,0 +1,117 @@ +import type { StreamChunk } from '@tanstack/ai' +import type { Cursor, SessionEvent } from './types' + +// ponytail: in-memory log capped at MAX_EVENTS per session. A reader whose +// cursor fell out of the window resumes from the oldest kept event. Durable +// cross-process delivery goes through StreamDurability in the protocol layer. +const MAX_EVENTS = 10_000 + +/** + * The ordered event stream of one session. Every operation publishes here. A + * host without a log uses {@link SessionFeed}. A durable host uses a feed + * that keeps the events in the session log. + */ +export interface EventFeed { + publish: (operationId: string, event: StreamChunk) => void + /** The cursor of the newest event, or `'0'` for an empty feed. */ + head: () => Cursor + /** + * Events after `from` (exclusive), then live events until `signal` aborts + * or the feed closes. Pass a `filter` to read one operation's events. + */ + read: (options: { + from?: Cursor + signal?: AbortSignal + filter?: (entry: SessionEvent) => boolean + /** Stop after `until()` returns true and no buffered event is left. */ + until?: () => boolean + }) => AsyncIterable + close: () => void +} + +/** + * The in-memory {@link EventFeed}. Cursors are opaque to callers and increase + * with each event. + */ +export class SessionFeed implements EventFeed { + private readonly entries: Array = [] + private sequence = 0 + private readonly waiters = new Set<() => void>() + private closed = false + + publish(operationId: string, event: StreamChunk): void { + this.sequence += 1 + this.entries.push({ cursor: String(this.sequence), operationId, event }) + if (this.entries.length > MAX_EVENTS) this.entries.shift() + this.wake() + } + + head(): Cursor { + return String(this.sequence) + } + + async *read(options: { + from?: Cursor + signal?: AbortSignal + filter?: (entry: SessionEvent) => boolean + /** Stop after `until()` returns true and no buffered event is left. */ + until?: () => boolean + }): AsyncIterable { + let after = Number(options.from ?? '0') + if (!Number.isFinite(after)) after = 0 + while (true) { + if (options.signal?.aborted) return + const next = this.entries.filter( + (entry) => + Number(entry.cursor) > after && + (options.filter ? options.filter(entry) : true), + ) + for (const entry of next) { + if (options.signal?.aborted) return + after = Number(entry.cursor) + yield entry + } + if (next.length === 0) { + // Skip past events the filter rejected, so they are not scanned again. + after = Math.max(after, this.sequence) + } + if (this.closed || options.until?.()) { + const rest = this.entries.filter( + (entry) => + Number(entry.cursor) > after && + (options.filter ? options.filter(entry) : true), + ) + for (const entry of rest) yield entry + return + } + // Events published while the reader handled `next` woke nobody, so + // scan again before waiting. An empty scan moves `after` to the head. + if (next.length > 0) continue + await this.waitForNext(options.signal) + } + } + + close(): void { + this.closed = true + this.wake() + } + + private wake(): void { + const waiters = [...this.waiters] + this.waiters.clear() + for (const wake of waiters) wake() + } + + private waitForNext(signal?: AbortSignal): Promise { + return new Promise((resolve) => { + if (signal?.aborted) return resolve() + const done = () => { + signal?.removeEventListener('abort', done) + this.waiters.delete(done) + resolve() + } + this.waiters.add(done) + signal?.addEventListener('abort', done, { once: true }) + }) + } +} diff --git a/packages/ai-harness/src/index.ts b/packages/ai-harness/src/index.ts new file mode 100644 index 0000000000..a6de481a21 --- /dev/null +++ b/packages/ai-harness/src/index.ts @@ -0,0 +1,170 @@ +export { defineHarness, isHarnessDefinition } from './define' +export type { + AnyHarness, + HarnessAgentsOf, + HarnessConfig, + HarnessDefinition, + HarnessDurability, + HarnessRouterContext, + HarnessRouting, + HarnessSubagents, + SessionMetadataOf, + SessionWithMetadata, +} from './define' +export { SessionMetadataError } from './session-entry' +export type { + MetadataOf, + SessionEntryChange, + SessionEntryUpdate, +} from './session-entry' + +export { durableTool } from './durable-tool' +export type { + DurableToolContext, + DurableToolOptions, + ToolStep, +} from './durable-tool' + +export { logMessageStore } from './log' +export type { ProjectOptions, ProjectRecord, ReduceOptions } from './log' +export type { LeaseOptions } from './resume' +export type { + LeaseStore, + TurnLease, + TurnLeaseKey, +} from '@tanstack/ai-persistence' + +export { definePlugin } from './plugins' +export type { + HarnessPlugin, + PluginContributions, + PluginDefinition, + PluginLifetime, + PluginPrompt, + AgentGroup, + PluginAgentActions, + PluginCommands, + PluginSetupContext, + PluginState, +} from './plugins' + +export { createExtensionPoint, createPluginEvent } from './extensions' +export type { ExtensionItem, ExtensionPoint, PluginEvent } from './extensions' + +export { checkConfigValue, configOption } from './config' +export type { ConfigOption } from './config' + +export { defineCommand } from './commands' +export type { + AnswerOf, + AnyCommand, + CommandContext, + CommandDefinition, + PluginSessionApi, + Question, +} from './commands' + +export { AuthRequiredError, scrubSecrets } from './auth' +export type { CredentialsAccess } from './auth' + +export { harnessContext } from './agents' +export type { + AgentInputOf, + AgentRegistryView, + AgentResultOf, + AgentStarter, + AnyAgent, + HarnessAgentContext, +} from './agents' + +export { createHarnessHost } from './host' +export type { + ForkSessionOptions, + HarnessHost, + HarnessHostOptions, + HarnessPersistence, + HostEvent, + HostSessionDeletedEvent, + HostSessionEvent, + HostSessions, + HostStatusEvent, + OpenSessionOptions, + ResumePendingOptions, +} from './host' +export type { + SessionIndexEntry, + SessionIndexListOptions, + SessionIndexPage, + SessionIndexStore, +} from '@tanstack/ai-persistence' + +export { HarnessSession } from './session' +export type { + AgentHandle, + AgentHandles, + AgentRunInfo, + AgentRunOptions, + AgentStartOptions, + DynamicAgentHandle, + SessionDescription, + SessionInspection, + SessionSnapshot, +} from './session' +export type { SessionUsage, UsageCounts } from './usage' + +export type { + FinishContext, + HarnessTurnOptions, + JoinCandidate, + JoinContext, + ModelErrorContext, + RecoverContext, + RecoverDecision, + RecoverHook, + TurnAdditions, +} from './turn' +export { isTransientModelError, retryTransientErrors } from './turn' + +export { HARNESS_EVENTS, InputRejectedError } from './types' +export type { + AgentRun, + BusyPolicy, + ChatTurnResult, + Cursor, + ForkPoint, + HarnessInput, + InputSettlement, + MediaKind, + MediaRecord, + Operation, + OperationKind, + OperationStatus, + Principal, + Receipt, + SessionEvent, + ThreadSettings, + ThreadSettingsChange, + TurnInfo, + TurnOverrides, + UserInput, + WaitingInput, +} from './types' + +export { + MEDIA_URL_PREFIX, + isMediaRecord, + kindOf, + mediaIdOf, + mediaOfMessage, + mediaPart, + mimeTypeOf, +} from './media-ref' + +export { MediaError } from './media' +export type { MediaOptions } from './media' + +export { harnessText } from './harness-text' +export type { HarnessTextOptions, RemoteHarness } from './harness-text' + +export { harnessAgent } from './harness-agent' +export { DEFAULT_SUBAGENT_LIMITS } from './session' diff --git a/packages/ai-harness/src/media-ref.ts b/packages/ai-harness/src/media-ref.ts new file mode 100644 index 0000000000..b3806cecad --- /dev/null +++ b/packages/ai-harness/src/media-ref.ts @@ -0,0 +1,166 @@ +// Browser-safe media references. Type-only imports, no server code, so the +// client and view entries can use these too. +import type { ContentPart, ModelMessage } from '@tanstack/ai' +import type { MediaKind, MediaRecord } from './types' + +/** The URL scheme of a content part that points to a harness media file. */ +export const MEDIA_URL_PREFIX = 'harness-media:' + +function isRecord(value: unknown): value is Record { + return typeof value === 'object' && value !== null +} + +function isMediaKind(value: unknown): value is MediaKind { + return ( + value === 'image' || + value === 'audio' || + value === 'video' || + value === 'document' + ) +} + +function isOptionalString(value: unknown) { + return value === undefined || typeof value === 'string' +} + +/** + * True for a value with the shape of a `MediaRecord`, for example the answer + * of an upload or the value of a `harness.media` event. + * + * @example + * const value: unknown = await response.json() + * if (isMediaRecord(value)) console.log(value.name) + */ +export function isMediaRecord(value: unknown): value is MediaRecord { + if (!isRecord(value)) return false + const hasOptionalIds = + isOptionalString(value.runId) && isOptionalString(value.subagentRunId) + return ( + hasOptionalIds && + typeof value.id === 'string' && + typeof value.threadId === 'string' && + isMediaKind(value.kind) && + typeof value.mimeType === 'string' && + typeof value.name === 'string' && + typeof value.size === 'number' && + (value.source === 'user' || value.source === 'generated') && + typeof value.createdAt === 'number' + ) +} + +/** + * The media kind of a MIME type, or `undefined` when the harness does not + * store that type. `image/*`, `audio/*`, and `video/*` map to their kind. + * `application/pdf` and `text/*` are documents. + * + * @example + * kindOf('image/png') // 'image' + * kindOf('application/zip') // undefined + */ +export function kindOf(mimeType: string) { + // A Content-Type header can carry parameters: `text/plain; charset=utf-8`. + const [type = ''] = mimeType.toLowerCase().split(';') + const [top = '', sub] = type.trim().split('/') + if (!sub) return undefined + switch (top) { + case 'image': + case 'audio': + case 'video': + return top + case 'text': + return 'document' + case 'application': + return sub === 'pdf' ? 'document' : undefined + default: + return undefined + } +} + +// ponytail: common types only. Send any other file with its MIME type. +const mimeTypes = new Map([ + ['png', 'image/png'], + ['jpg', 'image/jpeg'], + ['jpeg', 'image/jpeg'], + ['gif', 'image/gif'], + ['webp', 'image/webp'], + ['mp3', 'audio/mpeg'], + ['wav', 'audio/wav'], + ['ogg', 'audio/ogg'], + ['m4a', 'audio/mp4'], + ['flac', 'audio/flac'], + ['mp4', 'video/mp4'], + ['webm', 'video/webm'], + ['mov', 'video/quicktime'], + ['pdf', 'application/pdf'], + ['txt', 'text/plain'], + ['md', 'text/markdown'], + ['csv', 'text/csv'], + ['html', 'text/html'], +]) + +/** + * The MIME type of a file name, a path, or a URL path, by its extension. + * Returns `undefined` for no extension or an extension it does not know. + * It knows the common image, audio, video, PDF, and text types. + * + * @example + * mimeTypeOf('cat.PNG') // 'image/png' + * mimeTypeOf('notes') // undefined + */ +export function mimeTypeOf(fileName: string) { + const extension = /\.([^./\\]+)$/.exec(fileName)?.[1] + return extension === undefined + ? undefined + : mimeTypes.get(extension.toLowerCase()) +} + +/** + * The content part that sends a stored media file to a turn. Its source is + * the URL `harness-media:`. The harness swaps it for the bytes only when + * it calls the model, so the transcript stays small. + * + * @example + * session.prompt([{ type: 'text', content: 'What is this?' }, mediaPart(record)]) + */ +export function mediaPart(record: MediaRecord) { + const part: ContentPart = { + type: record.kind, + source: { + type: 'url', + value: `${MEDIA_URL_PREFIX}${record.id}`, + mimeType: record.mimeType, + }, + } + return part +} + +/** + * The media id of a part made by `mediaPart`, or `undefined` for any other + * value (text, a data source, a normal URL). + * + * @example + * const ids = parts.map(mediaIdOf).filter((id) => id !== undefined) + */ +export function mediaIdOf(part: unknown) { + if (!isRecord(part) || !isMediaKind(part.type) || !isRecord(part.source)) { + return undefined + } + const { type, value } = part.source + if (type !== 'url' || typeof value !== 'string') return undefined + if (!value.startsWith(MEDIA_URL_PREFIX)) return undefined + return value.slice(MEDIA_URL_PREFIX.length) || undefined +} + +/** + * The media records saved on a message in `metadata.harness.media`: the + * media a turn made, kept on its last assistant message. Entries with a bad + * shape are skipped. Returns `[]` when there are none. + * + * @example + * for (const media of mediaOfMessage(message)) console.log(media.name) + */ +export function mediaOfMessage(message: ModelMessage) { + const harness: unknown = message.metadata?.harness + const media = isRecord(harness) ? harness.media : undefined + return Array.isArray(media) ? media.filter(isMediaRecord) : [] +} diff --git a/packages/ai-harness/src/operation.ts b/packages/ai-harness/src/operation.ts new file mode 100644 index 0000000000..fad51fe2b6 --- /dev/null +++ b/packages/ai-harness/src/operation.ts @@ -0,0 +1,124 @@ +import type { StreamChunk } from '@tanstack/ai' +import type { EventFeed } from './feed' +import type { + Cursor, + Operation, + OperationKind, + OperationStatus, + Receipt, + SessionEvent, +} from './types' + +let counter = 0 + +/** A new operation id. It is also the AG-UI `runId` of the operation. */ +export function createOperationId(kind: OperationKind): string { + counter += 1 + return `op-${kind}-${Date.now().toString(36)}-${counter}` +} + +const TERMINAL: ReadonlySet = new Set([ + 'completed', + 'failed', + 'cancelled', + 'interrupted', +]) + +/** The session-side implementation of {@link Operation}. */ +export class OperationImpl implements Operation { + readonly id: string + /** The feed head when the operation was made. Its events come after it. */ + readonly startedCursor: Cursor + readonly abortController = new AbortController() + readonly receipt: Promise + private current: OperationStatus = 'accepted' + private readonly settled: Promise + private resolveResult!: (value: TResult) => void + private rejectResult!: (error: unknown) => void + private settleReceipt!: (receipt: Receipt) => void + + constructor( + readonly kind: OperationKind, + private readonly feed: EventFeed, + private readonly onCancel: ( + operation: OperationImpl, + ) => Promise, + readonly agent?: string, + /** A stored operation id, when a restart rebuilds this operation. */ + id?: string, + ) { + this.id = id ?? createOperationId(kind) + this.startedCursor = feed.head() + this.receipt = new Promise((resolve) => { + this.settleReceipt = resolve + }) + this.settled = new Promise((resolve, reject) => { + this.resolveResult = resolve + this.rejectResult = reject + }) + // Nobody may await a failed background operation. Do not report that as an + // unhandled rejection. `then` still rejects for callers that await. + this.settled.catch(() => {}) + } + + then( + onfulfilled?: ((value: TResult) => TResult1 | PromiseLike) | null, + onrejected?: ((reason: unknown) => TResult2 | PromiseLike) | null, + ): PromiseLike { + return this.settled.then(onfulfilled, onrejected) + } + + status(): OperationStatus { + return this.current + } + + isSettled(): boolean { + return TERMINAL.has(this.current) + } + + setStatus(status: OperationStatus): void { + this.current = status + } + + publish(event: StreamChunk): void { + this.feed.publish(this.id, event) + } + + /** Answer `receipt`. The first answer wins. */ + resolveReceipt(receipt: Receipt): void { + this.settleReceipt(receipt) + } + + finish(status: 'completed' | 'interrupted', result: TResult): void { + this.current = status + this.resolveResult(result) + } + + fail(status: 'failed' | 'cancelled', error: unknown): void { + this.current = status + this.rejectResult(error) + } + + events(options?: { + from?: Cursor + signal?: AbortSignal + }): AsyncIterable { + return this.feed.read({ + ...options, + filter: (entry) => entry.operationId === this.id, + until: () => this.isSettled(), + }) + } + + async *stream(options?: { + signal?: AbortSignal + }): AsyncIterable { + for await (const entry of this.events({ from: '0', ...options })) { + yield entry.event + } + } + + cancel(_reason?: string): Promise { + return this.onCancel(this) + } +} diff --git a/packages/ai-harness/src/plugins.ts b/packages/ai-harness/src/plugins.ts new file mode 100644 index 0000000000..e6dc4f6a63 --- /dev/null +++ b/packages/ai-harness/src/plugins.ts @@ -0,0 +1,967 @@ +import { CapabilityRegistry, createCapability } from '@tanstack/ai' +import { ResourceScope, disposeAll } from './resources' +import { assertMetadataSchema } from './session-entry' +import type { + AgentProduces, + AnyChatMiddleware, + AnyGenerationMiddleware, + AnyTextAdapter, + AnyTool, + Capability, + CapabilityHandle, + KeyedAdapter, + MetadataStore, + ModelMessage, + ProviderKeys, + SchemaInput, +} from '@tanstack/ai' +import type { + AgentInputOf, + AgentRegistry, + AgentRegistryView, + AgentResultOf, + AnyAgent, +} from './agents' +import type { AgentRun, Operation, TurnInfo } from './types' +import type { CredentialsAccess } from './auth' +import type { AnyCommand, PluginSessionApi } from './commands' +import type { ConfigOption } from './config' +import type { ExtensionItem, ExtensionPoint, PluginEvent } from './extensions' +import type { MetadataOf } from './session-entry' + +/** + * 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. + * - `run`: deprecated, the old name of `turn`. It works the same. + */ +export type PluginLifetime = 'session' | 'turn' | 'run' + +/** + * A prompt a plugin contributes. `id` must be unique in the session. A + * function `text` is called for each turn, so it can show current state. + */ +export interface PluginPrompt { + id: string + text: string | (() => string) +} + +/** What a plugin's `setup` returns. Every field is optional. */ +export interface PluginContributions { + /** Tools for the main model. */ + tools?: ReadonlyArray + /** System prompt text for every chat turn. A function runs for each turn. */ + prompts?: ReadonlyArray string) | PluginPrompt> + /** Chat middleware, the same type as `chat({ middleware })`. */ + middleware?: ReadonlyArray + /** + * Chat middleware for every agent run: subagents, background agents, and + * their children. Not the lead turn: add the same middleware to + * `middleware` for that. + */ + agentMiddleware?: ReadonlyArray + /** Middleware for the activities agents call (`ctx.generateImage`, ...). */ + generationMiddleware?: ReadonlyArray + /** Agents added to `session.agents`. `routing.router` can pick them. */ + agents?: ReadonlyArray + /** + * Agents the main model can call as tools, merged into the turn's + * `chat({ subagents })` after `defineHarness({ subagents })`. They are + * also added to `session.agents`. + */ + subagents?: ReadonlyArray + /** User actions, keyed by name. Run with `session.command(name, input)`. */ + commands?: Record + /** Session settings, keyed by name. Read with `ctx.config.get(name)`. */ + config?: Record + /** Items for extension points that other plugins read. */ + contribute?: ReadonlyArray + /** + * Pick the main-loop adapter for the next turn, or return `undefined` to + * keep the harness adapter. The last plugin that returns one wins. The + * session builds a `keyedAdapter(...)` with the user's key. `turn` is the + * turn that starts: its input, context, and sender. Its + * `overrides.adapter` wins over every pick. + */ + adapter?: ( + turn: TurnInfo, + ) => AnyTextAdapter | KeyedAdapter | undefined + /** + * Tools found at run time, for example the tools of an MCP server the user + * signed in to after the session opened. Called before each chat turn. A + * name that another tool already uses is skipped. + */ + discoverTools?: () => ReadonlyArray | Promise> + /** + * Change the tool list of a turn: every tool the model gets, after + * `discoverTools`. Return the new list. Runs before prompts resolve, so a + * prompt `text()` can describe the tools this returned. Code mode uses it to + * move tools behind `execute_typescript`. `model` is the model id of the + * turn's adapter, so the list can differ for each model. + */ + prepareTools?: (turn: { + tools: ReadonlyArray + model: string + }) => ReadonlyArray | Promise> +} + +/** Plugin state that survives restarts, stored in the metadata store. */ +export interface PluginState { + get: () => Promise + /** Change the state. On a write conflict, `change` runs again with fresh state. */ + update: (change: (current: T) => T) => Promise +} + +/** Run agents from a plugin: in the foreground, in the background, or as a group. */ +export interface PluginAgentActions { + run: { + ( + agent: TAgent, + input?: AgentInputOf, + ): Operation> + (name: string, input?: unknown): Operation + } + start: { + ( + agent: TAgent, + input?: AgentInputOf, + /** See `AgentStartOptions`. */ + options?: { wake?: boolean; resume?: boolean }, + ): AgentRun> + ( + name: string, + input?: unknown, + /** See `AgentStartOptions`. */ + options?: { wake?: boolean; resume?: boolean }, + ): AgentRun + } + /** + * Run children together. With `onFailure: 'cancel-siblings'` (default), one + * failure cancels the others. Every child settles before `group` returns. + */ + group: ( + options: { onFailure?: 'cancel-siblings' | 'collect' }, + body: (group: AgentGroup) => Promise, + ) => Promise +} + +/** The children of one `ctx.agents.group` call. */ +export interface AgentGroup { + run: { + ( + agent: TAgent, + input?: AgentInputOf, + ): Promise> + (name: string, input?: unknown): Promise + } + /** Like `run`, but never rejects: resolves with the result or the error. */ + runSettled: { + ( + agent: TAgent, + input?: AgentInputOf, + ): Promise< + { ok: true; value: AgentResultOf } | { ok: false; error: unknown } + > + ( + name: string, + input?: unknown, + ): Promise<{ ok: true; value: unknown } | { ok: false; error: unknown }> + } +} + +/** What the session gives plugins. */ +export interface PluginServices { + emit: (plugin: string, name: string, value: unknown) => void + on: (name: string, handler: (value: unknown) => void) => () => void + config: { get: (key: string) => unknown } + state: (plugin: string, initial: T) => PluginState + credentials: CredentialsAccess + keys: ProviderKeys + session: PluginSessionApi + /** The session API of a plugin with a `metadataSchema`: its writes check it too. */ + sessionWith: (metadataSchema: SchemaInput) => PluginSessionApi + agents: PluginAgentActions + /** Called after `ctx.commands` adds or removes a command. */ + commandsChanged?: () => void +} + +/** + * The commands of the session, for a plugin that adds and removes its own + * commands while the session runs, for example one command for each file in + * a folder. Commands that never change go in the `commands` of `setup`. + */ +export interface PluginCommands { + /** Is `name` a command of any plugin of this session? */ + has: (name: string) => boolean + /** + * Add this plugin's command `name`, or replace it. A name that another + * plugin owns throws. Session views get the new list at once. + */ + set: (name: string, command: AnyCommand) => void + /** Remove this plugin's command `name`. Another plugin's name does nothing. */ + delete: (name: string) => void + /** + * Resolves when every plugin of the session is set up. Wait for it before + * `has` or `set`: during `setup`, the plugins after this one have not + * added their commands yet. + */ + ready: Promise +} + +/** + * What a plugin's `setup` receives. `TMetadata` is the type of the session + * index `metadata`, from the `metadataSchema` of the plugin. + */ +export interface PluginSetupContext> { + /** Resources this plugin owns, with rollback and reverse-order cleanup. */ + resources: { + acquire: ( + open: () => T | Promise, + close: (resource: T) => unknown, + ) => Promise + signal: AbortSignal + } + /** Read a capability another plugin or middleware provided. Throws if absent. */ + get: (capability: Capability) => T + /** Read a capability, or `undefined` when nobody provided it. */ + getOptional: (capability: Capability) => T | undefined + /** Provide a capability this plugin declared in `provides`. */ + provide: (capability: Capability, value: T) => void + /** + * The agents of this session (harness agents plus earlier plugins' + * agents): find them, and run them from commands, tools, and hooks. + * `set` and `delete` change this plugin's agents while the session runs, + * like `ctx.commands`. + */ + agents: AgentRegistryView & + PluginAgentActions & { + /** + * Add this plugin's agent, or replace the one with the same name. A + * name that another plugin or the harness owns throws. + */ + set: ( + agent: AnyAgent, + options?: { + /** Also give the model a tool that runs this agent as a subagent. */ + subagent?: boolean + }, + ) => void + /** + * Remove this plugin's agent `name`, also from the subagent tools. + * Another owner's name does nothing. + */ + delete: (name: string) => void + } + /** + * The items other plugins contributed to `point`. The list fills while + * plugins set up, so read it at run time (in a tool, a command, or a + * middleware hook), not during `setup`. + */ + collect: (point: ExtensionPoint) => ReadonlyArray + /** Send a typed event to every plugin that listens, and to clients. */ + emit: (event: PluginEvent, value: T) => void + /** Listen for a typed event. Returns a function that stops listening. */ + on: (event: PluginEvent, handler: (value: T) => void) => () => void + /** This plugin's state, created with `initial` the first time. */ + state: (initial: T) => PluginState + /** Session settings. A change applies at the next turn. */ + config: { get: (key: string) => unknown } + /** + * Credentials of the running turn's sender. Outside a turn: of the + * principal that opened the session. They are read at each call, so code + * that runs after its turn ends (a tool of a background agent) gets the + * sender of the turn that runs at that time. There, use the keys of the + * agent run (`ctx.keys` in the agent) and, in a command, the `credentials` + * of its `run` context. + */ + credentials: CredentialsAccess + /** + * Model provider keys of the same principal as `credentials`: the key saved with + * `/connect `, else the provider's env var. Build a + * `keyedAdapter(...)` with `await ctx.keys.adapter(adapter)` just before + * the call. A missing key throws `AuthRequiredError`. + */ + keys: ProviderKeys + session: PluginSessionApi + /** Add and remove this plugin's commands while the session runs. */ + commands: PluginCommands + /** + * The turn this plugin is set up for: its input, context, and sender. Set + * only for a plugin with `lifetime: 'turn'`. A session plugin has no turn + * at setup. + */ + 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> { + /** A stable, unique name, for example `'acme/todos'`. */ + name: string + /** + * The shape of the session index `metadata` this plugin uses, as a + * Standard Schema. It types `ctx.session.entry()` and + * `ctx.session.updateEntry()`, and checks each write of this plugin. Pass + * the same schema as the harness `metadataSchema`: a harness whose + * 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`. */ + provides?: ReadonlyArray + /** Capabilities used when present. Never an error when missing. */ + optionalRequires?: ReadonlyArray + /** Agent outputs this plugin needs the session to have. */ + needs?: { produces?: ReadonlyArray } + setup?: ( + ctx: PluginSetupContext, + ) => PluginContributions | void | Promise +} + +const PLUGIN_BRAND = Symbol.for('tanstack.ai.harnessPlugin') + +/** + * A plugin made with {@link definePlugin}. `TMetadata` is the session + * metadata type it reads, from its `metadataSchema`. + */ +export type HarnessPlugin< + TDefinition extends PluginDefinition = PluginDefinition, + TMetadata = any, +> = Readonly & { + readonly [PLUGIN_BRAND]: true +} & PluginMetadataCheck + +/** + * @internal A type-only field: the metadata a plugin reads. The harness + * `plugins` option takes only plugins that can read the harness metadata. + */ +export interface PluginMetadataCheck { + readonly '~metadata'?: (metadata: TMetadata) => void +} + +/** + * Define a harness plugin. `setup` returns what the plugin contributes, and + * closures in it can use the resources it acquired. + * + * @example + * ```ts + * const today = definePlugin({ + * name: 'acme/today', + * setup: () => ({ prompts: [`Today is ${new Date().toDateString()}.`] }), + * }) + * ``` + */ +export function definePlugin< + TMetadataSchema extends SchemaInput | undefined = undefined, + const TDefinition extends PluginDefinition> = + PluginDefinition>, +>( + definition: TDefinition & { metadataSchema?: TMetadataSchema }, +): HarnessPlugin> { + if (definition.name.trim() === '') { + throw new Error('definePlugin requires a non-empty name') + } + assertMetadataSchema( + `definePlugin "${definition.name}"`, + definition.metadataSchema, + ) + return Object.freeze({ ...definition, [PLUGIN_BRAND]: true as const }) +} + +/** A tool, prompt, or middleware with the plugin that contributed it. */ +interface Owned { + value: T + owner: string +} + +/** Everything a set of mounted plugins contributes, plus its cleanup. */ +export interface MountedPlugins { + tools: Array + /** Strings, or functions to call for each turn. */ + prompts: Array string)> + commands: Map + config: Map + /** Contributions to extension points, by point name. */ + extensions: Map> + adapters: Array> + discoverers: Array<{ + discover: () => ReadonlyArray | Promise> + owner: string + }> + preparers: Array<{ + prepare: NonNullable + owner: string + }> + /** The `agents` of the plugins, in plugin order: root agents for routing. */ + agents: Array + /** Agents plugins give the main model, in plugin order. */ + subagents: Array + /** Who contributed what, for `session.inspect()`. */ + owners: { + plugins: Array<{ + name: string + lifetime: PluginLifetime + requires: Array + provides: Array + }> + tools: Array<{ name: string; owner: string }> + /** The owner of each item of `prompts`, at the same index. */ + prompts: Array<{ id: string; owner: string }> + /** The owner of each item of `middleware`, at the same index. */ + middleware: Array + /** The owner of each item of `adapters`, at the same index. */ + adapters: Array + } + middleware: Array + /** Chat middleware for every agent run. See `PluginContributions`. */ + agentMiddleware: Array + generationMiddleware: Array + /** + * Chat middleware that provides every plugin capability to each chat run, + * so any chat middleware can read it with `getX(ctx)`. Goes first. + */ + capabilityBridge: AnyChatMiddleware | undefined + /** The capability values, so a run mount can read session capabilities. */ + values: CapabilityValues + /** Dispose every plugin scope, newest first. Safe to call twice. */ + dispose: () => Promise +} + +export interface MountEnvironment { + threadId: string + registry: AgentRegistry + /** Names already taken by the harness itself. */ + harnessTools: ReadonlyArray + /** Capabilities provided by harness-level middleware. */ + harnessProvides: ReadonlyArray + /** Capability values from an outer mount (session plugins, for a run mount). */ + inherited?: CapabilityValues + /** Names taken by an outer mount (session plugins, for a run mount). */ + takenCommands?: ReadonlyMap + takenConfig?: ReadonlyMap + /** Extension items from an outer mount, visible to `collect`. */ + inheritedExtensions?: ReadonlyMap< + string, + ReadonlyArray<{ value: unknown; owner: string }> + > + /** Session services. Tests of the mount alone can leave them out. */ + services?: PluginServices + /** The turn of a run mount. Plugins read it as `ctx.turn`. */ + turn?: TurnInfo +} + +const unavailable = (what: string) => () => { + throw new Error(`${what} is only available inside a harness session.`) +} + +const NO_SESSION: PluginSessionApi = { + threadId: '', + principal: undefined, + snapshot: unavailable('ctx.session'), + prompt: unavailable('ctx.session'), + note: unavailable('ctx.session'), + transcript: unavailable('ctx.session'), + replaceTranscript: unavailable('ctx.session'), + entry: unavailable('ctx.session'), + updateEntry: unavailable('ctx.session'), + ask: unavailable('ctx.session'), + authRequired: unavailable('ctx.session'), + setConfig: unavailable('ctx.session'), + settings: unavailable('ctx.session'), +} + +const NO_SERVICES: PluginServices = { + emit: () => {}, + on: () => () => {}, + config: { get: () => undefined }, + state: unavailable('ctx.state'), + credentials: { + get: unavailable('ctx.credentials'), + require: unavailable('ctx.credentials'), + set: unavailable('ctx.credentials'), + delete: unavailable('ctx.credentials'), + list: unavailable('ctx.credentials'), + }, + keys: { + get: unavailable('ctx.keys'), + require: unavailable('ctx.keys'), + adapter: unavailable('ctx.keys'), + }, + session: NO_SESSION, + sessionWith: () => NO_SESSION, + agents: { + run: unavailable('ctx.agents.run') as PluginAgentActions['run'], + start: unavailable('ctx.agents.start') as PluginAgentActions['start'], + group: unavailable('ctx.agents.group'), + }, +} + +/** + * @internal The metadata store of the session, for first-party plugins that + * need it outside a chat run (in a command). The package does not export it. + */ +export const SessionMetadata = createCapability()( + 'tanstack/session-metadata', +) + +/** + * @internal How the snapshots plugin puts files back for `session.revert`. + * The package does not export it. + */ +export interface RevertFilesHandler { + /** + * Put back the files that the tool calls of `hidden` changed. Other files + * stay as they are. `records` reads the host records of one type from the + * session log. Resolves to what `unrevert` needs (JSON), or `undefined` + * when no file changed. + */ + revert: ( + hidden: ReadonlyArray, + records: (type: string) => Promise>>, + ) => Promise + /** Put the files back as they were before `revert`. */ + unrevert: (saved: unknown) => Promise +} + +/** @internal See {@link RevertFilesHandler}. */ +export const RevertFiles = createCapability()( + 'tanstack/revert-files', +) + +/** + * @internal Tells first-party plugins whether a revert stands. The session + * provides it. The package does not export it. + */ +export const RevertStanding = createCapability<() => boolean>()( + 'tanstack/revert-standing', +) + +/** + * @internal Aborts when the session closes, not on a reload. Work that + * belongs to the session, like background `bash` jobs, stops with it. The + * package does not export it. + */ +export const SessionSignal = createCapability()( + 'tanstack/session-signal', +) + +/** + * @internal Tells the session that a background job, like a `bash` job, + * started or ended. A durable session logs it, so recovery can note a job + * that a crash stopped. The package does not export it. + */ +export const BackgroundJobs = createCapability<{ + started: (jobId: string) => void + ended: (jobId: string) => void +}>()('tanstack/background-jobs') + +/** Capability values provided by plugins, keyed by handle. */ +export class CapabilityValues { + readonly context = { capabilities: new CapabilityRegistry() } + readonly handles: Array = [] + + constructor(private readonly parent?: CapabilityValues) {} + + provide(handle: Capability, value: T): void { + handle[1](this.context, value) + this.handles.push(handle) + } + + get(handle: Capability): T | undefined { + if (handle.has(this.context)) return handle[0](this.context) + return this.parent?.get(handle) + } + + /** Every handle provided here and in the parent chain. */ + all(): Array { + return [...(this.parent?.all() ?? []), ...this.handles] + } +} + +function checkCapabilityOrder( + plugins: ReadonlyArray, + available: Set, +): void { + const byName = new Map() + for (const handle of available) byName.set(handle.capabilityName, handle) + const names = new Set() + for (const plugin of plugins) { + if (names.has(plugin.name)) { + throw new Error(`Duplicate plugin: ${plugin.name}`) + } + names.add(plugin.name) + for (const handle of [ + ...(plugin.requires ?? []), + ...(plugin.provides ?? []), + ...(plugin.optionalRequires ?? []), + ]) { + const other = byName.get(handle.capabilityName) + if (other && other !== handle) { + throw new Error( + `Two different capabilities are named "${handle.capabilityName}" (plugin ${plugin.name}). Capability names must be unique.`, + ) + } + byName.set(handle.capabilityName, handle) + } + for (const handle of plugin.requires ?? []) { + if (!available.has(handle)) { + throw new Error( + `Plugin ${plugin.name} requires capability "${handle.capabilityName}", but no earlier plugin or harness middleware provides it. Add a provider before ${plugin.name}.`, + ) + } + } + for (const handle of plugin.provides ?? []) available.add(handle) + } +} + +function promptOf( + prompt: string | (() => string) | PluginPrompt, + owner: string, + index: number, +): PluginPrompt { + return typeof prompt === 'string' || typeof prompt === 'function' + ? { id: `${owner}#${index}`, text: prompt } + : prompt +} + +/** + * Set up plugins in order and collect what they contribute. + * + * Before any `setup` runs: duplicate plugin names, capability name clashes, + * and missing providers are errors. If a `setup` throws, or a contribution + * clashes (the same tool, prompt id, or agent name from two owners), every + * plugin set up so far is disposed newest first, and no chat turn starts. + */ +export async function mountPlugins( + plugins: ReadonlyArray, + env: MountEnvironment, +): Promise { + const available = new Set([ + ...env.harnessProvides, + ...(env.inherited?.all() ?? []), + ]) + checkCapabilityOrder(plugins, available) + + const values = new CapabilityValues(env.inherited) + const scopes: Array = [] + const tools: Array> = env.harnessTools.map((tool) => ({ + value: tool, + owner: 'the harness', + })) + const prompts: Array> = [] + const middleware: Array = [] + const middlewareOwners: Array = [] + const adapterOwners: Array = [] + const agentMiddleware: Array = [] + const generationMiddleware: Array = [] + const commands = new Map() + const config = new Map() + const extensions = new Map>() + for (const [point, items] of env.inheritedExtensions ?? []) { + extensions.set(point, [...items]) + } + const adapters: MountedPlugins['adapters'] = [] + const discoverers: MountedPlugins['discoverers'] = [] + const preparers: MountedPlugins['preparers'] = [] + const agents: Array = [] + const subagents: Array = [] + const services = env.services ?? NO_SERVICES + const remove = (name: string, list: Array) => { + const at = list.findIndex((item) => item.name === name) + if (at >= 0) list.splice(at, 1) + } + // `ctx.commands.ready`: resolves after the last plugin is set up. A failed + // mount leaves it pending, so no plugin acts on a session that never opens. + let markReady = () => {} + const ready = new Promise((resolve) => (markReady = resolve)) + const ownerOf = (name: string) => + commands.get(name)?.owner ?? env.takenCommands?.get(name)?.owner + + try { + for (const plugin of plugins) { + const scope = new ResourceScope() + scopes.push(scope) + const provided = new Set() + const declared = new Set([ + ...(plugin.requires ?? []), + ...(plugin.optionalRequires ?? []), + ]) + const contributions = await plugin.setup?.({ + resources: { + acquire: (open, close) => scope.acquire(open, close), + signal: scope.signal, + }, + get: (handle) => { + const value = values.get(handle) + if (value === undefined && !declared.has(handle)) { + throw new Error( + `Plugin ${plugin.name} reads capability "${handle.capabilityName}" without declaring it in requires.`, + ) + } + if (value === undefined) { + throw new Error( + `Capability "${handle.capabilityName}" was requested by ${plugin.name} but never provided.`, + ) + } + return value + }, + getOptional: (handle) => values.get(handle), + provide: (handle, value) => { + if (!(plugin.provides ?? []).includes(handle)) { + throw new Error( + `Plugin ${plugin.name} provides "${handle.capabilityName}" without declaring it in provides.`, + ) + } + values.provide(handle, value) + provided.add(handle) + }, + agents: { + list: () => env.registry.list(), + get: (name) => env.registry.get(name), + find: (query) => env.registry.find(query), + run: services.agents.run, + start: services.agents.start, + group: services.agents.group, + set: (agent, options) => { + env.registry.set(agent, plugin.name) + const list = options?.subagent ? subagents : agents + remove(agent.name, options?.subagent ? agents : subagents) + const at = list.findIndex((item) => item.name === agent.name) + list.splice(at < 0 ? list.length : at, 1, agent) + }, + delete: (name) => { + env.registry.delete(name, plugin.name) + // Still there: another owner has it. + if (env.registry.get(name)) return + remove(name, agents) + remove(name, subagents) + }, + }, + collect: (point: ExtensionPoint): ReadonlyArray => { + let items = extensions.get(point.name) + if (!items) { + items = [] + extensions.set(point.name, items) + } + const source = items + const live = () => source.map((item) => item.value as T) + // A live view: items contributed after this call appear too. Array + // methods like `filter` also check that an index exists, so every + // trap reads the items, not the empty target. + return new Proxy>([], { + get: (_target, key) => Reflect.get(live(), key), + has: (_target, key) => Reflect.has(live(), key), + ownKeys: () => Reflect.ownKeys(live()), + getOwnPropertyDescriptor: (_target, key) => + Reflect.getOwnPropertyDescriptor(live(), key), + }) + }, + emit: (event, value) => services.emit(plugin.name, event.name, value), + on: (event, handler) => { + const off = services.on(event.name, (value) => + handler(value as never), + ) + // The listener goes away with the plugin, for example on a reload. + scope.signal.addEventListener('abort', off) + return off + }, + state: (initial) => services.state(plugin.name, initial), + config: services.config, + credentials: services.credentials, + keys: services.keys, + session: plugin.metadataSchema + ? services.sessionWith(plugin.metadataSchema) + : services.session, + commands: { + has: (name) => ownerOf(name) !== undefined, + set: (name, command) => { + const owner = ownerOf(name) + if (owner !== undefined && owner !== plugin.name) { + throw new Error( + `Command "${name}" belongs to ${owner}. ${plugin.name} cannot replace it.`, + ) + } + commands.set(name, { command, owner: plugin.name }) + services.commandsChanged?.() + }, + delete: (name) => { + if (commands.get(name)?.owner !== plugin.name) return + commands.delete(name) + services.commandsChanged?.() + }, + ready, + }, + ...(env.turn ? { turn: env.turn } : {}), + }) + for (const handle of plugin.provides ?? []) { + if (!provided.has(handle)) { + throw new Error( + `Plugin ${plugin.name} declares capability "${handle.capabilityName}" in provides but did not provide it in setup.`, + ) + } + } + if (!contributions) continue + for (const tool of contributions.tools ?? []) { + const clash = tools.find((entry) => entry.value.name === tool.name) + if (clash) { + throw new Error( + `Duplicate tool "${tool.name}": first owner ${clash.owner}, second owner ${plugin.name}. No run started.`, + ) + } + tools.push({ value: tool, owner: plugin.name }) + } + ;(contributions.prompts ?? []).forEach((prompt, index) => { + const section = promptOf(prompt, plugin.name, index) + const clash = prompts.find((entry) => entry.value.id === section.id) + if (clash) { + throw new Error( + `Duplicate prompt id "${section.id}": first owner ${clash.owner}, second owner ${plugin.name}.`, + ) + } + prompts.push({ value: section, owner: plugin.name }) + }) + for (const item of contributions.middleware ?? []) { + middleware.push(item) + middlewareOwners.push(plugin.name) + } + agentMiddleware.push(...(contributions.agentMiddleware ?? [])) + generationMiddleware.push(...(contributions.generationMiddleware ?? [])) + for (const [name, command] of Object.entries( + contributions.commands ?? {}, + )) { + const clash = commands.get(name) ?? env.takenCommands?.get(name) + if (clash) { + throw new Error( + `Duplicate command "${name}": first owner ${clash.owner}, second owner ${plugin.name}.`, + ) + } + commands.set(name, { command, owner: plugin.name }) + } + for (const [key, option] of Object.entries(contributions.config ?? {})) { + const clash = config.get(key) ?? env.takenConfig?.get(key) + if (clash) { + throw new Error( + `Duplicate config key "${key}": first owner ${clash.owner}, second owner ${plugin.name}.`, + ) + } + config.set(key, { option, owner: plugin.name }) + } + for (const item of contributions.contribute ?? []) { + let items = extensions.get(item.point) + if (!items) { + items = [] + extensions.set(item.point, items) + } + items.push({ value: item.value, owner: plugin.name }) + } + if (contributions.adapter) { + adapters.push(contributions.adapter) + 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) { + for (const produces of plugin.needs?.produces ?? []) { + if (!env.registry.find({ produces })) { + throw new Error( + `Plugin ${plugin.name} needs an agent that produces "${produces}", but the session has none. Add one to agents or subagents.`, + ) + } + } + } + } catch (error) { + try { + await disposeAll(scopes) + } catch (cleanupError) { + throw new AggregateError( + [error, cleanupError], + 'Plugin setup failed, and cleanup also failed', + { cause: error }, + ) + } + throw error + } + + const bridgeHandles = values.handles + const capabilityBridge: AnyChatMiddleware | undefined = + bridgeHandles.length > 0 + ? { + name: 'harness:plugin-capabilities', + provides: bridgeHandles, + setup(ctx) { + for (const handle of bridgeHandles) { + handle[1](ctx, values.get(handle)) + } + }, + } + : undefined + + markReady() + let disposed: Promise | undefined + return { + tools: tools + .filter((entry) => entry.owner !== 'the harness') + .map((entry) => entry.value), + prompts: prompts.map((entry) => entry.value.text), + middleware, + agentMiddleware, + generationMiddleware, + commands, + config, + extensions, + adapters, + discoverers, + preparers, + agents, + subagents, + owners: { + plugins: plugins.map((plugin) => ({ + name: plugin.name, + lifetime: plugin.lifetime ?? 'session', + requires: (plugin.requires ?? []).map( + (handle) => handle.capabilityName, + ), + provides: (plugin.provides ?? []).map( + (handle) => handle.capabilityName, + ), + })), + tools: tools.map((entry) => ({ + name: entry.value.name, + owner: entry.owner, + })), + prompts: prompts.map((entry) => ({ + id: entry.value.id, + owner: entry.owner, + })), + middleware: middlewareOwners, + adapters: adapterOwners, + }, + capabilityBridge, + values, + dispose: () => (disposed ??= disposeAll(scopes)), + } +} diff --git a/packages/ai-harness/src/resources.ts b/packages/ai-harness/src/resources.ts new file mode 100644 index 0000000000..3230fc5044 --- /dev/null +++ b/packages/ai-harness/src/resources.ts @@ -0,0 +1,84 @@ +/** + * Resources one plugin owns for one lifetime (a session or a turn). + * + * - `acquire(open, close)` registers `close` only after `open` succeeds. + * - If the scope closes while `open` is still running, the late resource is + * closed as soon as it arrives. + * - `dispose()` closes resources newest first, runs every closer even when + * one throws, and returns the same promise to every caller. + */ +export class ResourceScope { + private readonly controller = new AbortController() + readonly signal: AbortSignal = this.controller.signal + private readonly closers: Array<() => unknown> = [] + private readonly pending = new Set>() + private closing?: Promise + private closed = false + + async acquire( + open: () => T | Promise, + close: (resource: T) => unknown, + ): Promise { + this.signal.throwIfAborted() + let release!: () => void + const pending = new Promise((resolve) => { + release = resolve + }) + this.pending.add(pending) + try { + const resource = await open() + if (this.closed) { + await close(resource) + throw this.signal.reason + } + this.closers.push(() => close(resource)) + return resource + } finally { + this.pending.delete(pending) + release() + } + } + + dispose(): Promise { + if (this.closing) return this.closing + this.closed = true + this.controller.abort(new Error('Resource scope disposed')) + // Defer, so a closer that calls dispose() again gets the same promise. + this.closing = Promise.resolve().then(async () => { + await Promise.all(this.pending) + const errors: Array = [] + for (const close of this.closers.reverse()) { + try { + await close() + } catch (error) { + errors.push(error) + } + } + this.closers.length = 0 + if (errors.length > 0) { + throw new AggregateError(errors, 'Resource cleanup failed') + } + }) + return this.closing + } +} + +/** + * Dispose scopes newest first. Every scope is disposed even when one fails. + * Throws an `AggregateError` with every failure, after all scopes ran. + */ +export async function disposeAll( + scopes: ReadonlyArray, +): Promise { + const errors: Array = [] + for (const scope of [...scopes].reverse()) { + try { + await scope.dispose() + } catch (error) { + errors.push(error) + } + } + if (errors.length > 0) { + throw new AggregateError(errors, 'Plugin cleanup failed') + } +} diff --git a/packages/ai-harness/src/session-entry.ts b/packages/ai-harness/src/session-entry.ts new file mode 100644 index 0000000000..10fa035f8a --- /dev/null +++ b/packages/ai-harness/src/session-entry.ts @@ -0,0 +1,109 @@ +import { isStandardSchema, validateWithStandardSchema } from '@tanstack/ai' +import type { InferSchemaType, SchemaInput } from '@tanstack/ai' +import type { SessionIndexEntry } from '@tanstack/ai-persistence' + +/** + * The type of the session metadata that `TSchema` checks. Without a schema, + * or with a schema that gives no type (a plain JSON Schema), any record. + */ +export type MetadataOf = 0 extends 1 & TSchema + ? Record + : TSchema extends SchemaInput + ? unknown extends InferSchemaType + ? Record + : InferSchemaType + : Record + +/** + * A change to a session index entry. Each field you pass replaces the stored + * one, so `metadata` replaces the whole object. + */ +export interface SessionEntryChange> { + title?: string + pinned?: boolean + metadata?: TMetadata +} + +/** + * A change, or a function that makes the change from the current entry. Use + * the function to change part of `metadata`: when another host writes the + * entry first, the function runs again on the new entry (with a store that + * has `upsertIf`), so neither change is lost. Return `undefined` to write + * nothing. + */ +export type SessionEntryUpdate< + TMetadata = Record, + TChange = SessionEntryChange, +> = TChange | ((entry: SessionIndexEntry) => TChange | undefined) + +/** Session metadata that a `metadataSchema` refused. Nothing was written. */ +export class SessionMetadataError extends Error { + readonly issues: ReadonlyArray<{ message: string; path?: Array }> + constructor( + issues: ReadonlyArray<{ message: string; path?: Array }>, + ) { + super( + `The session metadata is not valid: ${issues.map((issue) => issue.message).join(', ')}`, + ) + this.name = 'SessionMetadataError' + this.issues = issues + } +} + +/** + * @internal Throw when `schema` is not a Standard Schema: a plain JSON Schema + * cannot check a write, so it would let any metadata in. + */ +export function assertMetadataSchema(owner: string, schema: unknown) { + if (schema !== undefined && !isStandardSchema(schema)) { + throw new Error( + `${owner}: metadataSchema must be a Standard Schema (for example a zod object). A plain JSON Schema cannot check the metadata.`, + ) + } +} + +/** The output of `schema` for `metadata`. Throws when it refuses the value. */ +async function checked(schema: SchemaInput, metadata: unknown) { + const result = await validateWithStandardSchema>( + schema, + metadata, + ) + if (!result.success) throw new SessionMetadataError(result.issues) + return result.data +} + +/** + * @internal The entry after `update`. `undefined` when there is no entry or + * the update writes nothing. The harness schema checks `metadata`, and the + * entry keeps its output. `pluginSchema` checks the same value, but its + * output is not kept: a plugin schema with fewer keys would drop the keys of + * the harness. + */ +export async function changedEntry( + entry: SessionIndexEntry | undefined, + update: + | SessionEntryUpdate< + any, + SessionEntryChange & { usage?: SessionIndexEntry['usage'] } + > + | undefined, + harnessSchema: SchemaInput | undefined, + pluginSchema?: SchemaInput, +): Promise { + if (!entry) return undefined + const change = typeof update === 'function' ? update(entry) : update + if (!change) return undefined + let { metadata } = change + if (metadata !== undefined) { + if (pluginSchema) await checked(pluginSchema, metadata) + if (harnessSchema) metadata = await checked(harnessSchema, metadata) + } + return { + ...entry, + ...(change.title !== undefined ? { title: change.title } : {}), + ...(change.pinned !== undefined ? { pinned: change.pinned } : {}), + ...(metadata !== undefined ? { metadata } : {}), + ...(change.usage !== undefined ? { usage: change.usage } : {}), + updatedAt: Date.now(), + } +} diff --git a/packages/ai-harness/src/types.ts b/packages/ai-harness/src/types.ts new file mode 100644 index 0000000000..2b98b13676 --- /dev/null +++ b/packages/ai-harness/src/types.ts @@ -0,0 +1,421 @@ +import type { + AnyTextAdapter, + AnyTool, + ContentPart, + Interrupt, + KeyedAdapter, + ModelMessage, + PromptCacheOptions, + ReasoningOption, + RunAgentResumeItem, + StreamChunk, +} from '@tanstack/ai' + +/** A user message: plain text, or content parts (text, images, files). */ +export type UserInput = string | Array + +/** + * Settings for one chat turn: `session.prompt(message, { overrides })` or + * `session.followUp(message, { overrides })`. They apply to every model call + * of the turn, retries included. The rest of the session keeps its defaults. + * + * They live only in memory, on the queued turn, so they are not part of the + * input: the `inputId` duplicate check does not read them. A turn that + * recovery runs again after a restart uses the defaults. A steer that joins + * a running turn uses the overrides of that turn. + * + * @example + * ```ts + * session.prompt('Plan the migration.', { + * overrides: { reasoning: { level: 'high', summary: true }, tools: [planTool] }, + * }) + * ``` + */ +export interface TurnOverrides { + /** Replaces the harness adapter and every plugin pick for this turn. */ + adapter?: AnyTextAdapter | KeyedAdapter + /** Replaces `HarnessConfig.reasoning` for this turn. */ + reasoning?: ReasoningOption + /** Replaces the session and harness values for this turn, field by field. */ + promptCache?: PromptCacheOptions + /** + * Added to the static tools of this turn. A name that a static tool has + * fails the turn. Durable tools get `step` and `append`. + */ + tools?: ReadonlyArray +} + +/** + * What a `prompt` does when a chat turn is already running: + * - `queue` (default): runs as a new turn after the current one settles. + * - `steer`: joins the running turn at its next model call. + * - `reject`: refuses the prompt. + */ +export type BusyPolicy = 'queue' | 'steer' | 'reject' + +/** Kinds of work a session runs. */ +export type OperationKind = 'chat' | 'agent' | 'command' | 'compact' + +export type OperationStatus = + | 'accepted' + | 'running' + | 'interrupted' + | 'completed' + | 'failed' + | 'cancelled' + +/** The immediate answer to an input. `applied` arrives later as an event. */ +export interface Receipt { + inputId: string + status: 'accepted' | 'queued' | 'rejected' + operationId?: string + reason?: string +} + +/** Opaque position in a session's event stream. */ +export type Cursor = string + +/** One event in a session's ordered stream. */ +export interface SessionEvent { + cursor: Cursor + /** The operation that produced the event. */ + operationId: string + event: StreamChunk +} + +/** + * An input a client sends to a session. Stored in the inbox, or in the log of + * a durable host. + * + * `inputId` is an optional id that the caller chooses. A second input with the + * same id and the same payload gets the first input's receipt and does not run + * again. The same id with another payload is rejected with `'conflict'`. + * + * `context` is JSON data the client sends with a message, for example the + * screen the user is on. It is stored with the input, so a turn that runs + * again after a restart sees the same value. It is client data: do not trust + * it. Who sent the input is the principal from `authorize`, not a field here. + * + * In `cancelInput` and `setDelivery`, `inputId` is the input that waits. + * + * `ephemeral` of `prompt` and `continue` are messages for the model calls of + * that turn only, as in `session.prompt`. The inbox and the log do not keep + * them, and the duplicate check does not read them. Client data: do not + * trust it. + */ +export type HarnessInput = ( + | { + op: 'prompt' + message: UserInput + busy?: BusyPolicy + context?: unknown + ephemeral?: ReadonlyArray + } + | { op: 'steer'; message: UserInput; context?: unknown } + | { + op: 'followUp' + message: UserInput + context?: unknown + /** + * The harness wrote this input to wake the session, for example when a + * background agent ended. Its user message gets + * `metadata.harness.wake: true`. Only the harness sets it: the input + * of a client never has it. + */ + wake?: true + } + /** A turn from the stored transcript, with no new message. */ + | { + op: 'continue' + context?: unknown + ephemeral?: ReadonlyArray + } + | { op: 'resolve'; resume: Array } + | { + op: 'agent' + agent: string + input?: unknown + detached?: boolean + /** Run again after a host stop. See `AgentStartOptions.resume`. */ + resume?: boolean + /** `'none'`: no note for the main model. See `AgentRunOptions.attach`. */ + attach?: 'none' + /** The run whose agent code started this run. */ + parentRunId?: string + } + | { op: 'cancel'; operationId?: string } + /** Move a running tool call to the background. See `session.background`. */ + | { op: 'background'; toolCallId?: string } + | { op: 'cancelInput'; inputId: string } + | { + op: 'setDelivery' + inputId: string + delivery: WaitingInput['delivery'] + } + | { op: 'command'; name: string; input?: unknown } + | { op: 'answer'; questionId: string; value: unknown } + | { op: 'config'; key: string; value: unknown } + | { op: 'configure'; settings: ThreadSettingsChange } + | { op: 'reset'; note?: string } + | { op: 'revert'; messageId: string } + | { op: 'unrevert' } + | { + op: 'agentMessage' + /** A run of the agent: any run of its chain. */ + operationId: string + message: UserInput + /** Default `'steer'`. See `AgentRun.send`. */ + mode?: 'steer' | 'followUp' + /** The input id of the chain's first run. The session sets it. */ + run?: string + } +) & { inputId?: string } + +/** + * An input that waits to run. `session.inputs()` lists them. Change one with + * the `cancelInput` and `setDelivery` inputs. + */ +export interface WaitingInput { + inputId: string + /** + * - `steer`: it joins the running turn at the next model call. When the + * turn makes no more model calls, it runs as the next turn. + * - `queue`: it runs as its own turn, after the turns before it. + */ + delivery: 'steer' | 'queue' + /** The user message of the input. */ + message: UserInput +} + +/** + * Where a fork cuts the transcript, by message id. `before` keeps the + * messages before that message. `through` keeps that message too. + */ +export type ForkPoint = { before: string } | { through: string } + +/** + * The stored settings of one thread. They apply from the next turn, stay + * after a restart (with `stores.metadata`), and a fork copies them. A turn's + * `overrides` win over them. + */ +export interface ThreadSettings { + /** A name from `defineHarness({ models })`. */ + model?: string + /** The `reasoning` of every turn, as `chat({ reasoning })` takes it. */ + reasoning?: ReasoningOption + /** Text added as the last system prompt of every turn. */ + instructions?: string + /** Exactly these tools (by name), or every tool but the `remove` ones. */ + tools?: Array | { remove: Array } + /** Plugins whose tools, prompts, and middleware the turns leave out. */ + plugins?: { remove: Array } + /** + * The working folder of the thread, from the root of the workspace tools. + * A folder outside that root is refused when a tool uses it. + */ + cwd?: string +} + +/** A change to {@link ThreadSettings}: `null` clears a field, a missing field stays. */ +export type ThreadSettingsChange = { + [Key in keyof ThreadSettings]?: ThreadSettings[Key] | null +} + +/** How an input ended. `session.settled(inputId)` resolves to it. */ +export interface InputSettlement { + inputId: string + /** + * `interrupted`: the turn stopped for human input. The `resolve` that + * answers it is a new input with its own settlement. + */ + outcome: 'completed' | 'failed' | 'aborted' | 'interrupted' + /** The operation that ran the input, when one started. */ + operationId?: string + /** Set for `failed`. `code` is `'attempts_exhausted'`, `'timeout'`, or absent. */ + error?: { message: string; code?: string } +} + +/** + * An awaited operation rejects with this error when the session refused its + * input (for example `busy: 'reject'` while a turn runs, or an `inputId` + * conflict). `receipt` is the refused receipt. + */ +export class InputRejectedError extends Error { + readonly receipt: Receipt + + constructor(receipt: Receipt) { + super(rejectionMessage(receipt.reason)) + this.name = 'InputRejectedError' + this.receipt = receipt + } +} + +function rejectionMessage(reason: string | undefined) { + if (reason === 'busy') return 'A chat turn is already running.' + if (reason === 'conflict') { + return 'An earlier input has this inputId and another payload.' + } + return `The session rejected the input: ${reason ?? 'rejected'}.` +} + +/** Who sent an input, from the host's `authorize`. */ +export interface Principal { + id: string + name?: string + /** + * The organization of the principal. The credential scope gets it, so two + * organizations' credentials for one user stay apart. + */ + tenantId?: string +} + +/** + * A chat turn as plugins and the router see it: the input that started it, + * and who sent it. + */ +export interface TurnInfo { + /** The id of the turn operation. */ + operationId: string + /** The id of the turn input, when the turn has one. */ + inputId?: string + /** The user message of the input. `undefined` for a resolve and a continue. */ + message?: UserInput + /** The `context` of the input. Client data: do not trust it. */ + context?: unknown + /** Who sent the input. Default: the principal that opened the session. */ + principal?: Principal + /** The settings of this turn only. Not kept after a restart. */ + overrides?: TurnOverrides +} + +/** What a chat turn operation resolves to. */ +export interface ChatTurnResult { + /** + * The main model's text for this turn, without the text of agents it + * called. When a router sends the turn to agents, it is their text, as the + * transcript keeps it. + */ + text: string + /** Set when the turn stopped for outside input. */ + interrupts?: Array +} + +/** + * One accepted unit of work in a session. Await it for its result. + */ +export interface Operation extends PromiseLike { + /** Also the AG-UI `runId` of the operation. */ + readonly id: string + readonly kind: OperationKind + /** The agent name, for `kind: 'agent'`. */ + readonly agent?: string + status: () => OperationStatus + /** + * The answer to the input of this operation. It resolves after the session + * stored the input (`accepted` or `queued`), or refused it (`rejected`). + * With a durable host, a resolved receipt means that the input survives a + * restart. + */ + readonly receipt: Promise + /** This operation's events, from `from` (exclusive) onward. */ + events: (options?: { + from?: Cursor + signal?: AbortSignal + }) => AsyncIterable + /** This operation's raw AG-UI chunks, for existing transports. */ + stream: (options?: { signal?: AbortSignal }) => AsyncIterable + cancel: (reason?: string) => Promise +} + +/** A background agent run. `send` adds a message to it. */ +export interface AgentRun extends Operation { + /** + * Add a message to this run. `mode: 'steer'` (default) gives it to the + * run's next model call. When the run makes no further model call, it + * runs as a follow-up. `mode: 'followUp'` runs the agent again after the + * run ends, on the run's transcript with the message added, with its + * first input and options. A run that ended starts again at once. The + * receipt names the run that gets the message. + */ + send: ( + message: UserInput, + options?: { mode?: 'steer' | 'followUp'; inputId?: string }, + ) => Promise +} + +/** What a media file holds. It picks the content part a model gets. */ +export type MediaKind = 'image' | 'audio' | 'video' | 'document' + +/** A file in a thread's media store: sent by a user or made by an agent. */ +export interface MediaRecord { + id: string + threadId: string + kind: MediaKind + mimeType: string + name: string + /** Size in bytes. */ + size: number + source: 'user' | 'generated' + /** Epoch milliseconds. */ + createdAt: number + /** The operation that made generated media. */ + runId?: string + /** The child agent run that made it, when a subagent did. */ + subagentRunId?: string +} + +/** Names of the `CUSTOM` events a harness session adds to the stream. */ +export const HARNESS_EVENTS = { + operationStarted: 'harness.operation.started', + operationFinished: 'harness.operation.finished', + operationResumed: 'harness.operation.resumed', + configChanged: 'harness.config.changed', + /** `session.configure` changed the settings. The value has `settings`. */ + settingsChanged: 'harness.settings.changed', + /** A plugin added or removed a command with `ctx.commands`. */ + commandsChanged: 'harness.commands.changed', + /** + * `session.reload()` set the plugins up again. The value is empty, or has + * the `error` message when the new setup failed. Read the commands, the + * config, and the agents again. + */ + reloaded: 'harness.reloaded', + question: 'harness.question', + questionAnswered: 'harness.question.answered', + pluginEvent: 'harness.plugin.event', + authRequired: 'harness.auth_required', + inputAccepted: 'harness.input.accepted', + inputApplied: 'harness.input.applied', + inputRejected: 'harness.input.rejected', + /** An input ended. The value is an `InputSettlement`. */ + inputSettled: 'harness.input.settled', + /** + * A `setDelivery` input moved a waiting input. The value has `inputId` and + * `delivery`. + */ + inputDelivery: 'harness.input.delivery', + /** A media file was stored. The value is a `MediaRecord`. */ + media: 'harness.media', + /** + * A turn runs the model again after an error. The value has `operationId`, + * `retries`, `error`, and `continued`. `continued` is true when the model + * continues a partial answer. + */ + turnRetry: 'harness.turn.retry', + /** + * `session.reset()` started a fresh model context. The value has + * `inputId` and the `note`, when there is one. + */ + reset: 'harness.reset', + /** + * `session.revert` or `session.unrevert` changed what the transcript + * shows. The value has the `messageId` of the revert, or `null` when no + * revert stands. Read the transcript again. + */ + revert: 'harness.revert', + /** + * A model call reported its usage. The value has `model`, `sender` (when + * known), the `usage` of the call, and the new thread `total`. + */ + usage: 'harness.usage', +} as const diff --git a/packages/ai-harness/src/usage.ts b/packages/ai-harness/src/usage.ts new file mode 100644 index 0000000000..2066d4133f --- /dev/null +++ b/packages/ai-harness/src/usage.ts @@ -0,0 +1,96 @@ +import type { TokenUsage } from '@tanstack/ai' + +/** The token counts of some model calls, and their cost when known. */ +export interface UsageCounts { + /** How many model calls reported usage. */ + calls: number + promptTokens: number + completionTokens: number + totalTokens: number + /** Prompt tokens read from the cache of the provider. */ + cachedTokens: number + /** Prompt tokens written to the cache of the provider. */ + cacheWriteTokens: number + /** + * The sum of the cost that the providers reported, in their currency (USD + * for most). Set only when at least one call reported a cost. + */ + cost?: number +} + +/** + * The usage of a thread: every model call of its turns, their subagents, and + * its agent runs. `bySender` has the calls of a known sender only, keyed by + * the principal id, so its sum can be less than `total`. + */ +export interface SessionUsage { + total: UsageCounts + /** By `provider/model`, for example `anthropic/claude-sonnet-5-5`. */ + byModel: Record + /** By principal id. */ + bySender: Record +} + +/** One model call, as the log keeps it. A type, so it is a log record field set. */ +export type UsageCall = { + model: string + principal?: { id: string; tenantId?: string } + usage: Omit +} + +const emptyCounts = (): UsageCounts => ({ + calls: 0, + promptTokens: 0, + completionTokens: 0, + totalTokens: 0, + cachedTokens: 0, + cacheWriteTokens: 0, +}) + +export const emptyUsage = (): SessionUsage => ({ + total: emptyCounts(), + byModel: {}, + bySender: {}, +}) + +/** The counts that one call adds, from what the provider reported. */ +export function callUsage(usage: TokenUsage): UsageCall['usage'] { + return { + promptTokens: usage.promptTokens, + completionTokens: usage.completionTokens, + totalTokens: usage.totalTokens, + cachedTokens: usage.promptTokensDetails?.cachedTokens ?? 0, + cacheWriteTokens: usage.promptTokensDetails?.cacheWriteTokens ?? 0, + ...(typeof usage.cost === 'number' ? { cost: usage.cost } : {}), + } +} + +function add(counts: UsageCounts, usage: UsageCall['usage']) { + counts.calls += 1 + counts.promptTokens += usage.promptTokens + counts.completionTokens += usage.completionTokens + counts.totalTokens += usage.totalTokens + counts.cachedTokens += usage.cachedTokens + counts.cacheWriteTokens += usage.cacheWriteTokens + if (usage.cost !== undefined) counts.cost = (counts.cost ?? 0) + usage.cost +} + +/** Add one call to `totals`, in place. */ +export function addUsage(totals: SessionUsage, call: UsageCall) { + add(totals.total, call.usage) + add((totals.byModel[call.model] ??= emptyCounts()), call.usage) + if (call.principal) { + add((totals.bySender[call.principal.id] ??= emptyCounts()), call.usage) + } +} + +/** Is `value` usage totals that the harness wrote? */ +export function isSessionUsage(value: unknown): value is SessionUsage { + return ( + typeof value === 'object' && + value !== null && + 'total' in value && + 'byModel' in value && + 'bySender' in value + ) +} diff --git a/packages/ai-harness/src/utils.ts b/packages/ai-harness/src/utils.ts new file mode 100644 index 0000000000..533bad5e0d --- /dev/null +++ b/packages/ai-harness/src/utils.ts @@ -0,0 +1,3 @@ +export function isRecord(value: unknown): value is Record { + return typeof value === 'object' && value !== null +} diff --git a/packages/ai-harness/tsconfig.json b/packages/ai-harness/tsconfig.json new file mode 100644 index 0000000000..29112eff9f --- /dev/null +++ b/packages/ai-harness/tsconfig.json @@ -0,0 +1,8 @@ +{ + "extends": "../../tsconfig.base.json", + "compilerOptions": { + "outDir": "dist" + }, + "include": ["vite.config.ts", "./src", "./tests"], + "exclude": ["node_modules", "dist"] +} diff --git a/packages/ai-harness/vite.config.ts b/packages/ai-harness/vite.config.ts new file mode 100644 index 0000000000..1f3542380f --- /dev/null +++ b/packages/ai-harness/vite.config.ts @@ -0,0 +1,35 @@ +import { defineConfig, mergeConfig } from 'vitest/config' +import { tanstackViteConfig } from '@tanstack/vite-config' +import packageJson from './package.json' + +const config = defineConfig({ + test: { + name: packageJson.name, + dir: './', + watch: false, + globals: true, + environment: 'node', + include: ['src/**/*.test.ts', 'tests/**/*.test.ts'], + coverage: { + provider: 'v8', + reporter: ['text', 'json', 'html', 'lcov'], + exclude: [ + 'node_modules/', + 'dist/', + 'tests/', + '**/*.test.ts', + '**/*.config.ts', + ], + include: ['src/**/*.ts'], + }, + }, +}) + +export default mergeConfig( + config, + tanstackViteConfig({ + entry: ['./src/index.ts'], + srcDir: './src', + cjs: false, + }), +) diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index aa9e45eb08..8ca632dd08 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -2584,6 +2584,21 @@ importers: specifier: ^8.2.1 version: 8.2.1(@types/node@24.10.3)(esbuild@0.28.1)(jiti@2.7.0)(less@4.6.6)(sass@1.101.0)(terser@5.44.1)(tsx@4.23.12)(yaml@2.9.0) + packages/ai-harness: + devDependencies: + '@tanstack/ai': + specifier: workspace:* + version: link:../ai + '@tanstack/ai-persistence': + specifier: workspace:* + version: link:../ai-persistence + '@vitest/coverage-v8': + specifier: 4.1.10 + version: 4.1.10(vitest@4.1.11) + zod: + specifier: ^4.2.0 + version: 4.6.5 + packages/ai-isolate-cloudflare: dependencies: '@tanstack/ai-code-mode': @@ -3450,7 +3465,7 @@ importers: devDependencies: '@google/genai': specifier: ^2.10.0 - version: 2.10.0(@modelcontextprotocol/sdk@1.29.0(zod@4.3.6)) + version: 2.10.0(@modelcontextprotocol/sdk@1.29.0(zod@4.6.5)) '@tanstack/ai': specifier: workspace:* version: link:../ai @@ -3690,6 +3705,9 @@ importers: '@tanstack/ai-groq': specifier: workspace:* version: link:../../packages/ai-groq + '@tanstack/ai-harness': + specifier: workspace:* + version: link:../../packages/ai-harness '@tanstack/ai-llmgateway': specifier: workspace:* version: link:../../packages/ai-llmgateway @@ -20991,7 +21009,7 @@ snapshots: '@blaxel/core@0.3.11(@hey-api/openapi-ts@0.99.0(magicast@0.5.2)(typescript@7.0.2))': dependencies: '@hey-api/client-fetch': 0.10.2(@hey-api/openapi-ts@0.99.0(magicast@0.5.2)(typescript@7.0.2)) - '@modelcontextprotocol/sdk': 1.29.0(zod@4.3.6) + '@modelcontextprotocol/sdk': 1.29.0(zod@4.6.5) archiver: 7.0.1 axios: 1.18.1(debug@4.4.3(supports-color@7.2.0))(supports-color@7.2.0) dockerfile-ast: 0.7.1 @@ -21002,7 +21020,7 @@ snapshots: uuid: 11.1.1 ws: 8.21.0 yaml: 2.9.0 - zod: 4.3.6 + zod: 4.6.5 transitivePeerDependencies: - '@cfworker/json-schema' - '@hey-api/openapi-ts' @@ -22718,8 +22736,8 @@ snapshots: '@mistralai/mistralai@2.2.0': dependencies: ws: 8.21.0 - zod: 4.3.6 - zod-to-json-schema: 3.25.2(zod@4.3.6) + zod: 4.6.5 + zod-to-json-schema: 3.25.2(zod@4.6.5) transitivePeerDependencies: - bufferutil - utf-8-validate @@ -22732,11 +22750,11 @@ snapshots: eventsource-parser: 3.1.1 jose: 6.2.12 pkce-challenge: 5.0.1 - zod: 4.3.6 + zod: 4.6.5 '@modelcontextprotocol/core@2.1.0': dependencies: - zod: 4.3.6 + zod: 4.6.5 '@modelcontextprotocol/ext-apps@1.7.4(@modelcontextprotocol/sdk@1.29.0(supports-color@7.2.0)(zod@4.2.1))(react-dom@19.2.3(react@19.2.3))(react@19.2.3)(zod@3.25.76)': dependencies: @@ -22834,12 +22852,11 @@ snapshots: zod-to-json-schema: 3.25.2(zod@4.6.5) transitivePeerDependencies: - supports-color - optional: true '@modelcontextprotocol/server@2.1.0': dependencies: '@modelcontextprotocol/core': 2.1.0 - zod: 4.3.6 + zod: 4.6.5 '@msgpack/msgpack@3.1.3': {} @@ -23099,7 +23116,7 @@ snapshots: '@openrouter/sdk@0.13.20': dependencies: - zod: 4.3.6 + zod: 4.6.5 '@opentelemetry/api-logs@0.217.0': dependencies: @@ -28608,7 +28625,7 @@ snapshots: tar-stream: 3.1.7 undici: 7.28.0 xdg-app-paths: 5.1.0 - zod: 4.3.6 + zod: 4.6.5 transitivePeerDependencies: - bare-abort-controller - react-native-b4a @@ -38359,7 +38376,6 @@ snapshots: zod-to-json-schema@3.25.2(zod@4.6.5): dependencies: zod: 4.6.5 - optional: true zod@3.25.76: {} From 693f440a91ee19bcd2f16293fa1499b16258d528 Mon Sep 17 00:00:00 2001 From: Alem Tuzlak Date: Tue, 13 Oct 2026 22:40:28 +0200 Subject: [PATCH 005/131] feat(ai-harness): session engine and host, with the session split into modules --- packages/ai-harness/src/harness-agent.ts | 36 + packages/ai-harness/src/harness-text.ts | 290 +++++ packages/ai-harness/src/host.ts | 931 +++++++++++++++ packages/ai-harness/src/media.ts | 472 ++++++++ packages/ai-harness/src/session/agents.ts | 1029 ++++++++++++++++ packages/ai-harness/src/session/base.ts | 765 ++++++++++++ .../ai-harness/src/session/harness-session.ts | 810 +++++++++++++ packages/ai-harness/src/session/helpers.ts | 667 +++++++++++ packages/ai-harness/src/session/index.ts | 14 + packages/ai-harness/src/session/inputs.ts | 551 +++++++++ packages/ai-harness/src/session/requests.ts | 933 +++++++++++++++ packages/ai-harness/src/session/state.ts | 966 +++++++++++++++ packages/ai-harness/src/session/turn-parts.ts | 723 ++++++++++++ packages/ai-harness/src/session/turn.ts | 1043 +++++++++++++++++ packages/ai-harness/src/session/types.ts | 365 ++++++ packages/ai-harness/src/turn.ts | 264 +++++ 16 files changed, 9859 insertions(+) create mode 100644 packages/ai-harness/src/harness-agent.ts create mode 100644 packages/ai-harness/src/harness-text.ts create mode 100644 packages/ai-harness/src/host.ts create mode 100644 packages/ai-harness/src/media.ts create mode 100644 packages/ai-harness/src/session/agents.ts create mode 100644 packages/ai-harness/src/session/base.ts create mode 100644 packages/ai-harness/src/session/harness-session.ts create mode 100644 packages/ai-harness/src/session/helpers.ts create mode 100644 packages/ai-harness/src/session/index.ts create mode 100644 packages/ai-harness/src/session/inputs.ts create mode 100644 packages/ai-harness/src/session/requests.ts create mode 100644 packages/ai-harness/src/session/state.ts create mode 100644 packages/ai-harness/src/session/turn-parts.ts create mode 100644 packages/ai-harness/src/session/turn.ts create mode 100644 packages/ai-harness/src/session/types.ts create mode 100644 packages/ai-harness/src/turn.ts diff --git a/packages/ai-harness/src/harness-agent.ts b/packages/ai-harness/src/harness-agent.ts new file mode 100644 index 0000000000..fbb7dddc58 --- /dev/null +++ b/packages/ai-harness/src/harness-agent.ts @@ -0,0 +1,36 @@ +import { defineAgent } from '@tanstack/ai' +import { harnessText } from './harness-text' +import type { AnyHarness } from './define' +import type { HarnessHost } from './host' + +/** + * Use a whole harness as a child agent: put it in `subagents.agents` and the + * main model can call it as a tool, with its own tools, plugins, and history. + * + * @example + * ```ts + * defineHarness({ + * name: 'acme/lead', + * adapter, + * subagents: { agents: [harnessAgent(reviewer)] }, + * }) + * ``` + */ +export function harnessAgent( + harness: AnyHarness, + options: { host?: HarnessHost; name?: string; description?: string } = {}, +) { + const adapter = harnessText( + harness, + options.host ? { host: options.host } : {}, + ) + return defineAgent({ + // Tool names allow letters, digits, `_`, and `-`. + name: options.name ?? harness.name.replace(/[^a-zA-Z0-9_-]/g, '_'), + description: + options.description ?? + harness.description ?? + `The ${harness.name} harness`, + run: (ctx) => ctx.chat({ adapter }), + }) +} diff --git a/packages/ai-harness/src/harness-text.ts b/packages/ai-harness/src/harness-text.ts new file mode 100644 index 0000000000..09277f100f --- /dev/null +++ b/packages/ai-harness/src/harness-text.ts @@ -0,0 +1,290 @@ +import { EventType, uiMessagesToWire } from '@tanstack/ai' +import { createHarnessHost } from './host' +import { acceptedKinds } from './session' +import type { + AnyTextAdapter, + Modality, + ModelMessage, + StreamChunk, +} from '@tanstack/ai' +import { isHarnessDefinition } from './define' +import type { AnyHarness } from './define' +import type { HarnessHost } from './host' +import type { UserInput } from './types' + +export interface HarnessTextOptions { + /** The host that runs the inner sessions. Default: a memory host. */ + host?: HarnessHost + /** + * The kinds the inner model reads. Default: the list of the harness + * `adapter`. Set it for a harness whose plugin picks the model, or that + * has a `keyedAdapter`: neither list is known until a turn runs. + * `media.accepts` narrows it. + */ + inputModalities?: ReadonlyArray +} + +const isRecord = (value: unknown): value is Record => + typeof value === 'object' && value !== null + +/** The content of the last user message: its text, or its content parts. */ +function lastUserInput(messages: ReadonlyArray) { + const message = messages.findLast((entry) => entry.role === 'user') + const input: UserInput = message?.content ?? '' + return input +} + +/** + * What a harness reads: text, plus the kinds its session sends the model + * (`modalities`, else the adapter's list, narrowed by `media.accepts`). + * `undefined` when neither is known, so the outer harness sends every kind. + */ +function inputsOf( + harness: AnyHarness, + modalities: ReadonlyArray | undefined, +) { + // ponytail: the harness adapter only. A plugin that picks another model at + // runtime is not seen here, so the caller says it with `inputModalities`. + // A keyedAdapter has no model until a turn builds it, so it has no list. + const kinds = acceptedKinds( + modalities ?? harness.adapter?.inputModalities, + harness.media?.accepts, + ) + if (kinds === undefined) return undefined + const inputs: ReadonlyArray = [ + 'text', + ...kinds.filter((kind) => kind !== 'text'), + ] + return inputs +} + +/** Events of the inner turn that the outer chat shows as its own model output. */ +const FORWARDED = new Set([ + EventType.TEXT_MESSAGE_START, + EventType.TEXT_MESSAGE_CONTENT, + EventType.TEXT_MESSAGE_END, + EventType.REASONING_START, + EventType.REASONING_MESSAGE_START, + EventType.REASONING_MESSAGE_CONTENT, + EventType.REASONING_MESSAGE_END, + EventType.REASONING_END, +]) + +/** A harness served by `createHarnessHandler` (or `runCli --serve`) on another machine. */ +export interface RemoteHarness { + /** The handler base URL, for example `http://127.0.0.1:8787`. */ + url: string + /** The bearer token. */ + token?: string + fetch?: typeof fetch +} + +// Type-only fields. Every input kind type-checks, so `chat()` lets callers +// send files. The runtime `inputModalities` says what the harness reads. +const TYPES = { + providerOptions: {} as Record, + inputModalities: ['text', 'image', 'audio', 'video', 'document'] as const, + messageMetadataByModality: { + text: undefined as unknown, + image: undefined as unknown, + audio: undefined as unknown, + video: undefined as unknown, + document: undefined as unknown, + }, + toolCapabilities: [] as ReadonlyArray, + toolCallMetadata: undefined as unknown, + systemPromptMetadata: undefined as never, +} + +/** Stream the SSE `data:` payloads of a response. */ +async function* sseData(response: Response): AsyncGenerator { + if (!response.body) return + const reader = response.body.getReader() + const decoder = new TextDecoder() + let buffer = '' + while (true) { + const { value, done } = await reader.read() + if (done) break + buffer += decoder.decode(value, { stream: true }) + const blocks = buffer.split('\n\n') + buffer = blocks.pop() ?? '' + for (const block of blocks) { + const data = block + .split('\n') + .find((line) => line.startsWith('data: ')) + ?.slice(6) + if (data) yield JSON.parse(data) + } + } +} + +function remoteHarnessText(remote: RemoteHarness): AnyTextAdapter { + const base = remote.url.replace(/\/$/, '') + const doFetch = remote.fetch ?? fetch + // No runtime `inputModalities`: the model of the remote harness is not + // known here, so the outer harness sends every kind. + return { + kind: 'text', + name: 'harness-remote', + model: base, + '~types': TYPES, + structuredOutput: () => + Promise.reject( + new Error('harnessText does not support structured output.'), + ), + chatStream: (chatOptions) => + (async function* (): AsyncGenerator { + const threadId = chatOptions.threadId ?? 'default' + const runId = chatOptions.runId ?? `harness-${Date.now().toString(36)}` + // AG-UI parts: text parts carry `text`, media parts pass through. + const [message] = uiMessagesToWire([ + { role: 'user', content: lastUserInput(chatOptions.messages) }, + ]) + const response = await doFetch(`${base}/run`, { + method: 'POST', + headers: { + 'Content-Type': 'application/json', + ...(remote.token + ? { Authorization: `Bearer ${remote.token}` } + : {}), + }, + body: JSON.stringify({ + threadId, + runId, + // The remote session keeps its own history, so send the new message only. + messages: [{ ...message, id: `${runId}-user` }], + tools: [], + context: [], + state: {}, + forwardedProps: {}, + }), + }) + if (!response.ok) { + throw new Error( + `Remote harness failed (${response.status}): ${await response.text()}`, + ) + } + yield { + type: EventType.RUN_STARTED, + runId, + threadId, + timestamp: Date.now(), + } + for await (const data of sseData(response)) { + if (!isRecord(data) || typeof data.type !== 'string') continue + // The handler sends AG-UI chunks over SSE. + const chunk = data as StreamChunk + if ( + FORWARDED.has(chunk.type) && + !('subagentRunId' in chunk && chunk.subagentRunId) + ) { + yield chunk + } + if (chunk.type === EventType.RUN_ERROR) { + yield { + type: EventType.RUN_ERROR, + message: chunk.message, + timestamp: Date.now(), + } + return + } + } + yield { + type: EventType.RUN_FINISHED, + runId, + threadId, + timestamp: Date.now(), + metadata: { tanstack: { finishReason: 'stop' } }, + } + })(), + } +} + +/** + * Use a harness as the model of a `chat()` call, the way you would call a + * coding agent. Each outer thread gets its own inner session, which keeps + * its own transcript, tools, plugins, and agents. The inner turn gets the + * last user message with its content parts (images, audio, video, + * documents). The outer chat sees the inner turn's text and reasoning. + * + * `inputModalities` is what the harness reads: the `inputModalities` + * option, else its adapter's list, narrowed by `media.accepts`. It is + * `undefined` for a remote harness. + * + * @example + * ```ts + * const stream = chat({ adapter: harnessText(studio), messages, threadId }) + * ``` + */ +export function harnessText( + harness: AnyHarness, + options?: HarnessTextOptions, +): AnyTextAdapter +export function harnessText(remote: RemoteHarness): AnyTextAdapter +export function harnessText( + target: AnyHarness | RemoteHarness, + options: HarnessTextOptions = {}, +): AnyTextAdapter { + if (!isHarnessDefinition(target)) return remoteHarnessText(target) + const harness = target + let host = options.host + return { + kind: 'text', + name: 'harness', + model: harness.name, + inputModalities: inputsOf(harness, options.inputModalities), + '~types': TYPES, + chatStream: (chatOptions) => + (async function* (): AsyncGenerator { + host ??= createHarnessHost() + const threadId = `${chatOptions.threadId ?? 'default'}:${harness.name}` + const runId = chatOptions.runId ?? `harness-${Date.now().toString(36)}` + const signal = + isRecord(chatOptions.request) && + chatOptions.request.signal instanceof AbortSignal + ? chatOptions.request.signal + : undefined + const session = await host.open(harness, { threadId }) + const operation = session.prompt(lastUserInput(chatOptions.messages)) + signal?.addEventListener('abort', () => void operation.cancel(), { + once: true, + }) + + yield { + type: EventType.RUN_STARTED, + runId, + threadId, + timestamp: Date.now(), + } + let failed: string | undefined + for await (const chunk of operation.stream()) { + if ( + FORWARDED.has(chunk.type) && + !('subagentRunId' in chunk && chunk.subagentRunId) + ) { + yield chunk + } + if (chunk.type === EventType.RUN_ERROR) failed = chunk.message + } + if (failed !== undefined) { + yield { + type: EventType.RUN_ERROR, + message: failed, + timestamp: Date.now(), + } + return + } + yield { + type: EventType.RUN_FINISHED, + runId, + threadId, + timestamp: Date.now(), + metadata: { tanstack: { finishReason: 'stop' } }, + } + })(), + structuredOutput: () => + Promise.reject( + new Error('harnessText does not support structured output.'), + ), + } +} diff --git a/packages/ai-harness/src/host.ts b/packages/ai-harness/src/host.ts new file mode 100644 index 0000000000..a9a73bf3a3 --- /dev/null +++ b/packages/ai-harness/src/host.ts @@ -0,0 +1,931 @@ +import { EventType } from '@tanstack/ai' +import { memoryPersistence } from '@tanstack/ai-persistence' +import { SharedLog, loadLogState } from './log' +import { LEASE } from './resume' +import { HarnessSession } from './session' +import { HARNESS_EVENTS } from './types' +import { isRecord } from './utils' +import { changedEntry } from './session-entry' +import type { PromptCacheOptions, RunStore, StreamChunk } from '@tanstack/ai' +import type { + AIPersistence, + ArtifactStore, + BlobStore, + ChatTranscriptStores, + CredentialStore, + GenerationRunStore, + InboxStore, + InterruptStore, + LeaseStore, + LogStore, + MetadataStore, + WorkClaimStore, + SessionIndexEntry, + SessionIndexListOptions, + SessionIndexPage, + SessionIndexStore, +} from '@tanstack/ai-persistence' +import type { AnyHarness, SessionMetadataOf } from './define' +import type { SessionEntryChange, SessionEntryUpdate } from './session-entry' +import type { ProjectOptions, ReduceOptions } from './log' +import type { LeaseOptions } from './resume' +import type { ForkPoint, Principal } from './types' + +type SharedStores = { + credentials?: CredentialStore + artifacts?: ArtifactStore + blobs?: BlobStore + generationRuns?: GenerationRunStore + /** Claims on busy threads, for `resumePending`. */ + workClaims?: WorkClaimStore + /** The session index that `host.sessions` reads. */ + sessions?: SessionIndexStore +} + +/** A durable host needs run leases, host leases, or both. */ +type DurableLeases = + | { runs: RunStore; leases?: LeaseStore } + | { runs?: RunStore; leases: LeaseStore } + +/** + * The stores a host needs. Two shapes: + * + * - A message store, plus any of runs, interrupts, metadata, and inbox. + * Without an inbox, inputs are kept in memory and a restart loses the ones + * not yet applied. + * - Durable mode: a `log`, and `runs` or `leases`. The log holds the events, + * transcript, inputs, and tool steps of each thread, so this shape has no + * `messages` and no `inbox`. + * + * Without `artifacts`, `blobs`, and `generationRuns`, media is kept in memory. + * With `sessions`, the host keeps one index entry per thread (see + * `host.sessions`). + */ +export type HarnessPersistence = + | AIPersistence< + ChatTranscriptStores & { + inbox?: InboxStore + log?: never + leases?: never + } & SharedStores + > + | AIPersistence< + { + log: LogStore + interrupts?: InterruptStore + metadata?: MetadataStore + messages?: never + inbox?: never + } & DurableLeases & + SharedStores + > + +export interface HarnessHostOptions { + /** Where sessions keep their state. Default: in memory, lost on restart. */ + persistence?: HarnessPersistence + /** + * How a durable host folds your own log records into the model context. + * Only a host with `stores.log` reads it. + */ + project?: ProjectOptions + /** + * A fold of every host record of a log, across all its sessions, in log + * order. Read it with `host.logState(logId)`. Only a host with + * `stores.log` reads it. + */ + reduce?: ReduceOptions + /** + * How long a durable host merges streamed text before it appends it to the + * log, in milliseconds. Boundaries (a message start or end, a tool call, a + * run event) append at once. Default 100. + */ + coalesceMs?: number + /** How long a run lease lasts, and how often the host renews it. */ + lease?: LeaseOptions +} + +export interface OpenSessionOptions { + /** The conversation id. Opening the same id again returns the live session. */ + threadId: string + /** Who opened the session. Stored on inputs and run records. */ + principal?: Principal + /** + * The log of the session, on a durable host. Sessions with the same log id + * share one log and one writer: one append can hold records for several of + * them. Default: `threadId`. + */ + logId?: string + /** + * Prompt caching for this session. It overrides the harness `promptCache`, + * for example with a stable affinity key. Opening the same `threadId` again + * returns the live session, which keeps the value of the first open. + */ + promptCache?: PromptCacheOptions +} + +/** + * The session index of a host: one entry per 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. + */ +export interface HostSessions { + /** + * Entries, newest `updatedAt` first. Default: only top-level sessions + * (`parentThreadId: null`). Pass a thread id to get its child sessions. + * Pass `principal` to get only the sessions of that user. + */ + list: (options?: SessionIndexListOptions) => Promise + /** The entry of a thread, or `undefined`. */ + get: (threadId: string) => Promise + /** + * Set the title of a thread. Resolves to the changed entry, or to + * `undefined` for a thread that has no entry. + */ + rename: ( + threadId: string, + title: string, + ) => Promise + /** + * Change the entry of a thread from app code, with no session open: its + * `title`, `pinned`, or `metadata`. Each field you pass replaces the stored + * one. Pass a function of the current entry to change part of `metadata`: + * with a store that has `upsertIf`, it runs again when another host wrote + * first, so neither change is lost. The harness `metadataSchema` checks + * `metadata`, and types it. Resolves to the changed entry, or `undefined` + * for a thread that has no entry. Throws `SessionMetadataError` for + * metadata the schema refuses, and an error when the entry names another + * harness. + * + * @example + * ```ts + * await host.sessions.update(assistant, threadId, (entry) => ({ + * metadata: { ...entry.metadata, pinned: [...(entry.metadata?.pinned ?? []), file] }, + * })) + * ``` + */ + update: ( + harness: THarness, + threadId: string, + update: SessionEntryUpdate>, + ) => Promise> | undefined> + /** + * Remove the index entry of a thread. Only the entry: the transcript, the + * log, and the other data of the thread stay in their stores. + */ + delete: (threadId: string) => Promise + /** + * Fork a thread into a new thread with `host.fork`, so the fork gets the + * settings, plugin config, and media too. `{ through: id }` copies the + * messages up to and including `id`. `{ before: id }` copies the messages + * before `id`, so a fork before the first message has no messages. The new + * thread gets its own entry, with the title of the old one plus ` (fork)`, + * the harness, and the owner of the old one. Resolves to the new entry. + * Throws when the message id is not in the transcript, or when the entry + * of the thread names another harness. + * + * @example + * ```ts + * const entry = await host.sessions.fork(assistant, 't-1', { before: messageId }) + * const fork = await host.open(assistant, { threadId: entry.threadId }) + * ``` + */ + fork: ( + harness: AnyHarness, + threadId: string, + at: ForkPoint, + ) => Promise +} + +/** + * An open session of the host changed its status. `running`: an operation + * runs. `waiting`: a question or an interrupt waits for an answer. `idle`: + * nothing runs or waits. `at` is the time of the change, in epoch ms. + */ +export interface HostStatusEvent { + type: 'status' + threadId: string + status: 'idle' | 'running' | 'waiting' + at: number +} + +/** + * The host wrote the index entry of a thread: a new session, the end of a + * turn, a rename, or a fork. `entry` is the new entry. + */ +export interface HostSessionEvent { + type: 'session' + threadId: string + entry: SessionIndexEntry +} + +/** The host removed the index entry of a thread. `entry` is the removed entry. */ +export interface HostSessionDeletedEvent { + type: 'session-deleted' + threadId: string + entry: SessionIndexEntry +} + +/** What `host.events()` yields. */ +export type HostEvent = + | HostStatusEvent + | HostSessionEvent + | HostSessionDeletedEvent + +export interface ForkSessionOptions { + /** The thread to copy. */ + threadId: string + /** The new thread. It must have no messages yet. */ + newThreadId: string + /** + * The id of the last message to copy. `null` copies no messages, but still + * copies the settings and plugin config. Default: the whole transcript. An + * id the thread does not have throws. + */ + at?: string | null + /** Who opens the new thread, as in `open`. */ + principal?: Principal +} + +export interface ResumePendingOptions { + /** The harnesses whose threads this sweep may open, matched by name. */ + 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. + */ + close?: 'whenIdle' | 'never' + /** How many expired claims one call reads. Default 100. */ + limit?: number +} + +/** Runs sessions for one or more harnesses in this process. */ +export interface HarnessHost { + /** + * Open a session, or return the live one for this harness and thread. + * Session plugins are set up here, so a plugin error rejects the promise. + */ + open: ( + harness: THarness, + options: OpenSessionOptions, + ) => Promise> + /** + * Start a new thread as a copy of another one, and return its open + * session. It copies the transcript up to and including the message `at`, + * the media that those messages use, the stored settings, and the plugin + * config. It does not copy plugin state, pending interrupts, queued + * inputs, or running work. A `newThreadId` with messages is refused. + * + * @example + * ```ts + * const fork = await host.fork(assistant, { threadId: 't-1', newThreadId: 't-2', at: messageId }) + * ``` + */ + fork: ( + harness: THarness, + options: ForkSessionOptions, + ) => Promise> + /** + * Continue the work of hosts that stopped. Needs `stores.workClaims`. It + * opens each thread whose claim expired, for a harness in `harnesses`, and + * the session recovers its pending work. Two hosts can call it at once: + * each thread goes to one of them. Call it at boot, and from a cron job or + * a Durable Object alarm. Resolves with the threads it opened. + * + * Known limit: a thread that this host has open already is claimed, but + * its session does not recover again. Call `host.recover(threadId)` for it. + * + * @example + * ```ts + * await host.resumePending({ harnesses: [assistant] }) + * ``` + */ + resumePending: ( + options: ResumePendingOptions, + ) => Promise> + /** + * Run `session.recover()` on the open sessions of `threadId`, or on every + * open session without it. A turn that another host ran when its session + * opened runs here once that host's lease expires. + * + * @example + * ```ts + * await host.recover('thread-1') + * ``` + */ + recover: (threadId?: string) => Promise + /** + * Run `session.reload()` on every open session of `harness`. Rejects when + * one of them fails. The other sessions still reload. + * + * @example + * ```ts + * await host.reload(assistant) + * ``` + */ + 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`. + * + * @example + * ```ts + * await host.close({ recoverable: true }) + * ``` + */ + close: (options?: { recoverable?: boolean }) => Promise + /** + * The `reduce` fold of a log that has an open session in this host, or + * `undefined`. + */ + logState: (logId: string) => TLogState | undefined + /** The session index: list, rename, delete, and fork sessions. */ + sessions: HostSessions + /** + * The status changes and the session index changes of this host, for a + * list of sessions that stays current. First the current status of each + * open session, then each change as it occurs. Each call reads on its own. + * The read stops when `signal` aborts or the loop ends. + * + * @example + * ```ts + * for await (const event of host.events({ signal })) { + * if (event.type === 'status') console.log(event.threadId, event.status) + * } + * ``` + */ + events: (options?: { signal?: AbortSignal }) => AsyncIterable +} + +/** Throw when `entry` names another harness than `harness`. */ +function checkHarness( + entry: SessionIndexEntry | undefined, + harness: AnyHarness, +) { + if (entry?.harness !== undefined && entry.harness !== harness.name) { + throw new Error( + `Thread ${JSON.stringify(entry.threadId)} belongs to harness ${JSON.stringify(entry.harness)}, not ${JSON.stringify(harness.name)}.`, + ) + } +} + +let warned = false + +/** + * Writes to the session index. `upsert` replaces the whole entry, so a + * change reads the entry, changes a copy, and writes it back. The changes of + * one thread run one at a time, so no change loses the fields of another. + * A store with `getVersioned` and `upsertIf` also stops the lost writes of + * other hosts: a change that another host wrote first runs again on the new + * entry. `onChange` gets each write and each removal. + */ +function sessionIndex( + store: SessionIndexStore | undefined, + onChange: (event: HostSessionEvent | HostSessionDeletedEvent) => void, +) { + const tails = new Map>() + const serial = (threadId: string, task: () => Promise) => { + // The tail never rejects, so a failed change does not stop the next one. + const run = (tails.get(threadId) ?? Promise.resolve()).then(task) + const tail = run.then( + () => {}, + () => {}, + ) + tails.set(threadId, tail) + void tail.then(() => { + if (tails.get(threadId) === tail) tails.delete(threadId) + }) + return run + } + return { + get: async (threadId: string) => store?.get(threadId), + /** + * Write `change(entry)` for the thread. `change` gets `undefined` when + * the thread has no entry, and returns `undefined` to write nothing. + */ + update: ( + threadId: string, + change: ( + entry: SessionIndexEntry | undefined, + ) => + | SessionIndexEntry + | undefined + | Promise, + ) => + serial(threadId, async () => { + if (!store) return undefined + if (!store.getVersioned || !store.upsertIf) { + const next = await change(await store.get(threadId)) + if (next) { + await store.upsert(next) + onChange({ type: 'session', threadId, entry: next }) + } + return next + } + for (let attempt = 0; attempt < 5; attempt += 1) { + const current = await store.getVersioned(threadId) + const next = await change(current?.entry) + if (!next) return undefined + const written = await store.upsertIf(next, current?.revision ?? null) + if (!written.ok) continue + onChange({ type: 'session', threadId, entry: next }) + return next + } + throw new Error( + `The session index entry of ${threadId} changed on another host 5 times in a row.`, + ) + }), + remove: (threadId: string) => + serial(threadId, async () => { + if (!store) return + const entry = await store.get(threadId) + await store.delete(threadId) + if (entry) onChange({ type: 'session-deleted', threadId, entry }) + }), + } +} + +/** @internal How a session writes its index entry. */ +export type SessionIndexWriter = ReturnType + +/** + * The status rule: the status of `session` after `event`, or `undefined` + * when the event changes nothing. An operation that starts runs, and a + * question waits. When an operation ends or a question gets its answer, the + * session waits while a question or an interrupt still waits, runs while + * another operation runs, and else is idle. + */ +function statusAfter(session: HarnessSession, event: StreamChunk) { + if (event.type !== EventType.CUSTOM) return undefined + switch (event.name) { + case HARNESS_EVENTS.operationStarted: + return 'running' + case HARNESS_EVENTS.question: + return 'waiting' + case HARNESS_EVENTS.operationFinished: + case HARNESS_EVENTS.questionAnswered: { + const { pendingQuestions, pendingInterrupts, activeOperations } = + session.snapshot() + const isInterrupted = + isRecord(event.value) && event.value.status === 'interrupted' + const stillWaits = + isInterrupted || + pendingQuestions.length > 0 || + pendingInterrupts.length > 0 + if (stillWaits) return 'waiting' + return activeOperations.length > 0 ? 'running' : 'idle' + } + default: + return undefined + } +} + +/** How each host finds the open session of a chat turn. Not public API. */ +const turnFinders = new WeakMap< + object, + (harness: string, operationId: string) => Promise +>() + +/** + * @internal The open session of `harness` in `host` that has the chat turn + * `operationId`, running or ended. A turn's events live in the host that runs + * it, so the handler joins a run through this. + */ +export function sessionOfTurn( + host: HarnessHost, + harness: AnyHarness, + operationId: string, +): Promise { + return ( + turnFinders.get(host)?.(harness.name, operationId) ?? + Promise.resolve(undefined) + ) +} + +/** Throw a clear error for a durable store set that cannot work. */ +function checkDurableStores(persistence: HarnessPersistence) { + const { stores } = persistence + if (stores.log === undefined) return + if (stores.runs === undefined && stores.leases === undefined) { + throw new Error( + 'A durable host (stores.log) needs stores.runs or stores.leases: a lease finds a crashed host.', + ) + } + if (stores.messages !== undefined || stores.inbox !== undefined) { + throw new Error( + 'A durable host (stores.log) keeps the transcript and the inputs in the log. Remove stores.messages and stores.inbox. A reader outside a session can use logMessageStore({ store }).', + ) + } +} + +/** + * Create a host for harness sessions. + * + * @example + * ```ts + * const host = createHarnessHost({ persistence }) + * const session = await host.open(studio, { threadId: 'thread-1' }) + * const turn = await session.prompt('Write a haiku about the sea.') + * ``` + */ +export function createHarnessHost( + options: HarnessHostOptions = {}, +): HarnessHost { + if (!options.persistence && !warned) { + warned = true + console.warn( + '[@tanstack/ai-harness] No persistence given: sessions live in memory and are lost on restart.', + ) + } + const persistence: HarnessPersistence = + options.persistence ?? memoryPersistence() + checkDurableStores(persistence) + const memory = memoryPersistence().stores + // ponytail: a memory inbox when the stores have none. Pass `stores.inbox` + // to keep accepted inputs across restarts. A durable host keeps them in + // the log instead. + const inbox = persistence.stores.inbox ?? memory.inbox + // ponytail: memory credentials when the stores have none. Pass + // `stores.credentials` to keep sign-ins across restarts. + const credentials = persistence.stores.credentials ?? memory.credentials + // ponytail: memory media stores when the stores have none. Pass + // `stores.artifacts`, `stores.blobs`, and `stores.generationRuns` to keep + // media across restarts. + const media = { + stores: { + artifacts: persistence.stores.artifacts ?? memory.artifacts, + blobs: persistence.stores.blobs ?? memory.blobs, + generationRuns: + persistence.stores.generationRuns ?? memory.generationRuns, + }, + } + const logStore = persistence.stores.log + const { metadata, sessions: sessionStore } = persistence.stores + /** One per running `host.events()` read. */ + const listeners = new Set<(event: HostEvent) => void>() + const emit = (event: HostEvent) => { + for (const listener of listeners) listener(event) + } + /** The latest status of each open session, by session key. */ + const statuses = new Map() + const setStatus = ( + key: string, + threadId: string, + status: HostStatusEvent['status'], + ) => { + if (statuses.get(key)?.status === status) return + const event: HostStatusEvent = { + type: 'status', + threadId, + status, + at: Date.now(), + } + statuses.set(key, event) + emit(event) + } + const index = sessionIndex(sessionStore, emit) + const coalesceMs = options.coalesceMs ?? 100 + /** The shared log of each log id with an open session, by log id. */ + const logs = new Map>>() + /** The resolved shared logs, for `logState`. */ + const resolved = new Map>() + const sharedLog = (store: LogStore, logId: string) => { + const open = logs.get(logId) + if (open) return open + const loaded = loadLogState({ + store, + logId, + ...(metadata ? { metadata } : {}), + ...(options.project ? { project: options.project } : {}), + ...(options.reduce ? { reduce: options.reduce } : {}), + }).then((state) => { + const log = new SharedLog({ + store, + logId, + state, + coalesceMs, + ...(options.project ? { project: options.project } : {}), + ...(options.reduce ? { reduce: options.reduce } : {}), + ...(metadata ? { metadata } : {}), + // The next open folds the log again. + onIdle: () => { + if (logs.get(logId) === loaded) logs.delete(logId) + if (resolved.get(logId) === log) resolved.delete(logId) + }, + }) + resolved.set(logId, log) + return log + }) + loaded.catch(() => logs.delete(logId)) + logs.set(logId, loaded) + return loaded + } + const sessions = 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. + */ + const sweptOnly = new Set() + /** The closes of swept 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 host: HarnessHost = { + open(harness, { threadId, principal, logId: openLogId, promptCache }) { + const key = `${harness.name}\u0000${threadId}` + let session = sessions.get(key) + // The app opened a thread that a sweep opened: it stays open. + if (session) sweptOnly.delete(key) + if (!session) { + const logId = openLogId ?? threadId + const created = new HarnessSession({ + harness, + threadId, + logId, + persistence, + inbox, + credentials, + media, + hostId, + index, + ...(logStore + ? { + log: { + store: logStore, + ...(options.project ? { project: options.project } : {}), + open: async (sessionThread, onFailure) => { + let log = await sharedLog(logStore, logId) + // The log closed while this open waited, and the host + // dropped it. + while (!log.isOpen) log = await sharedLog(logStore, logId) + return log.view(sessionThread, onFailure) + }, + }, + } + : {}), + ...(options.lease ? { lease: options.lease } : {}), + ...(principal ? { principal } : {}), + ...(promptCache ? { promptCache } : {}), + onEvent: (event) => { + const status = statusAfter(created, event) + if (status) setStatus(key, threadId, 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 swept session under this key may still close: it lets go of + // the thread first. + const closed = closing.get(key) ?? Promise.resolve() + session = closed + .then(() => created.open()) + .then( + () => { + // A recovered turn can run already. Else the new session is + // idle. + if (!statuses.has(key)) setStatus(key, threadId, 'idle') + return created + }, + (error: unknown) => { + sessions.delete(key) + statuses.delete(key) + throw error + }, + ) + sessions.set(key, session) + } + return session as Promise> + }, + async fork(harness, { threadId, newThreadId, at, principal }) { + if (newThreadId === threadId) { + throw new Error('A fork needs a new thread: newThreadId is threadId.') + } + const source = await host.open(harness, { threadId }) + const transcript = await source.transcript() + const end = + at === undefined + ? transcript.length + : at === null + ? 0 + : transcript.findIndex((message) => message.id === at) + 1 + if (typeof at === 'string' && end === 0) { + throw new Error( + `Thread ${JSON.stringify(threadId)} has no message ${JSON.stringify(at)}.`, + ) + } + const isUsed = logStore + ? (await logStore.read(newThreadId, { limit: 1 })).length > 0 + : ((await persistence.stores.messages?.loadThread(newThreadId)) + ?.length ?? 0) > 0 + if (isUsed) { + throw new Error( + `Thread ${JSON.stringify(newThreadId)} is not empty. Fork into a new thread.`, + ) + } + const fork = await host.open(harness, { + threadId: newThreadId, + ...(principal ? { principal } : {}), + }) + await fork.adoptFork(source, transcript.slice(0, end)) + return fork + }, + async resumePending({ harnesses, close = 'whenIdle', limit = 100 }) { + const claims = persistence.stores.workClaims + if (!claims) throw new Error('resumePending needs stores.workClaims.') + const ttlMs = options.lease?.ttlMs ?? LEASE.ttlMs + const byName = new Map( + harnesses.map((harness) => [harness.name, harness]), + ) + const expired = await claims.listExpired({ now: Date.now(), limit }) + const opened: Array<{ threadId: string; harness: string }> = [] + for (const { threadId, harness: name } of expired) { + const harness = byName.get(name) + if (!harness) continue + // Compare-and-set: of two hosts that sweep at once, one gets it. + const isMine = await claims.claim({ + threadId, + harness: name, + ownerId: hostId, + until: Date.now() + ttlMs, + }) + if (!isMine) continue + const key = `${name}\u0000${threadId}` + // Only a session that this sweep opens closes when idle. A session + // that the app has open stays open. + const isOwn = close === 'whenIdle' && !sessions.has(key) + if (isOwn) sweptOnly.add(key) + const session = await host.open(harness, { threadId }) + 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 }) + } + return opened + }, + async recover(threadId) { + const live = await Promise.allSettled(sessions.values()) + await Promise.all( + live.flatMap((entry) => + entry.status === 'fulfilled' && + (threadId === undefined || entry.value.threadId === threadId) + ? [entry.value.recover()] + : [], + ), + ) + }, + async reload(harness) { + const prefix = `${harness.name}\u0000` + const live = await Promise.allSettled( + [...sessions].flatMap(([key, session]) => + key.startsWith(prefix) ? [session] : [], + ), + ) + const results = await Promise.allSettled( + live.flatMap((entry) => + entry.status === 'fulfilled' ? [entry.value.reload()] : [], + ), + ) + for (const result of results) { + if (result.status === 'rejected') throw result.reason + } + }, + async close(closeOptions) { + await Promise.allSettled(closing.values()) + const live = await Promise.allSettled(sessions.values()) + await Promise.all( + live + .filter((entry) => entry.status === 'fulfilled') + .map((entry) => entry.value.close(closeOptions)), + ) + }, + logState: (logId) => resolved.get(logId)?.state.reduced, + sessions: { + list: async (listOptions) => + (await sessionStore?.list({ + parentThreadId: null, + ...listOptions, + })) ?? { + entries: [], + }, + get: (threadId) => index.get(threadId), + rename: (threadId, title) => + index.update( + threadId, + (entry) => entry && { ...entry, title, updatedAt: Date.now() }, + ), + update: async (harness, threadId, update) => + // The schema checked the metadata, so it has the harness type. + (await index.update(threadId, (entry) => { + checkHarness(entry, harness) + return changedEntry( + entry, + update as SessionEntryUpdate>, + harness.metadataSchema, + ) + })) as SessionIndexEntry> | undefined, + delete: (threadId) => index.remove(threadId), + async fork(harness, threadId, at) { + // Check before a session opens. Without a sessions store there is no + // entry, so there is nothing to check. + const source = await index.get(threadId) + checkHarness(source, harness) + // `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) { + const opened = await host.open(harness, { threadId }) + 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 + } + } + const forked = await host.fork(harness, { + threadId, + newThreadId: `thread-${crypto.randomUUID()}`, + at: last, + ...(source?.principal ? { principal: source.principal } : {}), + }) + const now = Date.now() + // A sibling of the old thread, not a child: no parentThreadId. + const entry: SessionIndexEntry = { + threadId: forked.threadId, + createdAt: now, + updatedAt: now, + harness: harness.name, + ...(source?.title ? { title: `${source.title} (fork)` } : {}), + ...(source?.principal ? { principal: source.principal } : {}), + } + await index.update(entry.threadId, () => entry) + return entry + }, + }, + async *events({ signal } = {}) { + // The current statuses first. A change while the reader handles an + // event waits in the queue. + const queue: Array = [...statuses.values()] + let wake = () => {} + const listener = (event: HostEvent) => { + queue.push(event) + wake() + } + const onAbort = () => wake() + listeners.add(listener) + signal?.addEventListener('abort', onAbort) + try { + while (!signal?.aborted) { + const next = queue.shift() + if (next) { + yield next + continue + } + await new Promise((resolve) => { + wake = resolve + }) + } + } finally { + // A loop that ends (break, return, throw, abort) stops listening. + listeners.delete(listener) + signal?.removeEventListener('abort', onAbort) + } + }, + } + turnFinders.set(host, async (name, operationId) => { + for (const [key, opening] of sessions) { + if (!key.startsWith(`${name}\u0000`)) continue + const session = await opening.catch(() => undefined) + if (session?.operation(operationId)?.kind === 'chat') return session + } + return undefined + }) + return host +} diff --git a/packages/ai-harness/src/media.ts b/packages/ai-harness/src/media.ts new file mode 100644 index 0000000000..371575708f --- /dev/null +++ b/packages/ai-harness/src/media.ts @@ -0,0 +1,472 @@ +// The server side of harness media: the store, capture of generated media, +// and the chat middleware that gives the model the bytes. +import { generateTranscription } from '@tanstack/ai' +import { + resolveBlobRange, + retrieveBlob, + withGenerationPersistence, +} from '@tanstack/ai-persistence' +import { kindOf, mediaIdOf } from './media-ref' +import type { + AnyTranscriptionAdapter, + ChatMiddleware, + ChatMiddlewareContext, + ContentPart, + GenerationMiddleware, + Modality, + ModelMessage, + PersistedArtifactRef, +} from '@tanstack/ai' +import type { + AIPersistence, + ArtifactRecord, + ArtifactStore, + BlobBody, + BlobRange, + BlobStore, + GenerationRunStore, +} from '@tanstack/ai-persistence' +import type { MediaKind, MediaRecord } from './types' + +/** The `media` option of a harness: limits, and what the model reads. */ +export interface MediaOptions { + /** Refuse a file bigger than this many bytes. Default 100 MB. */ + maxBytes?: number + /** The kinds a user can send. Default: all four. */ + kinds?: ReadonlyArray + /** The kinds the model reads. Narrows what the adapter says it reads. */ + accepts?: ReadonlyArray + /** Turns audio into text first when the model cannot read audio. */ + transcribe?: AnyTranscriptionAdapter +} + +/** + * A media request that failed. `status` is the HTTP status a route answers: + * 404 (not found, or another thread's), 413 (too big), 415 (type not allowed). + */ +export class MediaError extends Error { + readonly status: 404 | 413 | 415 + constructor(status: 404 | 413 | 415, message: string) { + super(message) + this.name = 'MediaError' + this.status = status + } +} + +/** The stores media lives in. */ +export type MediaPersistence = AIPersistence<{ + artifacts: ArtifactStore + blobs: BlobStore +}> + +/** The media store of one thread. */ +export type MediaStore = ReturnType + +const DEFAULT_MAX_BYTES = 100 * 1024 * 1024 +const DATA_URL_MAX_BYTES = 1024 * 1024 +const ALL_KINDS: ReadonlyArray = [ + 'image', + 'audio', + 'video', + 'document', +] + +// Uploads are artifact records with the run id `upload:`, and their +// bytes use the default artifact key `artifacts//`. Generated media +// keeps what `withGenerationPersistence` gives it: a generation run id +// `[:...]:-`. +function uploadRunId(threadId: string) { + return `upload:${threadId}` +} + +function toMediaRecord(record: ArtifactRecord, kind: MediaKind) { + const isUpload = record.runId === uploadRunId(record.threadId) + const segments = isUpload ? [] : record.runId.split(':') + const [runId] = segments + const subagentRunId = segments.findLast((segment) => + segment.startsWith('subagent-'), + ) + const media: MediaRecord = { + id: record.artifactId, + threadId: record.threadId, + kind, + mimeType: record.mimeType, + name: record.name, + size: record.size, + source: isUpload ? 'user' : 'generated', + createdAt: record.createdAt, + ...(runId ? { runId } : {}), + ...(subagentRunId ? { subagentRunId } : {}), + } + return media +} + +/** The byte length of a body, or `undefined` for a stream. */ +function sizeOf(body: BlobBody) { + if (typeof body === 'string') return new TextEncoder().encode(body).byteLength + if (body instanceof Blob) return body.size + if (body instanceof ReadableStream) return undefined + return body.byteLength +} + +/** + * Base64 of `bytes`. It uses `btoa`, not the Node `Buffer`, so it runs on + * workerd too. + */ +export function toBase64(bytes: Uint8Array) { + const chunk = 0x8000 + let binary = '' + // Chunks keep `fromCharCode` under the engine's argument limit. + for (let offset = 0; offset < bytes.length; offset += chunk) { + binary += String.fromCharCode(...bytes.subarray(offset, offset + chunk)) + } + return btoa(binary) +} + +/** + * The media store of one thread, on the persistence `artifacts` and `blobs` + * stores. It sees only the media of `threadId`: an id of another thread reads + * as not found. + * + * `put` throws a `MediaError` with 413 (over `maxBytes`) or 415 (a MIME type + * with no media kind, or a kind not in `kinds`). `load` throws a `MediaError` + * with 404 when the record or its bytes are missing. + * + * @example + * const media = createMediaStore({ persistence, threadId }) + * const record = await media.put(bytes, { mimeType: 'image/png', name: 'cat.png' }) + * const url = await media.dataUrl(record.id) + */ +export function createMediaStore({ + persistence, + threadId, + options = {}, +}: { + persistence: MediaPersistence + threadId: string + options?: MediaOptions +}) { + const { artifacts, blobs } = persistence.stores + const maxBytes = options.maxBytes ?? DEFAULT_MAX_BYTES + const kinds = options.kinds ?? ALL_KINDS + const tooBig = () => + new MediaError( + 413, + `The file is bigger than the limit of ${maxBytes} bytes.`, + ) + + async function artifact(id: string) { + const record = await artifacts.get(id) + return record?.threadId === threadId ? record : null + } + + async function get(id: string) { + const record = await artifact(id) + const kind = record ? kindOf(record.mimeType) : undefined + return record && kind ? toMediaRecord(record, kind) : null + } + + async function load(id: string, range?: BlobRange) { + const record = await artifact(id) + const blob = record + ? await retrieveBlob( + persistence, + record, + range ? { range: resolveBlobRange(record.size, range) } : undefined, + ) + : null + if (!blob) throw new MediaError(404, `Media ${id} was not found.`) + return new Uint8Array(await blob.arrayBuffer()) + } + + return { + /** Store a user file. Returns its record. */ + async put(body: BlobBody, info: { mimeType: string; name: string }) { + const kind = kindOf(info.mimeType) + if (kind === undefined) { + throw new MediaError( + 415, + `Files of type ${info.mimeType} are not supported.`, + ) + } + if (!kinds.includes(kind)) { + throw new MediaError(415, `This harness does not accept ${kind} files.`) + } + const size = sizeOf(body) + if (size !== undefined && size > maxBytes) throw tooBig() + + const id = crypto.randomUUID() + const runId = uploadRunId(threadId) + const blobKey = `artifacts/${runId}/${id}` + // A stream has no length up front, so count it as the store reads it. + let seen = 0 + const counted = + body instanceof ReadableStream + ? body.pipeThrough( + new TransformStream({ + transform(bytes: Uint8Array, controller) { + seen += bytes.byteLength + if (seen > maxBytes) controller.error(tooBig()) + else controller.enqueue(bytes) + }, + }), + ) + : body + try { + await blobs.put(blobKey, counted, { contentType: info.mimeType }) + } catch (error) { + if (seen <= maxBytes) throw error + await blobs.delete(blobKey) + throw tooBig() + } + const record: ArtifactRecord = { + artifactId: id, + runId, + threadId, + blobKey, + name: info.name, + mimeType: info.mimeType, + size: size ?? seen, + createdAt: Date.now(), + } + await artifacts.save(record) + return toMediaRecord(record, kind) + }, + /** The record of `id`, or `null` when it is missing or another thread's. */ + get, + /** The bytes of `id`, or one `range` of them. */ + load, + /** A data URL for a file up to 1 MB. `undefined` above that. */ + async dataUrl(id: string) { + const record = await get(id) + if (!record || record.size > DATA_URL_MAX_BYTES) return undefined + return `data:${record.mimeType};base64,${toBase64(await load(id))}` + }, + } +} + +/** + * Generation middleware that keeps the media an agent makes and publishes a + * record for each new file. It is `withGenerationPersistence` plus a result + * transform that runs after it. + * + * A failed save or a failed `publish` does not fail the generation: the agent + * gets its result, and `onError` gets the error, so the caller can warn. + * + * @example + * const middleware = mediaCapture({ + * persistence, + * threadId, + * publish: (record) => console.log(`saved ${record.name}`), + * onError: (error) => console.warn('media not kept', error), + * }) + * await generateImage({ adapter, prompt, threadId, middleware }) + */ +export function mediaCapture({ + persistence, + threadId, + options = {}, + publish, + onError, +}: { + persistence: AIPersistence<{ + artifacts: ArtifactStore + blobs: BlobStore + generationRuns: GenerationRunStore + }> + threadId: string + options?: MediaOptions + publish: (record: MediaRecord) => void | Promise + onError?: (error: unknown) => void +}) { + const store = createMediaStore({ persistence, threadId, options }) + const persist = withGenerationPersistence(persistence, { + threadId, + maxArtifactBytes: options.maxBytes ?? DEFAULT_MAX_BYTES, + }) + const keep: GenerationMiddleware = { + ...persist, + async onStart(ctx) { + const from = ctx.resultTransforms.length + await persist.onStart?.(ctx) + // Guard the transforms that save the bytes, so a store failure (or a + // file over the limit) keeps the agent's result. + const added = ctx.resultTransforms.splice(from) + for (const transform of added) { + ctx.resultTransforms.push(async (result, transformCtx) => { + try { + return await transform(result, transformCtx) + } catch (error) { + onError?.(error) + return undefined + } + }) + } + }, + } + const announce: GenerationMiddleware = { + name: 'harness:media-capture', + onStart(ctx) { + const runId = ctx.runId ?? ctx.requestId + // `artifacts` is what the persistence transform above added. + ctx.resultTransforms.push( + async (result: { artifacts?: Array }) => { + const refs = (result.artifacts ?? []).filter( + (ref) => ref.role === 'output' && ref.runId === runId, + ) + for (const ref of refs) { + try { + const record = await store.get(ref.artifactId) + if (record) await publish(record) + } catch (error) { + onError?.(error) + } + } + return undefined + }, + ) + }, + } + return [keep, announce] +} + +type MediaContentPart = Exclude + +function textPart(content: string) { + const part: ContentPart = { type: 'text', content } + return part +} + +function cannotRead(model: string, kind: MediaKind) { + const fix = + kind === 'audio' + ? 'Add media.transcribe' + : `Use a model that reads ${kind} files` + return `${model} cannot read ${kind} files. ${fix}, or send a text summary.` +} + +function hasMediaRef(message: ModelMessage) { + return ( + Array.isArray(message.content) && + message.content.some((part) => mediaIdOf(part) !== undefined) + ) +} + +/** + * Chat middleware that gives the model the bytes of `harness-media:` parts. + * It changes only what the adapter gets (`providerMessages`), so the saved + * transcript keeps the small `harness-media:` URLs. + * + * - A part of a kind in `accepted` becomes a base64 data source. With + * `accepted` undefined, every kind is sent. + * - Audio not in `accepted` becomes a transcript when `transcribe` is set. + * - Any other kind not in `accepted` stops the run with a clear error before + * the model call. + * - A part whose file is gone becomes a short text note. + * + * List it after `withPersistence` and after any middleware that returns + * `messages`: a later `messages` result resets the provider messages. + * + * @example + * chat({ + * adapter, + * messages, + * middleware: [withPersistence(persistence), mediaMiddleware({ store })], + * }) + */ +export function mediaMiddleware({ + store, + accepted, + transcribe, +}: { + store: MediaStore + accepted?: ReadonlyArray + transcribe?: AnyTranscriptionAdapter +}) { + // ponytail: one cache per run, so the loop does not load or transcribe a + // file again at every model call. A new turn does it again for history. + const runs = new WeakMap>>() + + async function resolve( + ctx: ChatMiddlewareContext, + part: MediaContentPart, + id: string, + ) { + const record = await store.get(id) + if (!record) return textPart(`[${part.type} not found: ${id}]`) + const isAccepted = accepted === undefined || accepted.includes(record.kind) + const transcriber = + isAccepted || record.kind !== 'audio' ? undefined : transcribe + if (!isAccepted && transcriber === undefined) { + throw new Error(cannotRead(ctx.model, record.kind)) + } + const bytes = await store.load(id).catch((error: unknown) => { + if (error instanceof MediaError && error.status === 404) return undefined + throw error + }) + if (!bytes) return textPart(`[${record.kind} not found: ${record.name}]`) + if (transcriber === undefined) { + const swapped: ContentPart = { + ...part, + source: { + type: 'data', + value: toBase64(bytes), + mimeType: record.mimeType, + }, + } + return swapped + } + const transcript = await generateTranscription({ + adapter: transcriber, + audio: new Blob([bytes], { type: record.mimeType }), + ...(ctx.signal ? { abortSignal: ctx.signal } : {}), + }) + return textPart(`[audio transcript: ${record.name}] ${transcript.text}`) + } + + async function resolveMessage( + ctx: ChatMiddlewareContext, + cache: Map>, + message: ModelMessage, + ) { + const { content } = message + if (!Array.isArray(content) || !hasMediaRef(message)) return message + const parts: Array = [] + for (const part of content) { + const id = part.type === 'text' ? undefined : mediaIdOf(part) + if (part.type === 'text' || id === undefined) { + parts.push(part) + continue + } + let pending = cache.get(id) + if (!pending) { + pending = resolve(ctx, part, id) + cache.set(id, pending) + } + parts.push(await pending) + } + return { ...message, content: parts } + } + + const middleware: ChatMiddleware = { + name: 'harness:media', + // `init` runs before the user turn is stored, so a refusal stops the send + // early. `beforeModel` builds the provider messages again, so it runs + // there too (the cache makes that cheap). + async onConfig(ctx, config) { + const messages = config.providerMessages ?? config.messages + if (!messages.some(hasMediaRef)) return undefined + let cache = runs.get(ctx) + if (!cache) { + cache = new Map() + runs.set(ctx, cache) + } + const providerMessages: Array = [] + for (const message of messages) { + providerMessages.push(await resolveMessage(ctx, cache, message)) + } + return { providerMessages } + }, + } + return middleware +} diff --git a/packages/ai-harness/src/session/agents.ts b/packages/ai-harness/src/session/agents.ts new file mode 100644 index 0000000000..4ba5817791 --- /dev/null +++ b/packages/ai-harness/src/session/agents.ts @@ -0,0 +1,1029 @@ +import { + EventType, + compactForModel, + createSubagentId, + runAgentStream, + validateWithStandardSchema, +} from '@tanstack/ai' +import { withPersistence } from '@tanstack/ai-persistence' +import { withHarnessContext } from '../agents' +import { mediaMiddleware } from '../media' +import { OperationImpl } from '../operation' +import { checkpointMiddleware, holdRunLease, repairTranscript } from '../resume' +import { HARNESS_EVENTS } from '../types' +import { + agentKey, + chainFields, + asUserMessage, + chainInputOf, + runStatus, + loggedStatus, + createInputId, + createMessageId, + customEvent, + isSameSender, + resetCut, + referenceNote, + AGENT_STOPPED, + SHUTDOWN_REASON, +} from './helpers' +import { SessionTurnParts } from './turn-parts' +import type { + SubagentBudget, + ModelMessage, + RunRecord, + StreamChunk, + SubagentBinding, +} from '@tanstack/ai' +import type { InputState, LogWriter } from '../log' +import type { AgentStarter, AnyAgent } from '../agents' +import type { AnyHarness } from '../define' +import type { AgentGroup, MountedPlugins } from '../plugins' +import type { + AgentRun, + ChatTurnResult, + InputSettlement, + MediaRecord, + Principal, + Receipt, + UserInput, +} from '../types' +import type { + AgentStartOptions, + AgentChild, + DynamicAgentHandle, + AgentRunInfo, + AgentRunImpl, + AgentMessage, + AgentRunEntry, + AgentChain, + AgentRunStart, + AgentInputState, +} from './types' + +/** + * A layer of {@link HarnessSession}: agent runs started from code. + */ +export abstract class SessionAgents< + THarness extends AnyHarness = AnyHarness, +> extends SessionTurnParts { + /** + * The agent named `name`, for names known only at runtime (a slash + * command, a protocol input). `undefined` when no such agent exists. + */ + agent(name: string): DynamicAgentHandle | undefined { + if (!this.agentRegistry.get(name)) return undefined + return { + run: (input, options) => + this.runAgent(name, input, { ...options, wake: false }), + start: (input, options) => this.runAgent(name, input, options ?? {}), + } + } + + /** + * Add a message to the agent run `operationId`, as `AgentRun.send` does. + * `principal` is who sent it. A follow-up runs for its sender. An unknown + * run gets `not_running`. + */ + async sendToAgent( + operationId: string, + message: UserInput, + options?: { + mode?: 'steer' | 'followUp' + inputId?: string + principal?: Principal + }, + ): Promise { + const inputId = options?.inputId ?? createInputId() + const chain = this.chainOf.get(operationId) ?? this.loggedChain(operationId) + if (!chain) return { inputId, status: 'rejected', reason: 'not_running' } + const principal = options?.principal ?? this.principal + const mode = options?.mode ?? 'steer' + return this.admitting(async () => { + const admission = await this.accept( + inputId, + { + op: 'agentMessage', + operationId, + message, + mode, + run: chain.runInputId, + }, + principal, + ) + if (admission !== 'new') return this.duplicateReceipt(inputId, admission) + const sent = { inputId, message, ...(principal ? { principal } : {}) } + return this.keep({ inputId, ...this.deliver(chain, sent, mode) }) + }) + } + + /** + * 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. + */ + agentRun(operationId: string): AgentRun | undefined { + const chain = this.chainOf.get(operationId) ?? this.loggedChain(operationId) + return [...(chain?.runs ?? []), ...(chain?.followUps ?? [])].find( + (run) => run.operation.id === operationId, + )?.operation + } + + /** + * 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. + */ + agentRuns(): Array { + const live = [...this.agentChains.values()].flatMap((chain) => + [...chain.runs, ...chain.followUps].map( + ({ operation, principal }): AgentRunInfo => ({ + operationId: operation.id, + agent: chain.agent, + status: runStatus(operation.status()), + ...(chain.parentRunId ? { parentRunId: chain.parentRunId } : {}), + ...(principal ? { principal } : {}), + }), + ), + ) + const known = new Set(live.map((run) => run.operationId)) + return [ + ...this.loggedRuns().filter((run) => !known.has(run.operationId)), + ...live, + ] + } + + /** The agent runs in the log: each input that a run applied. */ + protected loggedRuns(): Array { + const inputs = this.writer?.state.inputs + if (!inputs) return [] + return [...inputs.values()].flatMap((input): Array => { + const first = chainInputOf(inputs, input) + if (!first || input.operationId === undefined) return [] + return [ + { + operationId: input.operationId, + agent: first.input.agent, + status: loggedStatus(input), + ...(first.input.parentRunId + ? { parentRunId: first.input.parentRunId } + : {}), + ...(input.principal ? { principal: input.principal } : {}), + }, + ] + }) + } + + /** + * The chain of a run in the log, with every run it had, so a follow-up + * continues its thread after a restart. `undefined` for an unknown run, + * and while a run of the chain has not ended: another host runs it. + */ + protected loggedChain(operationId: string): AgentChain | undefined { + const inputs = this.writer?.state.inputs + if (!inputs) return undefined + const own = [...inputs.values()].find( + (input) => input.operationId === operationId, + ) + const first = own && chainInputOf(inputs, own) + if (!first) return undefined + const known = this.agentChains.get(first.inputId) + if (known) return known + const logged = [...inputs.values()].filter( + (input) => + input.operationId !== undefined && + chainInputOf(inputs, input) === first, + ) + if (logged.some((input) => input.settlement === undefined)) { + return undefined + } + // ponytail: a run rebuilt from the log resolves to `undefined`. The log + // keeps how it ended, not its value. Read `stores.runs` if callers need + // the value. + const entries = logged.map((input) => { + const operation = this.agentOperation( + first.input.agent, + input.operationId, + ) + const settlement = input.settlement + if (settlement?.outcome === 'completed') { + operation.finish('completed', undefined) + } else { + operation.fail( + settlement?.outcome === 'aborted' ? 'cancelled' : 'failed', + new Error(settlement?.error?.message ?? 'The run did not complete.'), + ) + } + return { + operation, + ...(input.principal ? { principal: input.principal } : {}), + } + }) + const [firstRun, ...later] = entries + if (!firstRun) return undefined + const chain = this.createChain(chainFields(first), firstRun) + for (const entry of later) { + chain.runs.push(entry) + this.chainOf.set(entry.operation.id, chain) + } + chain.current = later.at(-1) ?? firstRun + return chain + } + + /** + * What every agent run of `operation` gets: the middleware of session + * plugins, then of run plugins, then the session's media middleware. It + * keeps the media the agents make, publishes a `harness.media` event for + * each file, and pushes its record to `captured`. The agents read the + * session's provider keys as `ctx.keys`. Their model calls count in the + * usage of `principal`. With `agents` (the subagents of a turn), the + * agent tools can start them in the background. + */ + protected binding( + operation: OperationImpl | OperationImpl, + captured: Array, + principal: Principal | undefined, + runPlugins?: MountedPlugins, + agents?: ReadonlyArray, + ) { + const options = this.harness.media + // The child runs as a background agent of this session, and its end + // starts a new turn. Core does not pass `start` to nested children. + const start: SubagentBinding['start'] = + agents && + (async (call, startOptions) => { + const agent = agents.find((entry) => entry.name === call.agent) + if (!agent) throw new Error(`Unknown agent "${call.agent}".`) + const subagentRunId = createSubagentId() + this.runAgent(agent, call.input, startOptions, undefined, { + subagentRunId, + ...(call.prompt !== undefined && { prompt: call.prompt }), + ...(call.parentToolCallId !== undefined && { + parentToolCallId: call.parentToolCallId, + }), + }) + return { subagentRunId } + }) + return { + generationMiddleware: [ + ...(this.sessionPlugins?.generationMiddleware ?? []), + ...(runPlugins?.generationMiddleware ?? []), + ...mediaCapture({ + persistence: this.media, + threadId: this.threadId, + options, + publish: (record) => { + captured.push(record) + operation.publish(customEvent(HARNESS_EVENTS.media, { ...record })) + }, + onError: (error) => + this.warn( + operation, + 'harness:media', + `The media was not kept. ${String(error)}`, + ), + }), + ], + chatMiddleware: [ + ...(this.sessionPlugins?.agentMiddleware ?? []), + ...(runPlugins?.agentMiddleware ?? []), + // A routed root agent gets the thread too, so a reset cuts it as well. + resetCut, + this.usageCounter(operation, principal), + // Last, because a later middleware that returns `messages` resets + // what the model gets. A child's adapter is not known here, so only + // `media.accepts` narrows what it reads. + mediaMiddleware({ + store: this.mediaStore, + accepted: options?.accepts, + transcribe: options?.transcribe, + }), + ], + keys: this.keys, + // No `promptCache`: a tool-mode child then uses the retention of the + // turn's chat() call, so a turn override reaches it. + ...(start && { start }), + } satisfies SubagentBinding + } + + protected override runAgent( + target: string | AnyAgent, + input: unknown, + options: AgentStartOptions, + parent?: AgentChain, + child?: AgentChild, + ): AgentRunImpl { + const name = typeof target === 'string' ? target : target.name + if (options.resume && !this.writer) { + throw new Error( + `Agent ${name}: resume: true needs a durable host (a host with stores.log).`, + ) + } + const operation = this.agentOperation(name) + this.markBusy() + // The agent runs for the sender of the turn that starts it, also after + // that turn ends. A child of agent code runs for its parent run's sender. + const principal = parent ? parent.current.principal : this.sender() + const sender = principal ? { principal } : {} + const fromParent = parent + ? { + parentRunId: parent.current.operation.id, + budget: parent.budget.child(), + } + : {} + const chain = this.createChain( + { + runInputId: createInputId(), + target, + input, + options, + ...sender, + ...fromParent, + }, + { operation, ...sender }, + ) + parent?.children.add(chain) + void this.executeAgent(operation, chain, { + inputId: chain.runInputId, + ...(child ? { child } : {}), + }) + return operation + } + + /** + * `harnessContext(ctx).agents` in the runs of `chain`: start a child in the + * background. + * The child counts against the chain's tree budget, and over it `start` + * throws. The child records the run that started it, and runs for that + * run's sender. + */ + protected agentStarter(chain: AgentChain): AgentStarter { + return { + start: (target, input, options) => { + const active = [...chain.children].filter( + (child) => !child.current.operation.isSettled(), + ).length + const refusal = chain.budget.reserve(active) + if (refusal !== undefined) throw new Error(refusal) + return this.runAgent(target, input, options ?? {}, chain) + }, + } + } + + /** A new agent run operation. `id`: a stored one, after a restart. */ + protected agentOperation(agent: string, id?: string): AgentRunImpl { + const operation = new OperationImpl( + 'agent', + this.feed, + (running) => this.cancel(running.id), + agent, + id, + ) + const run = Object.assign(operation, { + send: ( + message: UserInput, + options?: { mode?: 'steer' | 'followUp'; inputId?: string }, + ) => this.sendToAgent(operation.id, message, options), + }) + this.operations.set(run.id, run) + return run + } + + /** A new chain, with `first` as its first run. */ + protected createChain( + fields: Pick< + AgentChain, + | 'runInputId' + | 'target' + | 'input' + | 'options' + | 'principal' + | 'parentRunId' + > & { budget?: SubagentBudget }, + first: AgentRunEntry, + ): AgentChain { + const agent = + typeof fields.target === 'string' ? fields.target : fields.target.name + const chain = { + ...fields, + agent, + thread: `${this.threadId}:${agent}:${fields.runInputId}`, + budget: fields.budget ?? this.codeBudget(), + current: first, + runs: [first], + steers: [], + followUps: [], + children: new Set(), + } + this.agentChains.set(chain.runInputId, chain) + this.chainOf.set(first.operation.id, chain) + return chain + } + + /** + * Run again an agent run with `resume: true` whose host stopped: a first + * run, or a follow-up run of the chain of `first`. The new run applies the + * same input, so its attempt counts up, and it continues the chain's saved + * transcript. + */ + protected resumeAgent(input: InputState, first: AgentInputState) { + const operation = this.agentOperation(first.input.agent) + this.markBusy() + const sender = input.principal ? { principal: input.principal } : {} + const chain = this.createChain(chainFields(first), { operation, ...sender }) + this.feed.publish( + operation.id, + customEvent(HARNESS_EVENTS.operationResumed, { + operationId: operation.id, + ...(input.operationId ? { resumedFrom: input.operationId } : {}), + }), + ) + void this.executeAgent(operation, chain, { + inputId: input.inputId, + resumed: true, + ...(input.input.op === 'agentMessage' + ? { + message: { + inputId: input.inputId, + message: input.input.message, + ...sender, + }, + } + : {}), + }) + } + + /** + * What a resumable agent run adds on a durable host: checkpoints of its + * `ctx.chat` tool phases, and `harnessContext(ctx).step` values in the + * session log, under the key of the input the run applies. + */ + protected resumable(inputId: string) { + const { chatPersistence, writer } = this + if (!chatPersistence || !writer) return undefined + const { runs } = this.persistence.stores + const key = agentKey(inputId) + const checkpoint = checkpointMiddleware({ + ...(runs ? { runs } : {}), + messages: chatPersistence.stores.messages, + hostId: this.hostId, + ...(this.lease ? { lease: this.lease } : {}), + // As in a turn, each finished tool result goes to the session log at + // once, under a key of this run. A run again puts it back. + onToolResult: ({ toolCallId, message }) => + writer.append([ + { + type: 'harness.tool.result', + toolCallId: `${key}:${toolCallId}`, + message, + }, + ]), + }) + const durableTools = this.durableTools() + return { + middleware: [checkpoint, ...(durableTools ? [durableTools] : [])], + step: this.durableBinding(writer, key).step, + } + } + + /** + * The saved transcript of a resumable agent run whose host stopped. The + * chat runs that host left running end `failed`, and the last tool batch + * gets the same repair as in a turn: the results that finished, and an + * error for a cut call that must not run twice. + */ + protected async continueAgentThread(agentThread: string, inputId: string) { + const store = this.chatPersistence?.stores.messages + if (!store) return [] + const runs = this.persistence.stores.runs + const records = (await runs?.listByThread?.(agentThread)) ?? [] + for (const run of records.filter((item) => item.status === 'running')) { + await runs?.update(run.runId, { + status: 'failed', + finishedAt: Date.now(), + error: { message: AGENT_STOPPED }, + }) + } + const newest = [...records].sort((a, b) => b.startedAt - a.startedAt)[0] + const prefix = `${agentKey(inputId)}:` + const finished = new Map( + [...(this.writer?.state.toolResults ?? [])].flatMap(([id, message]) => + id.startsWith(prefix) + ? [[id.slice(prefix.length), message] as const] + : [], + ), + ) + await repairTranscript({ + messages: store, + threadId: agentThread, + pending: newest?.checkpoint?.pendingTools ?? [], + finished, + interrupted: this.harness.durability?.interruptedToolResult, + truncated: this.harness.durability?.truncatedToolResult, + }) + return store.loadThread(agentThread) + } + + /** Run a group of agents. Every child settles before the group returns. */ + protected override async agentGroup( + options: { onFailure?: 'cancel-siblings' | 'collect' }, + body: (group: AgentGroup) => Promise, + ): Promise { + const started: Array> = [] + const cancelOthers = (failed: OperationImpl) => { + if (options.onFailure === 'collect') return + for (const operation of started) { + if (operation !== failed && !operation.isSettled()) + void operation.cancel() + } + } + const group: AgentGroup = { + run: (target: string | AnyAgent, input?: unknown) => { + const operation = this.runAgent(target, input, { wake: false }) + started.push(operation) + return Promise.resolve(operation).catch((error: unknown) => { + cancelOthers(operation) + throw error + }) + }, + runSettled: (target: string | AnyAgent, input?: unknown) => { + const operation = this.runAgent(target, input, { wake: false }) + started.push(operation) + return Promise.resolve(operation).then( + (value) => ({ ok: true as const, value }), + (error: unknown) => ({ ok: false as const, error }), + ) + }, + } as AgentGroup + try { + return await body(group) + } finally { + await Promise.allSettled( + started.map((operation) => Promise.resolve(operation)), + ) + } + } + + /** + * One run of `chain`. Then the steers it never took run as follow-ups, + * before the follow-ups that wait, and the next follow-up starts. + */ + protected async executeAgent( + operation: AgentRunImpl, + chain: AgentChain, + start: AgentRunStart, + ): Promise { + try { + await this.runAgentOnce(operation, chain, start) + } finally { + // After a log failure, another host owns the thread. + if (!this.logFailure) { + chain.followUps.unshift( + ...chain.steers + .splice(0) + .map((steer) => this.followUpOf(chain, steer)), + ) + this.nextFollowUp(chain) + } + this.checkIdle() + } + } + + /** + * Give `sent` to `chain`. A steer waits for the running run's next model + * call. A follow-up, or a steer to a run that ended, runs after the + * current run. Returns the receipt fields. + */ + protected deliver( + chain: AgentChain, + sent: AgentMessage, + mode: 'steer' | 'followUp', + ) { + const { operation: running, principal } = chain.current + // A steer joins only a run of its own sender, so no message runs with + // another person's credentials. Any other message runs as a follow-up, + // for its sender. + const isOwnSteer = + mode === 'steer' && + !running.isSettled() && + isSameSender(sent.principal, principal) + if (isOwnSteer) { + chain.steers.push(sent) + return { status: 'accepted' as const, operationId: running.id } + } + const isWaiting = !running.isSettled() || chain.followUps.length > 0 + const next = this.followUpOf(chain, sent) + chain.followUps.push(next) + this.nextFollowUp(chain) + return { + status: isWaiting ? ('queued' as const) : ('accepted' as const), + operationId: next.operation.id, + } + } + + /** `sent` as a follow-up of `chain`, with the operation of its run. */ + protected followUpOf(chain: AgentChain, sent: AgentMessage) { + const operation = this.agentOperation(chain.agent) + this.chainOf.set(operation.id, chain) + return { ...sent, operation } + } + + /** Start the next follow-up of `chain` once its current run ended. */ + protected nextFollowUp(chain: AgentChain): void { + if (!chain.current.operation.isSettled()) return + const next = chain.followUps.shift() + if (!next) return + chain.current = { + operation: next.operation, + ...(next.principal ? { principal: next.principal } : {}), + } + chain.runs.push(chain.current) + this.markBusy() + void this.executeAgent(next.operation, chain, { + inputId: next.inputId, + message: next, + }) + } + + /** + * One run of `chain`. A new first run stores its `agent` input first. + * Every run keeps its messages in the chain's thread, so a run again + * after a host stop continues them. + */ + protected async runAgentOnce( + operation: AgentRunImpl, + chain: AgentChain, + start: AgentRunStart, + ): Promise { + const { agent: name, target, input, options } = chain + const { inputId } = start + // This run is the chain's current run. + const { principal } = chain.current + if (!start.resumed && !start.message) { + await this.accept( + inputId, + { + op: 'agent', + agent: name, + input, + ...(options.wake ? { detached: true } : {}), + ...(options.resume ? { resume: true } : {}), + ...(options.attach === 'none' ? { attach: 'none' as const } : {}), + ...(chain.parentRunId ? { parentRunId: chain.parentRunId } : {}), + }, + principal, + ) + } + const agent: AnyAgent | undefined = + typeof target === 'string' ? this.agentRegistry.get(target) : target + if (!agent) { + this.reject(inputId, 'unknown_agent') + operation.fail('failed', new Error(`Unknown agent: ${name}`)) + return + } + let checkedInput: unknown = input + if (agent.inputSchema !== undefined) { + const checked = await validateWithStandardSchema( + agent.inputSchema, + input ?? {}, + ) + if (!checked.success) { + const reason = `Input validation failed for agent ${name}: ${checked.issues + .map((issue) => issue.message) + .join(', ')}` + this.reject(inputId, 'invalid_input') + operation.fail('failed', new Error(reason)) + return + } + checkedInput = checked.data + } + + operation.setStatus('running') + await this.applied(inputId, operation.id) + const runs = this.persistence.stores.runs + await runs?.createOrResume({ + runId: operation.id, + threadId: this.threadId, + startedAt: Date.now(), + kind: 'agent', + agent: name, + ...storedPrincipal(principal), + }) + operation.publish({ + type: EventType.RUN_STARTED, + runId: operation.id, + threadId: this.threadId, + timestamp: Date.now(), + }) + this.publishStarted(operation) + const child = start.child ?? { subagentRunId: createSubagentId() } + await this.indexAgent(operation, child, principal) + + let text = '' + let result: unknown + let failure: string | undefined + const { subagentRunId, prompt } = child + let stopRunLease = () => {} + let releaseLease: (() => Promise) | undefined + try { + // A host that stops lets these leases expire. Then recovery fails the run. + stopRunLease = await holdRunLease( + runs, + operation.id, + this.hostId, + this.lease, + ) + releaseLease = await this.holdLease(inputId, operation) + // A first run starts from the session thread. A run again after a + // host stop continues the chain's thread. A follow-up run continues + // the chain's thread, with its message added. The message id is its + // input id, so a run again after a host stop does not add it twice. + const { message } = start + const saved = start.resumed + ? await this.continueAgentThread(chain.thread, inputId) + : message + ? ((await this.chatPersistence?.stores.messages.loadThread( + chain.thread, + )) ?? []) + : ((await this.messages.loadThread(this.threadId)) ?? []) + const messages = + message && !saved.some((item) => item.id === message.inputId) + ? [...saved, asUserMessage(message)] + : saved + const binding = this.binding(operation, [], principal) + const durable = options.resume ? this.resumable(inputId) : undefined + const stream = runAgentStream( + withHarnessContext(agent, { + ...(durable ? { step: durable.step } : {}), + agents: this.agentStarter(chain), + }), + { + input: checkedInput, + // A prompt is one more user message at the end, as in core. + messages: + prompt === undefined + ? messages + : [ + ...messages, + { role: 'user', content: prompt } satisfies ModelMessage, + ], + threadId: chain.thread, + runId: `${operation.id}:${subagentRunId}`, + parentRunId: operation.id, + subagentRunId, + abortSignal: operation.abortController.signal, + }, + undefined, + undefined, + // ponytail: the media of a background agent is kept and published, + // but not written to history: its messages live in its chain's + // thread. Write it to the session thread if a UI must show it after + // a restart. + { + ...binding, + chatMiddleware: [ + // Every run keeps its transcript in the chain's thread. + ...(this.chatPersistence + ? [withPersistence(this.chatPersistence)] + : []), + ...(durable?.middleware ?? []), + this.agentSteering(chain, inputId, operation), + ...binding.chatMiddleware, + ], + keys: this.keysOf(principal), + budget: chain.budget, + }, + ) + for await (const chunk of stream) { + operation.publish(chunk) + if ( + chunk.type === EventType.TEXT_MESSAGE_CONTENT && + 'subagentRunId' in chunk && + chunk.subagentRunId === subagentRunId + ) { + text += chunk.delta + } + if ( + chunk.type === EventType.SUBAGENT_FINISHED && + chunk.subagentRunId === subagentRunId + ) { + result = chunk.result + } + if ( + chunk.type === EventType.SUBAGENT_ERROR && + chunk.subagentRunId === subagentRunId + ) { + failure = chunk.message + } + } + } catch (error) { + failure = error instanceof Error ? error.message : String(error) + } + + const value = result !== undefined ? result : text + // The run record and the log keep how the agent ended, so recovery + // leaves it alone. After a log failure, another host owns the thread and + // decides how the run ended, so this host writes neither. + const end = async ( + run: Pick & + Partial>, + settlement: Pick, + ) => { + if (this.logFailure) return + await runs?.update(operation.id, { ...run, finishedAt: Date.now() }) + // A refused append stops the session. The operation still ends. The + // steers that joined the run settle with it, in one append. + const joined = this.joinedInputs(operation.id, inputId) + await this.settle( + ...[inputId, ...joined].map((id) => ({ + inputId: id, + operationId: operation.id, + ...settlement, + })), + ).catch(() => {}) + } + // A recoverable close: the run stops as on a crash, and the next host + // runs it again. + if (operation.abortController.signal.reason === SHUTDOWN_REASON) { + stopRunLease() + await this.giveBack(operation, releaseLease) + return + } + if (operation.abortController.signal.aborted) { + operation.publish({ + type: EventType.RUN_FINISHED, + runId: operation.id, + threadId: this.threadId, + outcome: { type: 'cancelled' }, + timestamp: Date.now(), + } as StreamChunk) + await end({ status: 'aborted' }, { outcome: 'aborted' }) + operation.fail('cancelled', new Error('Cancelled.')) + } else if (failure !== undefined) { + operation.publish({ + type: EventType.RUN_ERROR, + message: failure, + timestamp: Date.now(), + }) + const error = { message: failure } + await end({ status: 'failed', error }, { outcome: 'failed', error }) + await this.noteAgentEnd(name, 'failed', failure, options, principal) + operation.fail('failed', new Error(failure)) + } else { + operation.publish({ + type: EventType.RUN_FINISHED, + runId: operation.id, + threadId: this.threadId, + result: value, + timestamp: Date.now(), + } as StreamChunk) + await end( + { status: 'completed', result: compactForModel(value) }, + { outcome: 'completed' }, + ) + await this.noteAgentEnd(name, 'finished', value, options, principal) + operation.finish('completed', value) + } + stopRunLease() + await releaseLease?.() + this.publishFinished(operation) + } + + /** + * The index entry of an agent run: thread `subagent:`, with + * this thread as its parent and the harness of this thread. `upsert` + * replaces the whole entry, so the fields of other writers stay. A failed + * write does not fail the run. + */ + protected async indexAgent( + operation: OperationImpl, + child: AgentChild, + principal: Principal | undefined, + ) { + const threadId = `subagent:${child.subagentRunId}` + const now = Date.now() + await this.index + .update(threadId, (entry) => ({ + ...entry, + threadId, + parentThreadId: this.threadId, + ...(child.parentToolCallId !== undefined && { + parentToolCallId: child.parentToolCallId, + }), + harness: this.harness.name, + createdAt: entry?.createdAt ?? now, + updatedAt: now, + ...storedPrincipal(entry?.principal ?? principal), + })) + .catch((error: unknown) => + this.warn( + operation, + 'harness:sessions', + `The session index was not updated. ${String(error)}`, + ), + ) + } + + /** + * Tell the main model how an agent ended: a transcript note (unless + * `attach: 'none'`), and a new chat turn when `wake` is set. The wake turn + * runs as `principal`, the user who started the agent. + */ + protected async noteAgentEnd( + name: string, + ended: 'finished' | 'failed', + result: unknown, + options: AgentStartOptions, + principal?: Principal, + ): Promise { + const note = referenceNote(name, result, ended) + if ((options.attach ?? 'reference') === 'reference') { + await this.addNote(note) + } + if (options.wake) { + void this.queueFollowUp(`Background agent ${name} ${ended}: ${note}`, { + ...(principal ? { principal } : {}), + wake: true, + }) + } + } + + /** + * Fail an agent run whose host stopped. Its events end the operation for + * views, and the main model gets a note. + */ + // ponytail: `wake` comes from the log. Without a log there is no wake, only + // the note. Keep `wake` on the run record if a host without a log needs it. + protected async failStoppedAgent(stopped: { + operationId: string + agent: string + inputId?: string + wake?: boolean + /** Who started the agent. The wake turn runs as this user. */ + principal?: Principal + /** Why it failed. Default: the host stopped. */ + error?: InputSettlement['error'] + }): Promise { + const { operationId, inputId } = stopped + const error = stopped.error ?? { message: AGENT_STOPPED } + await this.persistence.stores.runs?.update(operationId, { + status: 'failed', + finishedAt: Date.now(), + error, + }) + if (inputId) { + // The steers that joined the run settle with it. + await this.settle( + ...[inputId, ...this.joinedInputs(operationId, inputId)].map((id) => ({ + inputId: id, + operationId, + outcome: 'failed' as const, + error, + })), + ) + } + this.feed.publish(operationId, { + type: EventType.RUN_ERROR, + message: error.message, + timestamp: Date.now(), + }) + this.feed.publish( + operationId, + customEvent(HARNESS_EVENTS.operationFinished, { + operationId, + status: 'failed', + }), + ) + await this.noteAgentEnd( + stopped.agent, + 'failed', + error.message, + { wake: stopped.wake === true }, + stopped.principal, + ) + } + + /** + * After a crash: queue the logged notes that are not in the transcript, + * and note each logged job that no host runs now. + */ + protected async recoverBackground(writer: LogWriter) { + const queued = new Set(this.pendingNotes.map((message) => message.id)) + for (const [id, note] of writer.state.notes) { + if (queued.has(id)) continue + this.pendingNotes.push({ id, role: 'assistant', content: note }) + } + for (const jobId of [...writer.state.jobs]) { + if (this.jobs.has(jobId)) continue + const id = createMessageId() + const note = `Background job ${jobId} stopped when the host restarted.` + this.pendingNotes.push({ id, role: 'assistant', content: note }) + // One batch: a crash keeps both records or neither. + await writer.append([ + { type: 'harness.note', noteId: id, note }, + { type: 'harness.job', jobId, ended: true }, + ]) + } + if (!this.activeTurn && !this.reverted) await this.flushNotes() + } +} diff --git a/packages/ai-harness/src/session/base.ts b/packages/ai-harness/src/session/base.ts new file mode 100644 index 0000000000..96dda2bc7c --- /dev/null +++ b/packages/ai-harness/src/session/base.ts @@ -0,0 +1,765 @@ +import { + EventType, + SubagentBudget, + combineStrategies, + maxIterations, + untilFinishReason, +} from '@tanstack/ai' +import { credentialsFor, providerKeysFor } from '../auth' +import { AgentRegistry } from '../agents' +import { SessionFeed } from '../feed' +import { engineMessageStore } from '../log' +import { createMediaStore } from '../media' +import { LEASE } from '../resume' +import { HARNESS_EVENTS } from '../types' +import { emptyUsage } from '../usage' +import { + DEFAULT_SUBAGENT_LIMITS, + createMessageId, + customEvent, + cacheObject, + scopeOf, + notOpenMessages, + tapFeed, +} from './helpers' +import type { OperationImpl } from '../operation' +import type { + AgentLoopStrategy, + AnyChatMiddleware, + ModelMessage, + ProviderKeys, + ResolvedPromptCache, + SchemaInput, + StreamChunk, +} from '@tanstack/ai' +import type { + AIPersistence, + ChatTranscriptStores, + InboxStore, + MessageStore, + SessionIndexStore, +} from '@tanstack/ai-persistence' +import type { AuthRequiredError, CredentialsAccess } from '../auth' +import type { EventFeed } from '../feed' +import type { LogWriter, RevertState } from '../log' +import type { DurableBind } from '../durable-tool' +import type { LeaseOptions } from '../resume' +import type { MediaStore } from '../media' +import type { PluginSessionApi } from '../commands' +import type { SessionUsage } from '../usage' +import type { AgentRegistryView, AnyAgent } from '../agents' +import type { AnyHarness } from '../define' +import type { HarnessPersistence, SessionIndexWriter } from '../host' +import type { + AgentGroup, + HarnessPlugin, + MountedPlugins, + PluginAgentActions, + PluginServices, + PluginState, +} from '../plugins' +import type { + ChatTurnResult, + Cursor, + InputSettlement, + Operation, + Principal, + Receipt, + SessionEvent, + ThreadSettings, +} from '../types' +import type { + AgentRunOptions, + AgentStartOptions, + AgentChild, + AgentHandles, + SessionInspection, + SessionDependencies, + QueuedTurn, + WaitingSteer, + InterruptedTurn, + AgentRunImpl, + AgentChain, +} from './types' + +/** + * The fields of a {@link HarnessSession}, and the helpers every layer uses. The + * session class is split into layers, one per file: base, inputs, state, + * requests, turn-parts, agents, turn, harness-session. + */ +export abstract class SessionBase { + readonly threadId: string + /** The log of this session. Sessions with the same log id share one log. */ + readonly logId: string + readonly agents: AgentHandles + /** The agents this session can run, for discovery. */ + readonly registry: AgentRegistryView + + protected readonly harness: THarness + protected readonly persistence: HarnessPersistence + protected readonly inbox: InboxStore + protected readonly principal: Principal | undefined + /** The session log writer of a durable host. It is also the feed. */ + protected writer: LogWriter | undefined + protected feed: EventFeed + protected readonly onEvent: (event: StreamChunk) => void + /** The transcript: `stores.messages`, or a view of the log. */ + protected messages: MessageStore + /** + * What `withPersistence` gets: the chat stores, with `messages`. With + * `sessions`, each child the model starts gets an index entry. + */ + protected chatPersistence: + | AIPersistence + | undefined + /** The message store the chat engine saves through, on a durable host. */ + protected engine: ReturnType | undefined + /** 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() + protected readonly agentRegistry = new AgentRegistry() + protected readonly operations = new Map>() + /** 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. */ + protected readonly chainOf = new Map() + protected readonly queue: Array = [] + protected readonly steerQueue: Array = [] + /** + * Waiting steers with an abort request. An id lands here before the abort + * append, so a join that runs during that append skips the steer. + */ + protected readonly abortedSteers = new Set() + /** The ids of the waiting steers that a join took. A cancel of one is refused. */ + protected readonly joining = new Set() + /** Notes that wait for the transcript. A durable host logged each one. */ + protected readonly pendingNotes: Array = [] + /** Background jobs that run on this host. Recovery skips them. */ + protected readonly jobs = new Set() + /** Running tool calls that support `detach`: id to the move. */ + protected readonly detachable = new Map void>() + protected activeTurn: OperationImpl | undefined + /** Renews this host's claim on the thread while the thread has work. */ + protected claimTimer: ReturnType | undefined + /** The claim write of the last `markBusy`. An input waits for it. */ + protected claimWrite: Promise = Promise.resolve() + /** Called each time the thread goes idle. See `onIdle`. */ + protected readonly idleListeners = new Set<() => void>() + /** Work inputs that are being stored and are not queued yet. */ + protected inputsInFlight = 0 + /** Who sent the input of the running turn. */ + protected turnPrincipal: Principal | undefined + /** + * The last turn stopped for these interrupts. `routed`: a routed turn + * stopped. The resolve continues its saved plan. `stores.metadata` keeps a + * copy, so a resolve after a restart continues the turn too. + */ + protected interrupted: InterruptedTurn | undefined + /** The usage totals of a host without a log. A durable host folds the log. */ + protected usageTotals: SessionUsage = emptyUsage() + /** Saves of `usageTotals` to `stores.metadata`, one after the other. */ + protected usageSaved: Promise = Promise.resolve() + protected activeResume = false + protected plugins: ReadonlyArray = [] + protected sessionPlugins: MountedPlugins | undefined + /** The last `reload()`. The calls run one at a time. */ + protected reloads: Promise = Promise.resolve() + /** + * The reloads and reverts that run. No new turn starts while one runs: a + * reload waits for the running turn, and a revert changes the files. + */ + protected holdTurns = 0 + /** Called when the running turn ends. A reload waits for it. */ + protected readonly turnEnds = new Set<() => void>() + protected closing: Promise | undefined + /** Aborts on close, not on a reload. See {@link SessionSignal}. */ + protected readonly lifetime = new AbortController() + /** The last `recover()`. The calls run one at a time. */ + protected recovery: Promise = Promise.resolve() + protected readonly onClose: () => void + protected checkpoint: AnyChatMiddleware | undefined + /** Gives each durable tool call its steps in the log. One per session. */ + protected bindTool: DurableBind | undefined + protected readonly hostId: string + protected readonly index: SessionIndexWriter + protected readonly lease: LeaseOptions | undefined + protected readonly listeners = new Map< + string, + Set<(value: unknown) => void> + >() + protected readonly configValues = new Map() + /** The stored settings of this thread: `session.configure`. */ + protected threadSettings: ThreadSettings = {} + protected readonly questions = new Map< + string, + { + message: string + schema: SchemaInput | undefined + secret: boolean + url: string | undefined + /** The operation whose events carry the question and its answer. */ + operationId: string + resolve: (value: unknown) => void + reject: (error: unknown) => void + } + >() + protected readonly stateDoc: Record = {} + /** Reads each plugin's saved state, so the first snapshot has it. */ + protected readonly stateLoaders = new Map Promise>() + protected readonly localState = new Map() + /** The revert that stands. See `revert`. */ + protected reverted: RevertState | undefined + /** The credentials of the running turn's sender, else of the session's principal. */ + protected readonly credentialAccess: CredentialsAccess + /** + * The connectors whose sign-in the user cancelled. Their `require` with + * `wait` fails the tool, as without `wait`, until the next new turn. + */ + protected readonly declinedSignIns = new Set() + /** + * Model provider keys of the same principal as `credentialAccess`: saved + * with `/connect `, else the env var. + */ + protected readonly keys: ProviderKeys + /** + * The credentials of one principal: the user who runs a command, or an + * agent that outlives a turn. A save there can answer that user's sign-ins. + */ + protected readonly credentialsOf: ( + principal: Principal | undefined, + ) => CredentialsAccess + /** The provider keys of one principal, for an agent that outlives a turn. */ + protected readonly keysOf: (principal: Principal | undefined) => ProviderKeys + protected readonly services: PluginServices + protected readonly media: SessionDependencies['media'] + protected readonly mediaStore: MediaStore + /** The prompt cache of every chat turn of this session. */ + protected readonly promptCache: ResolvedPromptCache + + constructor(deps: SessionDependencies) { + this.harness = deps.harness as THarness + this.threadId = deps.threadId + this.logId = deps.logId + // The session value wins over the harness value, field by field. + const own = cacheObject(deps.promptCache) + const shared = cacheObject(this.harness.promptCache) + this.promptCache = { + retention: own?.retention ?? shared?.retention ?? 'short', + key: own?.key ?? shared?.key ?? deps.threadId, + } + this.persistence = deps.persistence + this.inbox = deps.inbox + this.media = deps.media + this.mediaStore = createMediaStore({ + persistence: deps.media, + threadId: deps.threadId, + options: this.harness.media, + }) + this.principal = deps.principal + this.onEvent = deps.onEvent + this.feed = tapFeed(new SessionFeed(), deps.onEvent) + this.onClose = deps.onClose + this.hostId = deps.hostId + this.index = deps.index + this.lease = deps.lease + this.log = deps.log + // A durable host gets its log view in `open()`. + this.messages = deps.persistence.stores.messages ?? notOpenMessages + const onMissing = (error: AuthRequiredError) => + this.feed.publish( + 'session', + customEvent(HARNESS_EVENTS.authRequired, { + connector: error.connector, + ...(error.url ? { url: error.url } : {}), + }), + ) + // The scope is read at each call, so it follows the running turn. + this.credentialAccess = credentialsFor( + deps.credentials, + () => scopeOf(deps.threadId, this.sender()), + onMissing, + { + // A tool of a running turn can wait, unless the user cancelled the + // sign-in for this turn. + canWait: (id) => + this.activeTurn !== undefined && !this.declinedSignIns.has(id), + onSet: () => this.resumeSignIns(), + }, + ) + this.keys = providerKeysFor(this.credentialAccess, onMissing) + this.credentialsOf = (principal) => + credentialsFor( + deps.credentials, + scopeOf(deps.threadId, principal), + onMissing, + { onSet: () => this.resumeSignIns() }, + ) + this.keysOf = (principal) => + providerKeysFor(this.credentialsOf(principal), onMissing) + this.services = { + emit: (plugin, name, value) => this.emitPluginEvent(plugin, name, value), + on: (name, handler) => { + let set = this.listeners.get(name) + if (!set) { + set = new Set() + this.listeners.set(name, set) + } + set.add(handler) + return () => set.delete(handler) + }, + config: { get: (key) => this.configValue(key) }, + state: (plugin, initial) => this.pluginState(plugin, initial), + credentials: this.credentialAccess, + keys: this.keys, + session: this.pluginApi(), + sessionWith: (metadataSchema) => + this.pluginApi(undefined, metadataSchema), + agents: { + run: ((target: string | AnyAgent, input?: unknown) => + this.runAgent(target, input, { + wake: false, + })) as PluginAgentActions['run'], + start: (( + target: string | AnyAgent, + input?: unknown, + options?: AgentStartOptions, + ) => + this.runAgent( + target, + input, + options ?? {}, + )) as PluginAgentActions['start'], + group: (options, body) => this.agentGroup(options, body), + }, + commandsChanged: () => + this.feed.publish( + 'session', + customEvent(HARNESS_EVENTS.commandsChanged, {}), + ), + } + this.registry = this.agentRegistry + for (const agent of this.harness.agents ?? []) { + this.agentRegistry.add(agent, 'the harness') + } + for (const agent of this.harness.subagents?.agents ?? []) { + this.agentRegistry.add(agent, 'the harness') + } + this.agents = new Proxy({} as AgentHandles, { + get: (_target, name) => { + if (typeof name !== 'string') return undefined + return { + run: (input?: unknown, options?: AgentRunOptions) => + this.runAgent(name, input, { ...options, wake: false }), + start: (input?: unknown, options?: AgentStartOptions) => + this.runAgent(name, input, options ?? {}), + } + }, + }) + } + + /** An operation of this session by id, running or settled. */ + operation(id: string): Operation | undefined { + return this.operations.get(id) + } + + // =========================== + // Events and state + // =========================== + + /** The ordered events of every operation, from `from` (exclusive). */ + events(options?: { + from?: Cursor + signal?: AbortSignal + }): AsyncIterable { + return this.feed.read(options ?? {}) + } + + /** + * The token usage of this thread: every model call of its turns, their + * subagents, and its agent runs, in total, by `provider/model`, and by + * sender. `cost` is the sum of what the providers reported. A durable host + * keeps the totals in the log, a host with `stores.metadata` keeps them + * there, and other hosts keep them in memory. + * + * @example + * ```ts + * const { total, bySender } = session.usage() + * console.log(total.totalTokens, bySender['user-1']?.cost) + * ``` + */ + usage(): SessionUsage { + return structuredClone(this.writer?.state.usage ?? this.usageTotals) + } + + /** The resolved plugin plan: order, owners, and extension contributors. */ + inspect(): SessionInspection { + const mounted = this.sessionPlugins + const extensionPoints: Record> = {} + for (const [point, items] of mounted?.extensions ?? []) { + extensionPoints[point] = [...new Set(items.map((item) => item.owner))] + } + return { + plugins: mounted?.owners.plugins ?? [], + tools: mounted?.owners.tools ?? [], + prompts: mounted?.owners.prompts ?? [], + commands: [...(mounted?.commands ?? [])].map(([name, entry]) => ({ + name, + owner: entry.owner, + })), + config: [...(mounted?.config ?? [])].map(([key, entry]) => ({ + key, + owner: entry.owner, + })), + extensionPoints, + agents: this.agentRegistry.list().map((agent) => agent.name), + } + } + + protected configValue(key: string): unknown { + if (this.configValues.has(key)) return this.configValues.get(key) + return this.sessionPlugins?.config.get(key)?.option.default + } + + /** Put each plugin's saved state in the snapshot before the first read. */ + protected async loadPluginState() { + for (const [plugin, load] of this.stateLoaders) { + if (!(plugin in this.stateDoc)) this.stateDoc[plugin] = await load() + } + } + + /** + * Who sent the input of the running turn, else the principal that opened + * the session. Turns run one at a time, so one value is enough. + */ + protected sender(): Principal | undefined { + return this.turnPrincipal ?? this.principal + } + + protected emitPluginEvent( + plugin: string, + name: string, + value: unknown, + ): void { + this.feed.publish( + 'session', + customEvent(HARNESS_EVENTS.pluginEvent, { plugin, name, value }), + ) + for (const handler of this.listeners.get(name) ?? []) { + try { + handler(value) + } catch { + // One broken listener must not stop the others. + } + } + } + + protected pluginState(plugin: string, initial: T): PluginState { + const metadata = this.persistence.stores.metadata + const namespace = `plugin:${plugin}` + const key = this.threadId + const read = async (): Promise<{ value: T; revision: string | null }> => { + if (metadata?.getVersioned) { + const stored = await metadata.getVersioned(namespace, key) + // The store holds what this plugin wrote. + return stored + ? { value: stored.value as T, revision: stored.revision } + : { value: initial, revision: null } + } + if (metadata) { + const stored = await metadata.get(namespace, key) + return { + value: stored === null ? initial : (stored as T), + revision: null, + } + } + return { + value: this.localState.has(namespace) + ? (this.localState.get(namespace) as T) + : initial, + revision: null, + } + } + this.stateLoaders.set(plugin, async () => (await read()).value) + const publish = (value: T) => { + this.stateDoc[plugin] = value + this.feed.publish('session', { + type: EventType.STATE_SNAPSHOT, + snapshot: { plugins: { ...this.stateDoc } }, + timestamp: Date.now(), + }) + } + return { + get: async () => (await read()).value, + update: async (change) => { + for (let attempt = 0; attempt < 5; attempt += 1) { + const current = await read() + const next = change(structuredClone(current.value)) + if (metadata?.setIf) { + const written = await metadata.setIf( + namespace, + key, + next, + current.revision, + ) + if (!written.ok) continue + } else if (metadata) { + await metadata.set(namespace, key, next) + } else { + this.localState.set(namespace, next) + } + publish(next) + return next + } + throw new Error(`Plugin ${plugin}: state update conflicted 5 times.`) + }, + } + } + + /** + * The thread has work: claim it in `stores.workClaims` for this host, and + * renew the claim while the work runs, so a sweep finds the thread after a + * crash. A claim that another host holds is left alone: the log still keeps + * one writer. A closing session claims nothing. + */ + protected markBusy(): void { + const claims = this.persistence.stores.workClaims + if (!claims || this.claimTimer || this.closing) return + const ttlMs = this.lease?.ttlMs ?? LEASE.ttlMs + const claim = async () => { + try { + const isHeld = await claims.claim({ + threadId: this.threadId, + harness: this.harness.name, + ownerId: this.hostId, + until: Date.now() + ttlMs, + }) + // Renew only a claim this host holds. + if (!isHeld) this.stopClaimTimer() + } catch (error) { + // The work goes on without a claim. + this.warnClaim(error) + } + } + const timer = setInterval( + () => void claim(), + this.lease?.renewMs ?? LEASE.renewMs, + ) + // A claim timer must not keep a CLI or a test process alive. + if (typeof timer === 'object' && 'unref' in timer) timer.unref() + this.claimTimer = timer + this.claimWrite = claim() + } + + protected stopClaimTimer(): void { + if (this.claimTimer) clearInterval(this.claimTimer) + this.claimTimer = undefined + } + + /** Give the thread's claim back. A failed write is a warning. */ + protected async releaseClaim(): Promise { + try { + await this.persistence.stores.workClaims?.release( + this.threadId, + this.hostId, + ) + } catch (error) { + this.warnClaim(error) + } + } + + /** + * A claim write failed. Without a claim, a sweep cannot find this work + * after a crash, so clients get a warning. The work goes on. + */ + protected warnClaim(error: unknown): void { + this.feed.publish( + 'session', + customEvent('harness.plugin.warning', { + plugin: 'harness:work-claims', + message: error instanceof Error ? error.message : String(error), + }), + ) + } + + /** + * No input being stored, no queued or running turn, no waiting steer, and + * no running agent. + */ + protected 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, + ) + if (hasFollowUps) return false + return ( + this.inputsInFlight === 0 && + this.activeTurn === undefined && + this.queue.length === 0 && + this.steerQueue.length === 0 && + ![...this.operations.values()].some( + (operation) => operation.kind === 'agent' && !operation.isSettled(), + ) + ) + } + + /** + * Run `work`, a public entry that stores a work input and queues its work. + * The thread is busy until `work` ends, so it is not idle while the input + * is stored. + */ + protected async admitting(work: () => Promise) { + this.inputsInFlight += 1 + try { + return await work() + } finally { + this.inputsInFlight -= 1 + this.checkIdle() + } + } + + /** + * When the thread is idle, give its claim back and tell the idle + * listeners. A thread that waits for a human (an approval, a sign-in) is + * idle. + */ + protected checkIdle(): void { + if (!this.isIdle()) return + this.stopClaimTimer() + // ponytail: one release write per idle thread, also when this host + // holds no claim (a sweep may hold one for it). Track the claim if the + // writes matter. + void this.releaseClaim() + for (const listener of [...this.idleListeners]) listener() + } + + /** + * @internal Call `listener` each time the thread goes idle, and once soon + * when it is idle now. Returns a function that stops the calls. The host + * uses it to close a session that `resumePending` opened. + */ + onIdle(listener: () => void): () => void { + this.idleListeners.add(listener) + queueMicrotask(() => this.checkIdle()) + return () => { + this.idleListeners.delete(listener) + } + } + + protected warn( + operation: OperationImpl | OperationImpl, + plugin: string, + error: unknown, + ) { + operation.publish( + customEvent('harness.plugin.warning', { + plugin, + message: error instanceof Error ? error.message : String(error), + }), + ) + } + + /** + * A tree budget for a child started from code. The child counts as the + * first call, and its own children count against the same limits. + */ + protected codeBudget(): SubagentBudget { + const root = SubagentBudget.root(this.limits()) + root.reserve(0) + return root.child() + } + + // =========================== + // Agents + // =========================== + + /** The limits for children of this session, with the harness defaults. */ + /** + * chat() stops after 5 model calls by default. An agent that reads, + * searches, and calls tools needs more before it can answer. Without + * `durability.truncatedToolResult`, an answer cut at the output limit ends + * the turn. With it, the model gets that text for each cut call and is + * called again. + */ + protected loopStrategy(): AgentLoopStrategy { + const strategy = this.harness.agentLoopStrategy ?? maxIterations(50) + return this.harness.durability?.truncatedToolResult === undefined + ? combineStrategies([strategy, untilFinishReason(['length'])]) + : strategy + } + + protected limits() { + return this.harness.subagents?.limits ?? DEFAULT_SUBAGENT_LIMITS + } + + /** + * Add an assistant note to the transcript: at once when no turn runs, else + * before the next turn starts. + */ + protected async addNote(note: string) { + const id = createMessageId() + this.pendingNotes.push({ id, role: 'assistant', content: note }) + // A durable host logs the note first, so a crash does not lose it. + await this.writer?.append([{ type: 'harness.note', noteId: id, note }]) + // While a revert stands, a note waits: the next turn or unrevert adds it. + if (!this.activeTurn && !this.reverted) await this.flushNotes() + } + + /** Write queued notes to the transcript while no turn is writing it. */ + protected async flushNotes(): Promise { + const messages = this.messages + if (this.pendingNotes.length === 0) return + const notes = this.pendingNotes.splice(0) + const history = await messages.loadThread(this.threadId) + await messages.saveThread(this.threadId, [...history, ...notes]) + } + + /** + * Track a background job of this host. A durable host also logs its start + * and end, so recovery can note a job that a crash stopped. + */ + protected logJob(jobId: string, ended: boolean) { + if (ended) this.jobs.delete(jobId) + else this.jobs.add(jobId) + // No one waits for a job record. A failed log write fails the session. + this.writer + ?.append([{ type: 'harness.job', jobId, ended }]) + .catch(() => undefined) + } + + // Implemented by a later layer of the class. + protected abstract drain(): void + protected abstract agentGroup( + options: { onFailure?: 'cancel-siblings' | 'collect' }, + body: (group: AgentGroup) => Promise, + ): Promise + protected abstract pluginApi( + principal?: Principal, + metadataSchema?: SchemaInput, + ): PluginSessionApi + protected abstract resumeSignIns(): Promise + protected abstract runAgent( + target: string | AnyAgent, + input: unknown, + options: AgentStartOptions, + parent?: AgentChain, + child?: AgentChild, + ): AgentRunImpl +} diff --git a/packages/ai-harness/src/session/harness-session.ts b/packages/ai-harness/src/session/harness-session.ts new file mode 100644 index 0000000000..f98060f633 --- /dev/null +++ b/packages/ai-harness/src/session/harness-session.ts @@ -0,0 +1,810 @@ +import { EventType } from '@tanstack/ai' +import { REDACTED_THINKING_ID_PREFIX } from '@tanstack/ai/adapter-internals' +import { LogConflictError } from '@tanstack/ai-persistence' +import { engineMessageStore, revertOf, sessionMessageStore } from '../log' +import { mountPlugins } from '../plugins' +import { + CUT_OFF_NOTE, + checkpointMiddleware, + findCrashedRuns, + repairSteps, + repairTranscript, +} from '../resume' +import { HARNESS_EVENTS } from '../types' +import { + chainInputOf, + createMessageId, + customEvent, + errorText, + inputKey, + CHAT_OPS, + CONTROL_OPS, + hasFinalAnswer, + tapFeed, +} from './helpers' +import { SessionTurn } from './turn' +import type { ModelMessage } from '@tanstack/ai' +import type { InboxEntry } from '@tanstack/ai-persistence' +import type { InputState, LogWriter } from '../log' +import type { AnyHarness } from '../define' +import type { RecoverContext, RecoverDecision } from '../turn' +import type { HarnessInput, InputSettlement } 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()`. + */ +export class HarnessSession< + THarness extends AnyHarness = AnyHarness, +> extends SessionTurn { + /** @internal Mount session plugins and replay inputs left in the inbox. */ + async open(): Promise { + // First: plugins publish events while they mount. + await this.openLog() + try { + await this.mountSessionPlugins() + await this.loadSettings() + this.reverted = this.writer + ? this.writer.state.revert + : revertOf( + await this.persistence.stores.metadata?.get( + 'harness:revert', + this.threadId, + ), + ) + await this.loadInterrupted() + // Before recovery: a turn that recovery runs updates this entry. + await this.openEntry() + if (this.writer) { + await this.recoverFromLog(this.writer) + return + } + await this.loadUsage() + await this.recoverCrashedTurn() + await this.recoverInbox() + } catch (error) { + // A session that fails to open gives its log view back, so the thread + // can open again. + this.feed.close() + throw error + } + } + + /** Read `harness.plugins()` and set up its session plugins. */ + protected async mountSessionPlugins(): Promise { + this.plugins = this.harness.plugins?.() ?? [] + this.sessionPlugins = await mountPlugins( + this.plugins.filter( + (plugin) => (plugin.lifetime ?? 'session') === 'session', + ), + { + threadId: this.threadId, + registry: this.agentRegistry, + harnessTools: this.harness.tools ?? [], + harnessProvides: (this.harness.middleware ?? []).flatMap( + (middleware) => middleware.provides ?? [], + ), + services: this.services, + inherited: this.storeCapabilities(), + }, + ) + await this.loadConfig() + await this.loadPluginState() + } + + /** + * Tear the plugins down and set them up again, with the list that + * `harness.plugins()` gives now. It waits for the running turn, and no new + * turn starts until it ends. The transcript, the log, the thread settings, + * and the inbox stay. Clients get a `harness.reloaded` event. + * + * When the new setup throws, the session keeps working with no plugins, + * and the promise rejects with the error. Call `reload()` again after a + * fix. + * + * @example + * ```ts + * await session.reload() + * ``` + */ + reload(): Promise { + const run = this.reloads.then(() => this.runReload()) + this.reloads = run.catch(() => {}) + return run + } + + protected async runReload(): Promise { + if (this.closing) return + this.holdTurns++ + try { + while (this.activeTurn) { + await new Promise((resolve) => this.turnEnds.add(resolve)) + } + const old = this.sessionPlugins + // No plugins until the new setup works, so a failed setup leaves none. + this.sessionPlugins = undefined + this.plugins = [] + this.agentRegistry.deletePluginAgents() + await old?.dispose() + await this.mountSessionPlugins() + this.feed.publish('session', customEvent(HARNESS_EVENTS.reloaded, {})) + } catch (error) { + this.sessionPlugins = undefined + this.plugins = [] + this.agentRegistry.deletePluginAgents() + this.feed.publish( + 'session', + customEvent(HARNESS_EVENTS.reloaded, { error: errorText(error) }), + ) + throw error + } finally { + this.holdTurns-- + this.drain() + } + } + + /** + * Recover the work in the log again, as `open()` does. `open()` skips a + * turn while another host holds its lease. Call this after that lease + * expires, and the turn runs here. The work that this session runs or + * queues stays as it is. A session without `stores.log` does nothing. + * + * @example + * ```ts + * await session.recover() + * ``` + */ + recover(): Promise { + const { writer } = this + if (!writer || this.closing) return Promise.resolve() + const run = this.recovery.then(async () => { + // First the records that other hosts wrote since this session read. + await writer.catchUp() + if (this.logFailure) throw this.logFailure + if (this.closing) return + await this.recoverFromLog(writer) + }) + this.recovery = run.catch(() => {}) + return run + } + + /** + * A durable host gives the session its view of the log, and the session + * writes through it: the log is the event feed and the transcript. Then + * both modes build the stores that `withPersistence` and the checkpoints + * get. + */ + protected async openLog(): Promise { + const { stores } = this.persistence + if (this.log) { + const { store, project } = this.log + this.writer = await this.log.open(this.threadId, (error) => + this.stopOnLogFailure(error), + ) + this.feed = tapFeed(this.writer, this.onEvent) + const view = { + writer: this.writer, + store, + ...(project ? { project } : {}), + } + this.messages = sessionMessageStore(view) + this.engine = engineMessageStore(view) + } + const engineMessages = this.engine ?? this.messages + this.chatPersistence = { + stores: { + messages: engineMessages, + ...(stores.runs ? { runs: stores.runs } : {}), + ...(stores.interrupts ? { interrupts: stores.interrupts } : {}), + ...(stores.metadata ? { metadata: stores.metadata } : {}), + ...(stores.sessions ? { sessions: stores.sessions } : {}), + }, + } + const { writer } = this + this.checkpoint = checkpointMiddleware({ + ...(stores.runs ? { runs: stores.runs } : {}), + messages: engineMessages, + hostId: this.hostId, + ...(this.lease ? { lease: this.lease } : {}), + // A durable host keeps each finished tool result, so a crash later in + // the batch does not lose it. + ...(writer + ? { + onToolStart: ({ toolCallId, name, replay }) => + writer.append([ + { type: 'harness.tool.started', toolCallId, name, replay }, + ]), + onToolResult: ({ toolCallId, message }) => + writer.append([ + { type: 'harness.tool.result', toolCallId, message }, + ]), + } + : {}), + }) + if (writer) { + this.bindTool = (toolCallId) => this.durableBinding(writer, toolCallId) + } + } + + /** + * The log refused a write: another host wrote to this log, or the store + * failed. The state of this session is not known any more, so it stops. + * The next `host.open` folds the log again. + */ + protected stopOnLogFailure(error: unknown): void { + if (this.logFailure) return + this.logFailure = new Error( + error instanceof LogConflictError + ? 'Another host wrote to this log, so this session stopped. Open the thread again.' + : `The session log failed, so this session stopped: ${error instanceof Error ? error.message : String(error)}`, + ) + for (const operation of this.operations.values()) { + if (!operation.isSettled()) + operation.abortController.abort(this.logFailure) + } + void this.close() + } + + /** + * Continue the newest chat turn that a crashed host left running. Older + * crashed turns are marked failed, and so are crashed agent runs. + */ + protected async recoverCrashedTurn(): Promise { + const runs = await findCrashedRuns( + this.persistence.stores.runs, + this.threadId, + ) + for (const run of runs.filter((record) => record.kind === 'agent')) { + await this.failStoppedAgent({ + operationId: run.runId, + agent: run.agent ?? 'agent', + }) + } + const crashed = runs.filter((record) => record.kind !== 'agent') + const newest = crashed.sort((a, b) => b.startedAt - a.startedAt)[0] + for (const record of crashed) { + await this.persistence.stores.runs?.update(record.runId, { + status: 'failed', + finishedAt: Date.now(), + error: { + message: + record === newest + ? 'The host stopped. The session continued this turn in a new run.' + : 'The host stopped during this turn.', + }, + }) + } + if (!newest || this.interrupted) return + await repairTranscript({ + messages: this.messages, + threadId: newest.threadId, + pending: newest.checkpoint?.pendingTools ?? [], + interrupted: this.harness.durability?.interruptedToolResult, + truncated: this.harness.durability?.truncatedToolResult, + }) + const operation = this.createTurnOperation() + this.feed.publish( + operation.id, + customEvent(HARNESS_EVENTS.operationResumed, { + operationId: operation.id, + resumedFrom: newest.runId, + }), + ) + this.enqueueTurn({ operation }) + } + + /** Re-run turns that were accepted but never applied before a restart. */ + protected async recoverInbox(): Promise { + const pending: Array = await this.inbox.listPending( + this.threadId, + ) + for (const entry of pending) { + const input = entry.input as HarnessInput + if ( + input.op === 'prompt' || + input.op === 'followUp' || + input.op === 'steer' + ) { + const operation = this.createTurnOperation() + this.admitted.set(entry.inputId, inputKey(input, entry.principal)) + this.bindTurn(entry.inputId, operation) + this.enqueueTurn({ + operation, + message: input.message, + inputId: entry.inputId, + principal: entry.principal, + context: input.context, + ...(input.op === 'followUp' && input.wake ? { wake: true } : {}), + }) + } else if (input.op === 'reset') { + this.queueReset(entry.inputId, input.note, entry.principal, 'end') + } else { + this.reject(entry.inputId, 'expired_on_restart') + } + } + } + + /** The recovery decision for `input`: the hook's, or `decision`. */ + protected async recoverDecision( + writer: LogWriter, + input: InputState, + decision: RecoverDecision, + interruptedTools: RecoverContext['interruptedTools'] = [], + ): Promise { + const recover = this.harness.durability?.recover + if (!recover) return decision + const answer = await recover({ + session: this, + input: { + inputId: input.inputId, + input: input.input, + attempt: input.attempt, + ...(input.timeoutAt !== undefined + ? { timeoutAt: input.timeoutAt } + : {}), + abortRequested: input.abortRequested, + ...(input.operationId ? { operationId: input.operationId } : {}), + retries: input.retries ?? 0, + }, + messages: [...writer.state.messages], + interruptedTools, + decision, + }) + return answer ?? decision + } + + /** + * The part of a cut answer that the log has: the answer text and the + * signed thinking blocks that operation `operationId` streamed after the + * last transcript record. A thinking block with no signature, tool call + * arguments, and agent events are not in it. + */ + protected async cutOffAnswer(writer: LogWriter, operationId: string) { + const from = writer.state.transcriptSeq + if (from === undefined) return { text: '', thinking: [] } + let text = '' + // By reasoning message id, in stream order. + const blocks = new Map() + const block = (id: string) => { + const found = blocks.get(id) ?? { content: '' } + blocks.set(id, found) + return found + } + for await (const { event } of writer.read({ + from: String(from), + filter: (entry) => entry.operationId === operationId, + until: () => true, + })) { + if ('subagentRunId' in event && event.subagentRunId) continue + if (event.type === EventType.TEXT_MESSAGE_CONTENT) text += event.delta + if (event.type === EventType.REASONING_MESSAGE_CONTENT) { + block(event.messageId).content += event.delta + } + if ( + event.type === EventType.REASONING_ENCRYPTED_VALUE && + event.subtype === 'message' + ) { + block(event.entityId).signature = event.encryptedValue + } + } + // As the transcript keeps thinking: a redacted block has no text. + const thinking = [...blocks].flatMap(([id, { content, signature }]) => + !signature + ? [] + : id.startsWith(REDACTED_THINKING_ID_PREFIX) + ? [{ content: '', signature, redacted: true }] + : [{ content, signature }], + ) + return { text, thinking } + } + + /** + * A `cancelInput` or `setDelivery` whose own record is in the log, but not + * its change: a crash cut it. Apply it now, before any input runs. When + * its input does not wait any more, reject it with `not_waiting`. Other + * ops are not changed. + */ + protected async recoverControl(writer: LogWriter, logged: InputState) { + const { input } = logged + if (input.op !== 'cancelInput' && input.op !== 'setDelivery') return + const target = writer.state.inputs.get(input.inputId) + const isWaiting = target?.status === 'pending' && 'message' in target.input + if (!isWaiting) { + this.reject(logged.inputId, 'not_waiting') + return + } + await this.applyControl( + logged.inputId, + input.op === 'cancelInput' + ? { type: 'harness.input.abort', inputId: input.inputId } + : { + type: 'harness.input.delivery', + inputId: input.inputId, + delivery: input.delivery, + }, + ) + } + + /** + * Recover a durable session from its log, input by input, in admission + * order. A chat turn whose host stopped (its run lease expired) settles + * `completed` when the log already has its final answer. Else it settles + * `aborted` when an abort was asked, settles `failed` when no attempt or no + * time is left, and else runs again as the next attempt. An agent run (a + * first run or a follow-up run) whose host stopped runs again with + * `resume: true`, else settles `failed` with the steers that joined it. A + * message that waited for such a run goes to the run again, else settles + * `aborted`. An input that never ran runs now. One that a `setDelivery` + * sent to `steer` joins the turn that runs, when one does. A `cancelInput` + * or `setDelivery` that a crash cut before it was applied applies first. + * Other inputs are rejected with `expired_on_restart`. An input that this + * session stored or took already is skipped, so `recover()` can run this + * again. + */ + protected async recoverFromLog(writer: LogWriter): Promise { + const runs = this.persistence.stores.runs + const maxAttempts = this.harness.durability?.maxAttempts ?? 10 + // True for an input that this session stored or took. Else the session + // takes it now, so `recover()` does not take it again. + const isTaken = (input: InputState) => { + if (this.admitted.has(input.inputId)) return true + this.admitted.set(input.inputId, inputKey(input.input, input.principal)) + return false + } + await this.recoverBackground(writer) + const logged = [...writer.state.inputs.values()] + for (const input of logged) { + const isControl = + input.status === 'pending' && CONTROL_OPS.has(input.input.op) + if (isControl && !isTaken(input)) await this.recoverControl(writer, input) + } + // Read again: the pass above changed the inputs it applied to. + const inputs = [...writer.state.inputs.values()] + for (const input of inputs) { + // A reset that did not settle applies now. Its marker goes in once. + if (input.input.op === 'reset') { + const isOpen = + (input.status === 'pending' || input.status === 'applied') && + !isTaken(input) + if (isOpen && input.abortRequested) { + await this.settle({ inputId: input.inputId, outcome: 'aborted' }) + } else if (isOpen) { + this.queueReset( + input.inputId, + input.input.note, + input.principal, + 'end', + ) + } + continue + } + const isChat = CHAT_OPS.has(input.input.op) + if (input.status === 'pending') { + // The pass above handled a control input. + if (CONTROL_OPS.has(input.input.op) || isTaken(input)) continue + // A message to an agent run goes to the run that recovery continues. + // That run's input is older, so its chain exists here. + if (input.input.op === 'agentMessage') { + const chain = this.agentChains.get(input.input.run ?? '') + if (chain) { + this.deliver( + chain, + { + inputId: input.inputId, + message: input.input.message, + ...(input.principal ? { principal: input.principal } : {}), + }, + input.input.mode ?? 'steer', + ) + } else { + // ponytail: a follow-up that waited for a run that had ended when + // the host stopped settles aborted too. Start it from the chain + // in the log if an app needs it. + await this.settle({ inputId: input.inputId, outcome: 'aborted' }) + } + continue + } + if ( + !isChat || + !('message' in input.input || input.input.op === 'continue') + ) { + this.reject(input.inputId, 'expired_on_restart') + continue + } + // An input that was cancelled before it ran settles aborted. + const decision = await this.recoverDecision( + writer, + input, + input.abortRequested + ? { action: 'settle', outcome: 'aborted' } + : { action: 'run' }, + ) + if (decision.action === 'settle') { + await this.settle({ + inputId: input.inputId, + outcome: decision.outcome, + ...(decision.error ? { error: decision.error } : {}), + }) + continue + } + if (input.input.op === 'continue') { + // It runs from the transcript, as `continue()` does. + const operation = this.createTurnOperation() + this.bindTurn(input.inputId, operation) + this.enqueueTurn({ + operation, + inputId: input.inputId, + principal: input.principal, + context: input.input.context, + overrides: decision.overrides, + isContinue: true, + }) + continue + } + const steer: WaitingSteer = { + inputId: input.inputId, + message: input.input.message, + principal: input.principal, + context: input.input.context, + overrides: decision.overrides, + ...(input.input.op === 'followUp' && input.input.wake + ? { wake: true } + : {}), + } + // Else it runs as its own turn, as `steer()` does when no turn runs. + if (input.delivery === 'steer' && this.activeTurn) { + this.steerQueue.push(steer) + continue + } + this.enqueueTurn(this.steerTurn(steer)) + continue + } + // An agent run: a first run, or a follow-up run of its chain. + const first = chainInputOf(writer.state.inputs, input) + if (input.status !== 'applied' || !(isChat || first)) continue + const leases = this.persistence.stores.leases + const run = input.operationId ? await runs?.get(input.operationId) : null + const now = Date.now() + // A host lease store decides when there is one. Else the run lease does. + const isAlive = + leases && input.operationId + ? await leases.isAlive({ + threadId: this.threadId, + inputId: input.inputId, + operationId: input.operationId, + attempt: input.attempt, + }) + : run?.status === 'running' && + run.leaseExpiresAt !== undefined && + run.leaseExpiresAt >= now + // Another host still drives this turn. + if (isAlive || isTaken(input)) continue + if (first) { + // A run that ended before its settle record landed did not stop. + if (!run || run.status === 'running') { + const isResumable = first.input.resume === true + if (isResumable && input.attempt < maxAttempts) { + if (input.operationId) { + await runs?.update(input.operationId, { + status: 'failed', + finishedAt: now, + error: { + message: + 'The host stopped. The session continued this agent run in a new run.', + }, + }) + } + this.resumeAgent(input, first) + continue + } + await this.failStoppedAgent({ + operationId: input.operationId ?? '', + agent: first.input.agent, + inputId: input.inputId, + wake: first.input.detached, + ...(input.principal ? { principal: input.principal } : {}), + ...(isResumable + ? { + error: { + message: `The agent run stopped the host ${input.attempt} times.`, + code: 'attempts_exhausted', + }, + } + : {}), + }) + } + continue + } + if (run?.status === 'running') { + await runs?.update(run.runId, { + status: 'failed', + finishedAt: now, + error: { message: 'The host stopped during this turn.' }, + }) + } + const operationId = input.operationId + // The inputs that joined this turn settle with it, in one append. + const joined = this.joinedInputs(operationId ?? '', input.inputId) + const ended = (settlement: Omit) => + this.settle( + ...[input.inputId, ...joined].map((inputId) => ({ + inputId, + ...(operationId ? { operationId } : {}), + ...settlement, + })), + ) + // A final answer in the log comes first: the turn ended. + const fallback: RecoverDecision = hasFinalAnswer( + writer.state.messages, + input.appliedAt ?? 0, + ) + ? { action: 'settle', outcome: 'completed' } + : input.abortRequested + ? { action: 'settle', outcome: 'aborted' } + : input.attempt >= maxAttempts + ? { + action: 'settle', + outcome: 'failed', + error: { + message: `The input stopped the host ${input.attempt} times.`, + code: 'attempts_exhausted', + }, + } + : input.timeoutAt !== undefined && now >= input.timeoutAt + ? { + action: 'settle', + outcome: 'failed', + error: { + message: 'The input passed its time limit.', + code: 'timeout', + }, + } + : { action: 'run' } + const pending = [ + ...(run?.checkpoint?.pendingTools ?? []), + ...[...writer.state.started] + .filter(([toolCallId]) => !writer.state.toolResults.has(toolCallId)) + .map(([toolCallId, tool]) => ({ toolCallId, ...tool })), + ] + // Before the decision, so the hook has it when it settles the input. A + // resolve repairs nothing. + const interruptedTools = + input.input.op === 'resolve' + ? [] + : repairSteps( + writer.state.messages, + pending, + writer.state.toolResults, + ).filter((step) => 'reason' in step) + const decision = await this.recoverDecision( + writer, + input, + fallback, + interruptedTools, + ) + if (decision.action === 'settle') { + await ended({ + outcome: decision.outcome, + ...(decision.error ? { error: decision.error } : {}), + }) + continue + } + if (input.input.op === 'resolve') { + const store = this.persistence.stores.interrupts + const records = store + ? await Promise.all( + input.input.resume.map((item) => store.get(item.interruptId)), + ) + : [] + const consumed = + records.length > 0 && + records.every( + (record) => + record !== null && + record.threadId === this.threadId && + record.status !== 'pending', + ) + const newPhase = + this.interrupted && this.interrupted.runId === operationId + if ( + consumed || + newPhase || + (!this.interrupted && run?.status === 'completed') + ) { + await ended({ + outcome: this.interrupted ? 'interrupted' : 'completed', + }) + continue + } + if (!this.interrupted) { + await ended({ + outcome: 'failed', + error: { message: 'The interrupted turn is no longer available.' }, + }) + continue + } + const operation = this.createTurnOperation() + this.bindTurn(input.inputId, operation) + this.feed.publish( + operation.id, + customEvent(HARNESS_EVENTS.operationResumed, { + operationId: operation.id, + ...(operationId ? { resumedFrom: operationId } : {}), + }), + ) + // The same turn as a live `resolve()`: it answers the stopped turn, + // as the input's sender, with the context of the stopped turn. + this.enqueueTurn({ + operation, + inputId: input.inputId, + resume: input.input.resume, + answers: this.interrupted, + principal: input.principal, + context: this.interrupted.context, + overrides: decision.overrides, + }) + continue + } + // Before the repair, which can write a transcript record. + const cutOff = this.harness.durability?.continueCutOff + const partial = + cutOff && operationId + ? await this.cutOffAnswer(writer, operationId) + : undefined + await repairTranscript({ + messages: this.messages, + threadId: this.threadId, + pending, + finished: writer.state.toolResults, + interrupted: this.harness.durability?.interruptedToolResult, + truncated: this.harness.durability?.truncatedToolResult, + }) + if (partial && (partial.text !== '' || partial.thinking.length > 0)) { + const history = await this.messages.loadThread(this.threadId) + const notes = [ + (typeof cutOff === 'object' && cutOff.note) || CUT_OFF_NOTE, + ].flat() + // One append, so the transcript never ends with the cut answer. + await this.messages.saveThread(this.threadId, [ + ...history, + { + id: createMessageId(), + role: 'assistant', + content: partial.text || null, + ...(partial.thinking.length > 0 + ? { thinking: partial.thinking } + : {}), + }, + ...notes.map( + (content): ModelMessage => ({ + id: createMessageId(), + role: 'user', + content, + }), + ), + ]) + } + const operation = this.createTurnOperation() + this.bindTurn(input.inputId, operation) + this.feed.publish( + operation.id, + customEvent(HARNESS_EVENTS.operationResumed, { + operationId: operation.id, + ...(operationId ? { resumedFrom: operationId } : {}), + }), + ) + // The transcript has the message, so the turn adds none. Plugins and + // the router still get what the input sent. + const sent = input.input + this.enqueueTurn({ + operation, + inputId: input.inputId, + principal: input.principal, + ...('message' in sent ? { sentMessage: sent.message } : {}), + ...('context' in sent ? { context: sent.context } : {}), + overrides: decision.overrides, + }) + } + } +} diff --git a/packages/ai-harness/src/session/helpers.ts b/packages/ai-harness/src/session/helpers.ts new file mode 100644 index 0000000000..92057f2eb3 --- /dev/null +++ b/packages/ai-harness/src/session/helpers.ts @@ -0,0 +1,667 @@ +import { + EventType, + INTERRUPT_PAYLOAD_METADATA_KEY, + compactForModel, + provideDetachableRun, + readInterruptBinding, +} from '@tanstack/ai' +import { isRecord } from '../utils' +import type { + AnyChatMiddleware, + AnyTool, + Interrupt, + Modality, + ModelMessage, + PromptCacheOptions, + Scope, + StreamChunk, + SubagentsBag, +} from '@tanstack/ai' +import type { LogRecord, MessageStore } from '@tanstack/ai-persistence' +import type { EventFeed } from '../feed' +import type { InputState } from '../log' +import type { MountedPlugins } from '../plugins' +import type { + MediaKind, + OperationStatus, + Principal, + ThreadSettings, +} from '../types' +import type { + AgentRunInfo, + InterruptedTurn, + AgentMessage, + AgentInputState, +} from './types' + +/** Where `stores.metadata` keeps the interrupted turn of each thread. */ +export const INTERRUPTED = 'harness:interrupted' + +/** Where `stores.metadata` keeps the usage totals of a host without a log. */ +export const USAGE = 'harness:usage' + +/** + * The key of a resumable agent run in the session log: its steps, and the + * tool results of its chat. Not a tool call id, so a turn never reads them. + */ +export const agentKey = (inputId: string) => `agent:${inputId}` + +/** The chain fields that the `agent` input of a first run keeps. */ +export const chainFields = (first: AgentInputState) => ({ + runInputId: first.inputId, + target: first.input.agent, + input: first.input.input, + options: { + wake: first.input.detached === true, + ...(first.input.resume ? { resume: true } : {}), + ...(first.input.attach ? { attach: first.input.attach } : {}), + }, + ...(first.principal ? { principal: first.principal } : {}), + ...(first.input.parentRunId ? { parentRunId: first.input.parentRunId } : {}), +}) + +/** A message to an agent run as a user message. Its id is its input id. */ +export const asUserMessage = (sent: AgentMessage): ModelMessage => ({ + id: sent.inputId, + role: 'user', + content: sent.message, +}) + +/** True for the `agent` input of a first run. */ +export const isAgentInput = ( + input: InputState | undefined, +): input is AgentInputState => input?.input.op === 'agent' + +/** + * The `agent` input of the chain that `input` runs: the input itself, or + * the first run of a follow-up message. `undefined` for other inputs, and + * for a steer that joined a run. + */ +export const chainInputOf = ( + inputs: ReadonlyMap, + input: InputState, +): AgentInputState | undefined => { + if (isAgentInput(input)) return input + if (input.input.op !== 'agentMessage' || input.into !== undefined) { + return undefined + } + const first = inputs.get(input.input.run ?? '') + return isAgentInput(first) ? first : undefined +} + +/** The `agentRuns()` status of a live run. */ +export const runStatus = (status: OperationStatus): AgentRunInfo['status'] => { + switch (status) { + case 'accepted': + return 'queued' + case 'interrupted': + return 'running' + case 'running': + case 'completed': + case 'failed': + case 'cancelled': + return status + } +} + +/** The `agentRuns()` status of a run in the log, from how it ended. */ +export const loggedStatus = (input: InputState): AgentRunInfo['status'] => { + const outcome = input.settlement?.outcome + switch (outcome) { + case undefined: + return 'running' + case 'aborted': + return 'cancelled' + case 'failed': + return 'failed' + case 'completed': + case 'interrupted': + return 'completed' + } +} + +/** True for an interrupted turn as `stores.metadata` gives it back. */ +export function isInterruptedTurn(value: unknown): value is InterruptedTurn { + if (!isRecord(value) || typeof value.runId !== 'string') return false + if (!Array.isArray(value.interrupts)) return false + if (!value.interrupts.every((item) => isRecord(item))) return false + const { routed, principal } = value + const isPrincipal = + principal === undefined || + (isRecord(principal) && typeof principal.id === 'string') + return ( + isPrincipal && + (routed === undefined || + (isRecord(routed) && + Array.isArray(routed.messages) && + typeof routed.root === 'boolean')) + ) +} + +/** The connector a sign-in interrupt (`require` with `wait`) waits for. */ +export function signInConnector(interrupt: Interrupt | undefined) { + if (interrupt?.reason !== 'auth_required') return undefined + const payload = interrupt.metadata?.[INTERRUPT_PAYLOAD_METADATA_KEY] + const request = isRecord(payload) ? payload.request : undefined + return isRecord(request) && typeof request.connector === 'string' + ? request.connector + : undefined +} + +/** Limits for a harness's children when `subagents.limits` is not set. */ +export const DEFAULT_SUBAGENT_LIMITS = { + maxDepth: 2, + maxConcurrent: 3, + maxCalls: 12, +} + +export function createInputId(): string { + return `in-${Date.now().toString(36)}-${Math.random().toString(36).slice(2, 9)}` +} + +export function createMessageId(): string { + return `msg-${Date.now().toString(36)}-${Math.random().toString(36).slice(2, 9)}` +} + +export function customEvent( + name: string, + value: Record, +): StreamChunk { + return { type: EventType.CUSTOM, name, value, timestamp: Date.now() } +} + +export const errorText = (error: unknown) => + error instanceof Error ? error.message : String(error) + +/** A `promptCache` option as an object. A string is the retention alone. */ +export const cacheObject = (option: PromptCacheOptions | undefined) => + typeof option === 'string' ? { retention: option, key: undefined } : option + +/** The credential scope of `principal` in a thread, with only the set fields. */ +export const scopeOf = ( + threadId: string, + principal: Principal | undefined, +): Scope => ({ + threadId, + ...(principal ? { userId: principal.id } : {}), + ...(principal?.tenantId ? { tenantId: principal.tenantId } : {}), +}) + +/** What the log and the inbox keep of a principal: its id and tenant. */ +/** The `metadata.tanstack` of a logged event, if it has one. */ +export const tanstackOf = (event: Record) => + isRecord(event.metadata) && isRecord(event.metadata.tanstack) + ? event.metadata.tanstack + : undefined + +/** The tool call an interrupt waits on: an approval or a client tool. */ +export const toolCallOf = (interrupt: Interrupt) => { + const binding = readInterruptBinding(interrupt) + return binding && 'toolCallId' in binding ? binding.toolCallId : undefined +} + +export const storedPrincipal = (principal: Principal | undefined) => + principal + ? { + principal: { + id: principal.id, + ...(principal.tenantId ? { tenantId: principal.tenantId } : {}), + }, + } + : {} + +/** One sender: the same id and tenant, or no principal on both. */ +export const isSameSender = ( + a: Pick | undefined, + b: Pick | undefined, +) => a?.id === b?.id && a?.tenantId === b?.tenantId + +/** + * What makes two inputs with one id the same input: the payload and the + * sender. The same id from another sender is a conflict, not a retry. + */ +export const inputKey = ( + input: unknown, + principal: Pick | undefined, +) => JSON.stringify([input, principal?.id ?? null, principal?.tenantId ?? null]) + +export const isPlainObject = ( + value: unknown, +): value is Record => + isRecord(value) && + [Object.prototype, null].includes(Object.getPrototypeOf(value)) + +/** + * The chat() context of a turn. Two plain objects merge, and the harness + * value wins, so a client cannot replace a server value. Else the harness + * value when it is set, else the input's. + */ +export function turnContext(harness: unknown, input: unknown) { + if (isPlainObject(harness) && isPlainObject(input)) { + return { ...input, ...harness } + } + return harness !== undefined ? harness : input +} + +/** Where `stores.metadata` keeps the stored settings of each thread. */ +export const SETTINGS = 'harness:settings' + +export const SETTING_KEYS = new Set([ + 'model', + 'reasoning', + 'instructions', + 'tools', + 'plugins', + 'cwd', +]) + +export const isNames = (value: unknown): value is Array => + Array.isArray(value) && value.every((item) => typeof item === 'string') + +/** The names a setting can point at: the harness models and plugins. */ +export interface SettingNames { + models: ReadonlySet + plugins: ReadonlySet +} + +/** One checked setting value. Throws with a short reason for a bad one. */ +export function checkSetting( + key: string, + value: unknown, + names: SettingNames, +): unknown { + if (key === 'model') { + if (typeof value === 'string' && names.models.has(value)) return value + throw new Error( + `Unknown model ${JSON.stringify(value)}. Add it to defineHarness({ models }).`, + ) + } + if (key === 'instructions' || key === 'cwd') { + if (typeof value === 'string') return value + throw new Error(`The setting "${key}" must be a string.`) + } + if (key === 'reasoning') { + // chat() checks the level itself. + const isOption = + typeof value === 'string' || + (isRecord(value) && typeof value.level === 'string') + if (isOption) return value + throw new Error('The setting "reasoning" must be a level or { level }.') + } + if (key === 'tools') { + if (isNames(value)) return [...value] + if (isRecord(value) && isNames(value.remove)) { + return { remove: [...value.remove] } + } + throw new Error( + 'The setting "tools" must be a list of tool names or { remove: [names] }.', + ) + } + // `plugins` + if (!isRecord(value) || !isNames(value.remove)) { + throw new Error('The setting "plugins" must be { remove: [plugin names] }.') + } + const unknown = value.remove.find((name) => !names.plugins.has(name)) + if (unknown !== undefined) { + throw new Error(`Unknown plugin ${JSON.stringify(unknown)}.`) + } + return { remove: [...value.remove] } +} + +/** + * `settings` with `change` applied: `null` clears a field, a missing field + * stays. Throws with a short reason for an unknown field or a bad value. + */ +export function changeSettings( + settings: ThreadSettings, + change: unknown, + names: SettingNames, +): ThreadSettings { + if (!isRecord(change)) throw new Error('The settings must be an object.') + const next: Record = { ...settings } + for (const [key, value] of Object.entries(change)) { + if (!SETTING_KEYS.has(key)) { + throw new Error(`Unknown setting ${JSON.stringify(key)}.`) + } + if (value === null) delete next[key] + else if (value !== undefined) next[key] = checkSetting(key, value, names) + } + // `checkSetting` gave each field its type. + return next as ThreadSettings +} + +/** + * `mounted` without what the `removed` plugins give a turn: their tools, + * prompts, middleware, adapter picks, discoverers, and preparers. + */ +export function withoutPlugins( + mounted: MountedPlugins, + removed: ReadonlySet, +): MountedPlugins { + const { owners } = mounted + const toolOwners = new Map( + owners.tools.map((entry) => [entry.name, entry.owner]), + ) + const isKept = (owner: string | undefined) => + owner === undefined || !removed.has(owner) + return { + ...mounted, + tools: mounted.tools.filter((tool) => isKept(toolOwners.get(tool.name))), + prompts: mounted.prompts.filter((_prompt, index) => + isKept(owners.prompts[index]?.owner), + ), + middleware: mounted.middleware.filter((_item, index) => + isKept(owners.middleware[index]), + ), + adapters: mounted.adapters.filter((_pick, index) => + isKept(owners.adapters[index]), + ), + discoverers: mounted.discoverers.filter((item) => isKept(item.owner)), + preparers: mounted.preparers.filter((item) => isKept(item.owner)), + } +} + +/** + * The tools that the `tools` setting keeps: the listed ones, or all but the + * removed ones. The tools named in `keep` (the turn's own) always stay. + */ +export function settingTools( + tools: Array, + setting: ThreadSettings['tools'], + keep: ReadonlySet, +) { + if (setting === undefined) return tools + const isList = Array.isArray(setting) + const names = new Set(isList ? setting : setting.remove) + return tools.filter( + (tool) => keep.has(tool.name) || names.has(tool.name) === isList, + ) +} + +/** + * The metadata of the user message of a wake turn: the harness sent it (a + * background job ended), not the user. The model gets it as usual. A view + * does not show it as something the user said. + */ +export const WAKE_METADATA = { harness: { wake: true } } + +/** + * The transcript message of a `reset()`: a user message with + * `metadata.harness.reset`. Its content is the note, which the model sees + * first. The id comes from the input, so the marker goes in once. + */ +export function resetMarker( + inputId: string, + note: string | undefined, +): ModelMessage { + return { + id: `reset-${inputId}`, + role: 'user', + content: note ?? '', + metadata: { harness: { reset: note !== undefined ? { note } : {} } }, + } +} + +export const isResetMarker = (message: ModelMessage) => + isRecord(message.metadata?.harness) && + isRecord(message.metadata.harness.reset) + +/** + * What the model sees of `messages`: only the messages after the last reset + * marker. A marker with a note stays, as a plain user message. Returns + * `messages` itself when it has no marker. + */ +export function resetContext( + messages: Array, +): Array { + const at = messages.findLastIndex(isResetMarker) + const marker = messages[at] + if (!marker) return messages + const after = messages.slice(at + 1) + if (marker.content === '') return after + const { metadata: _reset, ...note } = marker + return [note, ...after] +} + +/** + * Gives the model only the context after the last reset. It changes only + * what the adapter gets (`providerMessages`), so every save keeps the + * messages before the reset. + */ +export const resetCut: AnyChatMiddleware = { + name: 'harness:reset', + onConfig: (ctx, config) => { + if (ctx.phase !== 'init' && ctx.phase !== 'beforeModel') return undefined + const messages = config.providerMessages ?? config.messages + const cut = resetContext(messages) + return cut === messages ? undefined : { providerMessages: cut } + }, +} + +/** + * Gives the next model call the `pending` messages after its context, and + * empties `pending`. It changes only what the adapter gets + * (`providerMessages`), so no save keeps them. + */ +export function ephemeralCall(pending: Array): AnyChatMiddleware { + return { + name: 'harness:ephemeral', + onConfig: (ctx, config) => { + if (ctx.phase !== 'beforeModel' || pending.length === 0) return undefined + const messages = config.providerMessages ?? config.messages + return { providerMessages: [...messages, ...pending.splice(0)] } + }, + } +} + +/** + * Gives each model call of a turn the `ephemeral` messages of its input, + * after `start`: the last message of the context when the turn started. The + * messages that the turn adds come after them. It changes only what the + * adapter gets (`providerMessages`), so no save keeps them. + */ +export function turnEphemeral( + ephemeral: ReadonlyArray, + start: ModelMessage, +): AnyChatMiddleware { + const isStart = (message: ModelMessage) => + start.id !== undefined + ? message.id === start.id + : JSON.stringify(message) === JSON.stringify(start) + return { + name: 'harness:turn-ephemeral', + onConfig: (ctx, config) => { + if (ctx.phase !== 'beforeModel') return undefined + const messages = config.providerMessages ?? config.messages + // Count from the end: a reset cut or a middleware can change the start + // of what the adapter gets. When a compaction dropped `start`, they go + // first. + const after = + config.messages.length - 1 - config.messages.findLastIndex(isStart) + const at = Math.max(0, messages.length - after) + return { + providerMessages: [ + ...messages.slice(0, at), + ...ephemeral, + ...messages.slice(at), + ], + } + }, + } +} + +/** A short transcript note about how an agent ended, for the next turn. */ +export function referenceNote( + agent: string, + result: unknown, + ended: 'finished' | 'failed' = 'finished', +): string { + const body = + typeof result === 'string' + ? result + : JSON.stringify(compactForModel(result)) + const clipped = body.length > 2000 ? `${body.slice(0, 2000)}...` : body + return `[${agent} ${ended}] ${clipped}` +} + +/** Why recovery fails an agent run whose host stopped. */ +export const AGENT_STOPPED = 'The host stopped during this agent run.' + +/** + * The kinds a model reads: the adapter's input list, narrowed by + * `media.accepts`. Either one alone when only one is known, and `undefined` + * (send every kind) when neither is. + */ +export function acceptedKinds( + modalities: ReadonlyArray | undefined, + accepts: ReadonlyArray | undefined, +) { + if (modalities === undefined || accepts === undefined) + return accepts ?? modalities + return accepts.filter((kind) => modalities.includes(kind)) +} + +/** The inputs that run as chat turns. They settle, and so do agent inputs. */ +export const CHAT_OPS = new Set([ + 'prompt', + 'steer', + 'followUp', + 'continue', + 'resolve', +]) + +/** The inputs that leave work to do: the thread is busy until they end. */ +export const WORK_OPS = new Set([ + 'prompt', + 'steer', + 'followUp', + 'continue', + 'resolve', + 'reset', + 'agent', + 'agentMessage', +]) + +/** The inputs that change a waiting input. */ +export const CONTROL_OPS = new Set(['cancelInput', 'setDelivery']) + +/** Why a turn stops when its input passes its time limit. */ +export const TIMEOUT_REASON = 'harness:input-timeout' + +/** What the model gets after its partial answer, on a `'continue'`. */ +export const CONTINUE_NOTE = + 'Your last answer stopped early because of an error. Continue from the exact point where it stopped. Do not repeat the text you already wrote.' + +/** + * Why work stops at `close({ recoverable: true })`. A user cancel stops it + * with `RUN_CANCEL_REASON`. + */ +export const SHUTDOWN_REASON = 'harness:shutdown' + +/** + * A shutdown abort is a detach, not an end: `withPersistence` then writes no + * aborted state to the run record. It goes before `withPersistence`, which + * reads the detach in its own `onAbort`. + */ +export const detachOnShutdown: AnyChatMiddleware = { + name: 'harness:detach-on-shutdown', + onAbort: (ctx, info) => { + if (info.reason === SHUTDOWN_REASON) provideDetachableRun(ctx, true) + }, +} + +/** Why a turn fails when no option gives it an adapter. */ +export const NO_MODEL = + 'This turn has no model. Set `adapter` in defineHarness, return one from a plugin `adapter()`, or pass `overrides.adapter` to the turn.' + +/** + * The text of the last assistant message: the answer of an operation rebuilt + * from the log, or of a routed turn. + */ +export function lastAssistantText(messages: ReadonlyArray) { + const last = messages.findLast((message) => message.role === 'assistant') + return typeof last?.content === 'string' ? last.content : '' +} + +/** + * True when chat() runs its routed path for this bag: the router picks, and + * the engine does not run. So the middleware that gives the engine its + * history and its run lease does not run, and the turn does that work. + */ +export function isRoutedBag( + bag: Pick | undefined, +) { + return bag !== undefined && bag.agents.length > 0 && bag.router !== undefined +} + +/** + * Middleware that keeps the override tools of a turn. A middleware `onConfig` + * can return a tool list without them, and this adds back each one whose + * name is missing. + */ +export function keepTurnTools(kept: ReadonlyArray) { + return { + name: 'harness:turn-tools', + onConfig: (ctx, config) => { + if (ctx.phase !== 'init' && ctx.phase !== 'beforeModel') return + const names = new Set(config.tools.map((tool) => tool.name)) + const missing = kept.filter((tool) => !names.has(tool.name)) + if (missing.length === 0) return + return { tools: [...config.tools, ...missing] } + }, + } satisfies AnyChatMiddleware +} + +/** The turn has its final answer: an assistant message with no open tool call after it was applied. */ +export function hasFinalAnswer( + messages: ReadonlyArray, + appliedAt: number, +) { + const last = messages.at(-1) + return ( + messages.length > appliedAt && + last?.role === 'assistant' && + (last.toolCalls?.length ?? 0) === 0 + ) +} + +/** Throw for a record that a host may not write. */ +export function checkHostRecords(records: ReadonlyArray) { + for (const record of records) { + if (record.type.startsWith('harness.')) { + throw new Error( + `The record type ${JSON.stringify(record.type)} is reserved for the harness.`, + ) + } + if (record.thread !== undefined && typeof record.thread !== 'string') { + throw new Error( + 'The thread field of a log record must be a string: the thread id of a session of the same log.', + ) + } + } +} + +export const notOpen = () => + Promise.reject(new Error('The session is not open yet.')) + +/** The message store of a durable session before `open()` reads its log. */ +export const notOpenMessages: MessageStore = { + loadThread: notOpen, + saveThread: notOpen, +} + +/** `feed`, and `onEvent` after each publish. This is how the host sees events. */ +export function tapFeed( + feed: EventFeed, + onEvent: (event: StreamChunk) => void, +) { + return { + publish: (operationId, event) => { + feed.publish(operationId, event) + onEvent(event) + }, + head: () => feed.head(), + read: (options) => feed.read(options), + close: () => feed.close(), + } satisfies EventFeed +} diff --git a/packages/ai-harness/src/session/index.ts b/packages/ai-harness/src/session/index.ts new file mode 100644 index 0000000000..7f2ced8b79 --- /dev/null +++ b/packages/ai-harness/src/session/index.ts @@ -0,0 +1,14 @@ +export { HarnessSession } from './harness-session' +export { DEFAULT_SUBAGENT_LIMITS, acceptedKinds } from './helpers' +export type { + AgentRunOptions, + AgentStartOptions, + AgentHandle, + DynamicAgentHandle, + AgentHandles, + AgentRunInfo, + SessionSnapshot, + SessionInspection, + SessionDescription, + SessionDependencies, +} from './types' diff --git a/packages/ai-harness/src/session/inputs.ts b/packages/ai-harness/src/session/inputs.ts new file mode 100644 index 0000000000..e5700096f4 --- /dev/null +++ b/packages/ai-harness/src/session/inputs.ts @@ -0,0 +1,551 @@ +import { EventType } from '@tanstack/ai' +import { isRecord } from '../utils' +import { HARNESS_EVENTS, InputRejectedError } from '../types' +import { + createInputId, + customEvent, + storedPrincipal, + inputKey, + CHAT_OPS, + WORK_OPS, + lastAssistantText, + checkHostRecords, +} from './helpers' +import { SessionBase } from './base' +import type { OperationImpl } from '../operation' +import type { LogRecord } from '@tanstack/ai-persistence' +import type { HarnessRecord } from '../log' +import type { AnyHarness } from '../define' +import type { + ChatTurnResult, + HarnessInput, + InputSettlement, + Receipt, + SessionEvent, + WaitingInput, +} from '../types' +import type { QueuedTurn, WaitingSteer } from './types' + +/** + * A layer of {@link HarnessSession}: admitting inputs, receipts, and + * settlements. + */ +export abstract class SessionInputs< + THarness extends AnyHarness = AnyHarness, +> extends SessionBase { + /** + * Append host records to the session log, in one batch after the events + * that wait. Only a durable host (with `stores.log`) has a log. A `type` + * that starts with `harness.` is refused: the harness owns those. + * A record can name another session of the same log with `thread`. The + * append is all or nothing. + * + * With `project` on the host, a record can change the model context. The + * running turn sees the change at its next model call. + * + * @example + * ```ts + * await session.append([{ type: 'app.signal', text: 'The build failed.' }]) + * ``` + */ + async append(records: ReadonlyArray): Promise { + if (!this.writer) { + throw new Error( + 'session.append needs a durable host (a host with stores.log).', + ) + } + checkHostRecords(records) + await this.writer.append(records) + } + + /** + * Cancel one operation, or the running chat turn. On a durable host the + * abort request is stored first, so a turn that a crash stops later settles + * `aborted` and does not run again. + */ + /** + * Move a running tool call to the background: the call returns at once, + * and the job keeps running. When it ends, a note tells the model and + * wakes an idle session. Without `toolCallId`, it moves every running + * call that supports it, such as `bash` and the agent tools. + * With no such call, the receipt is rejected with `not_running`. + * + * @example + * ```ts + * await session.background() + * ``` + */ + async background(toolCallId?: string): Promise { + const inputId = createInputId() + const ids = + toolCallId === undefined ? [...this.detachable.keys()] : [toolCallId] + const moves = ids.flatMap((id) => this.detachable.get(id) ?? []) + if (moves.length === 0) { + return { inputId, status: 'rejected', reason: 'not_running' } + } + for (const id of ids) this.detachable.delete(id) + for (const move of moves) move() + return { inputId, status: 'accepted' } + } + + /** + * The inputs that wait, in the order they run: first the steers for the + * running turn, then the queued turns. Change one with `cancelInput` or + * `setDelivery`. + * + * @example + * ```ts + * for (const input of session.inputs()) { + * if (input.delivery === 'queue') await session.cancelInput(input.inputId) + * } + * ``` + */ + inputs() { + const waiting: Array = [] + // A steer with an abort request never runs. It settles with the turn. + for (const steer of this.steerQueue) { + if (this.isAbortRequested(steer.inputId)) continue + waiting.push({ + inputId: steer.inputId, + delivery: 'steer', + message: steer.message, + }) + } + // A resolve and a turn that recovery runs again have no new message. + for (const turn of this.queue) { + if (turn.inputId === undefined || turn.message === undefined) continue + waiting.push({ + inputId: turn.inputId, + delivery: 'queue', + message: turn.message, + }) + } + return waiting + } + + /** + * How a chat input ended: `completed`, `failed`, `aborted`, or + * `interrupted` (the turn waits for human input). It waits until the input + * ends. On a durable host it reads the log, so it also works after a + * restart and from another host that opens the thread. There it also + * takes the input id of an agent run from the log. + * + * Rejects with `InputRejectedError` when the session refused the input, + * and with an error for an id this session does not know. + * + * @example + * ```ts + * const turn = session.prompt('Summarize the report.', { inputId: 'req-42' }) + * const { outcome } = await session.settled('req-42') + * ``` + */ + async settled(inputId: string): Promise { + const known = this.knownSettlement(inputId) + if (known) return known + if (!this.isChatInput(inputId)) { + throw new Error( + `The session knows no chat input with the id ${JSON.stringify(inputId)}.`, + ) + } + const settledEvent = (entry: SessionEvent) => + entry.event.type === EventType.CUSTOM && + entry.event.name === HARNESS_EVENTS.inputSettled && + isRecord(entry.event.value) && + entry.event.value.inputId === inputId + for await (const _entry of this.feed.read({ + from: this.feed.head(), + filter: settledEvent, + })) { + const settlement = this.knownSettlement(inputId) + if (settlement) return settlement + } + throw new Error('The session closed before the input ended.') + } + + protected enqueueTurn(turn: QueuedTurn): void { + // Recovery queues turns without `accept`, so they claim here too. + this.markBusy() + this.queue.push(turn) + this.drain() + } + + /** + * Where a turn that runs next goes: after a resolve that answers the turn + * that just ended, and after the resets that wait, in their order. + */ + protected frontOfQueue(): number { + const at = this.queue.findIndex((turn) => !turn.answers && !turn.reset) + return at < 0 ? this.queue.length : at + } + + protected isAbortRequested(inputId: string) { + return ( + this.writer?.state.inputs.get(inputId)?.abortRequested === true || + this.abortedSteers.has(inputId) + ) + } + + /** + * `steer` joined the running turn `operationId`: it settles with that + * turn, and a prompt that joined gets the turn's result. + */ + protected join( + operationId: string, + steer: { inputId: string; operation?: OperationImpl }, + ) { + const joined = this.turnJoins.get(operationId) ?? [] + joined.push(steer.inputId) + this.turnJoins.set(operationId, joined) + const running = this.activeTurn + if (steer.operation && running) { + const follower = steer.operation + this.operations.delete(follower.id) + running.then( + (result) => follower.finish('completed', result), + (error: unknown) => follower.fail('failed', error), + ) + } + if (!this.writer) void this.inbox.markApplied(steer.inputId, operationId) + this.feed.publish( + operationId, + customEvent(HARNESS_EVENTS.inputApplied, { + inputId: steer.inputId, + operationId, + }), + ) + } + + /** The inputs that joined the turn of `hostInput` (or operation). */ + protected joinedInputs(operationId: string, hostInput: string | undefined) { + if (this.writer && hostInput !== undefined) { + return [...this.writer.state.inputs.values()] + .filter( + (input) => input.status === 'joined' && input.into === hostInput, + ) + .map((input) => input.inputId) + } + return this.turnJoins.get(operationId) ?? [] + } + + /** + * Take a waiting input out of the steers or the queue, so it does not run + * from there. `undefined` when it does not wait: it started, a join holds + * it, it has an abort request, or the id is not known. + */ + protected takeWaiting(inputId: string) { + const steer = this.steerQueue.find( + (item) => + item.inputId === inputId && + !this.joining.has(inputId) && + !this.isAbortRequested(inputId), + ) + // This shifts the later steers. `claimJoins` sees that and skips them. + if (steer) { + this.steerQueue.splice(this.steerQueue.indexOf(steer), 1) + return steer + } + const turn = this.queue.find((item) => item.inputId === inputId) + const message = turn?.message + if (!turn || message === undefined) return undefined + this.queue.splice(this.queue.indexOf(turn), 1) + const taken: WaitingSteer = { + inputId, + message, + operation: turn.operation, + overrides: turn.overrides, + ephemeral: turn.ephemeral, + principal: turn.principal, + context: turn.context, + ...(turn.wake ? { wake: true } : {}), + } + return taken + } + + /** + * Mark a `cancelInput` or `setDelivery` input applied. On a durable host, + * `change` lands in the same append, so a crash never keeps one without + * the other. + */ + protected async applyControl(inputId: string, change: HarnessRecord) { + if (!this.writer) { + await this.inbox.markApplied(inputId, 'session') + return + } + await this.writer.append([ + change, + { + type: 'harness.input.applied', + inputId, + operationId: 'session', + attempt: 1, + }, + ]) + } + + // =========================== + // Inbox and lifecycle events + // =========================== + + /** + * Store an input, unless its id is known. The check before the first + * `await` is synchronous, so a second call with the same id in the same + * tick is a duplicate too. `principal` is who sent it. + */ + protected async accept( + inputId: string, + input: HarnessInput, + principal = this.principal, + ): Promise<'new' | 'duplicate' | 'conflict'> { + const sender = storedPrincipal(principal) + const payload = inputKey(input, sender.principal) + const local = this.admitted.get(inputId) + if (local !== undefined) return local === payload ? 'duplicate' : 'conflict' + const logged = this.writer?.state.inputs.get(inputId) + if (logged) { + return inputKey(logged.input, logged.principal) === payload + ? 'duplicate' + : 'conflict' + } + this.admitted.set(inputId, payload) + // Before the input is stored, so a crash after the write leaves a claim + // that a sweep finds. + if (WORK_OPS.has(input.op)) { + this.markBusy() + await this.claimWrite + } + const at = Date.now() + if (this.writer) { + await this.writer.append([ + { type: 'harness.input', inputId, input, at, ...sender }, + ]) + } else { + const stored = await this.inbox.append({ + inputId, + threadId: this.threadId, + input, + createdAt: at, + ...sender, + }) + // The inbox had the id before this session: a restart. + if (stored.createdAt !== at || stored.status !== 'pending') { + return inputKey(stored.input, stored.principal) === payload + ? 'duplicate' + : 'conflict' + } + } + this.feed.publish( + 'session', + customEvent(HARNESS_EVENTS.inputAccepted, { inputId, op: input.op }), + ) + return 'new' + } + + /** + * Does `stores.runs` have a run with this id, on any thread? A client run + * id must not reuse it. A store that fails counts as yes, so nothing runs. + */ + protected async isRunTaken(runId: string): Promise { + const runs = this.persistence.stores.runs + if (!runs) return false + return runs.get(runId).then( + (run) => run !== null, + () => true, + ) + } + + /** Record that `operationId` runs the input. A durable host counts attempts. */ + protected async markApplied( + inputId: string, + operationId: string, + ): Promise { + if (!this.writer) { + await this.inbox.markApplied(inputId, operationId) + return + } + const logged = this.writer.state.inputs.get(inputId) + const timeoutMs = this.harness.durability?.timeoutMs + // Stamped once, at the first apply. Retries do not move it. + const timeoutAt = + logged?.timeoutAt ?? + (timeoutMs !== undefined ? Date.now() + timeoutMs : undefined) + await this.writer.append([ + { + type: 'harness.input.applied', + inputId, + operationId, + attempt: (logged?.attempt ?? 0) + 1, + ...(timeoutAt !== undefined ? { timeoutAt } : {}), + }, + ]) + } + + protected async applied(inputId: string, operationId: string): Promise { + await this.markApplied(inputId, operationId) + this.feed.publish( + operationId, + customEvent(HARNESS_EVENTS.inputApplied, { inputId, operationId }), + ) + } + + /** Refuse an input. Returns its rejected receipt. */ + protected reject(inputId: string, reason: string) { + if (this.writer) { + void this.writer + .append([{ type: 'harness.input.rejected', inputId, reason }]) + .catch(() => {}) + } else { + void this.inbox.markRejected(inputId, reason) + } + const receipt: Receipt = { inputId, status: 'rejected', reason } + this.receipts.set(inputId, receipt) + this.feed.publish( + 'session', + customEvent(HARNESS_EVENTS.inputRejected, { inputId, reason }), + ) + return receipt + } + + /** + * Record how inputs ended, all in one append: a turn and the inputs that + * joined it settle together. Clients get a `harness.input.settled` event. + */ + protected async settle( + ...settlements: Array + ): Promise { + if (this.writer) { + await this.writer.append( + settlements.map((settlement) => ({ + type: 'harness.input.settled', + ...settlement, + })), + ) + } + for (const settlement of settlements) { + this.settlements.set(settlement.inputId, settlement) + this.feed.publish( + settlement.operationId ?? 'session', + customEvent(HARNESS_EVENTS.inputSettled, { ...settlement }), + ) + } + } + + /** How a chat input ended, when it did. Throws for a refused input. */ + protected knownSettlement(inputId: string): InputSettlement | undefined { + const logged = this.writer?.state.inputs.get(inputId) + if (logged?.settlement) return logged.settlement + const refused = + logged?.status === 'rejected' + ? { inputId, status: 'rejected' as const, reason: logged.reason ?? '' } + : this.receipts.get(inputId) + if (refused?.status === 'rejected') throw new InputRejectedError(refused) + return this.settlements.get(inputId) + } + + protected isChatInput(inputId: string) { + const input = this.writer?.state.inputs.get(inputId)?.input + // A reset, an agent run, and a message to an agent run in the log + // settle their input too. + if (input) + return ( + CHAT_OPS.has(input.op) || + input.op === 'reset' || + input.op === 'agent' || + input.op === 'agentMessage' + ) + return this.turnOperations.has(inputId) || this.receipts.has(inputId) + } + + /** Link a turn operation and its input. */ + protected bindTurn( + inputId: string, + operation: OperationImpl, + ) { + this.turnOperations.set(inputId, operation) + this.operationInputs.set(operation.id, inputId) + } + + /** 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 + } + + /** Resolve a turn's receipt, and keep it for a duplicate. */ + protected answerTurn( + operation: OperationImpl, + receipt: Receipt, + ) { + operation.resolveReceipt(this.keep(receipt)) + } + + protected refuse(operation: OperationImpl, receipt: Receipt) { + this.answerTurn(operation, receipt) + operation.fail('failed', new InputRejectedError(receipt)) + } + + /** What a duplicate or a conflicting chat input gets back. */ + protected duplicateReceipt( + inputId: string, + admission: 'duplicate' | 'conflict', + ): Receipt { + if (admission === 'conflict') { + return { inputId, status: 'rejected', reason: 'conflict' } + } + const kept = this.receipts.get(inputId) + if (kept) return kept + const operationId = this.writer?.state.inputs.get(inputId)?.operationId + return { + inputId, + status: 'accepted', + ...(operationId ? { operationId } : {}), + } + } + + /** End `operation` with a stored settlement. */ + protected settleOperation( + operation: OperationImpl, + settlement: InputSettlement, + ) { + const text = lastAssistantText(this.writer?.state.messages ?? []) + switch (settlement.outcome) { + case 'completed': + operation.finish('completed', { text }) + return + case 'interrupted': + operation.finish('interrupted', { text, interrupts: [] }) + return + case 'aborted': + operation.fail('cancelled', new Error('Cancelled.')) + return + case 'failed': + operation.fail( + 'failed', + new Error(settlement.error?.message ?? 'The input failed.'), + ) + return + } + } + + protected publishStarted( + operation: OperationImpl | OperationImpl, + ): void { + operation.publish( + customEvent(HARNESS_EVENTS.operationStarted, { + operationId: operation.id, + kind: operation.kind, + ...(operation.agent ? { agent: operation.agent } : {}), + }), + ) + } + + protected publishFinished( + operation: OperationImpl | OperationImpl, + ): void { + operation.publish( + customEvent(HARNESS_EVENTS.operationFinished, { + operationId: operation.id, + status: operation.status(), + }), + ) + } +} diff --git a/packages/ai-harness/src/session/requests.ts b/packages/ai-harness/src/session/requests.ts new file mode 100644 index 0000000000..cace43749d --- /dev/null +++ b/packages/ai-harness/src/session/requests.ts @@ -0,0 +1,933 @@ +import { RUN_CANCEL_REASON } from '@tanstack/ai' +import { changedEntry } from '../session-entry' +import { OperationImpl } from '../operation' +import { HARNESS_EVENTS } from '../types' +import { + signInConnector, + createInputId, + customEvent, + errorText, + inputKey, + SHUTDOWN_REASON, +} from './helpers' +import { SessionState } from './state' +import type { + ModelMessage, + RunAgentResumeItem, + SchemaInput, +} from '@tanstack/ai' +import type { SessionIndexEntry } from '@tanstack/ai-persistence' +import type { PluginSessionApi, Question } from '../commands' +import type { AnyHarness } from '../define' +import type { + BusyPolicy, + ChatTurnResult, + HarnessInput, + Operation, + Principal, + Receipt, + TurnOverrides, + UserInput, + WaitingInput, +} from '../types' +import type { QueuedTurn, WaitingSteer } from './types' + +/** + * A layer of {@link HarnessSession}: prompt, steer, follow-up, continue, reset, + * resolve, and cancel. + */ +export abstract class SessionRequests< + THarness extends AnyHarness = AnyHarness, +> extends SessionState { + // =========================== + // Inputs + // =========================== + + /** + * Start a chat turn, or queue it while one runs (see `busy`). + * + * `inputId` is an id you choose. A second prompt with the same id and the + * same message returns the first input's operation and does not run again. + * The same id with another message is rejected with `'conflict'`. + * `await operation.receipt` resolves when the input is stored. + * + * `overrides` changes the adapter, reasoning, prompt cache, or tools of + * this turn only (see `TurnOverrides`). A steer that joins the running + * turn uses the overrides of that turn, and its own are ignored. + * + * `principal` is who sent the message. Default: the principal that opened + * the session. The log keeps it, and the turn runs with its credentials + * and provider keys. A steer that joins the running turn runs with the + * credentials of that turn's sender. + * + * `context` is JSON data from the client, stored with the input. Tools + * get it in their context (see `HarnessConfig.context`). Do not trust it. + * + * `ephemeral` are messages that each model call of this turn gets, after + * the transcript and the new user message, before the messages that the + * turn adds (answers, tool results, steers). No store keeps them: not the + * transcript, and not the log. Like `overrides`, they are not part of the + * input: the `inputId` duplicate check does not read them, so a retry with + * other ephemeral messages gets the first input's operation and its own + * are dropped. A turn that recovery runs again has none. A steer that joins + * the running turn ignores them. + */ + prompt( + message: UserInput, + options?: { + busy?: BusyPolicy + inputId?: string + overrides?: TurnOverrides + principal?: Principal + context?: unknown + ephemeral?: ReadonlyArray + /** @internal The AG-UI `runId` of a `POST run` request. The turn runs as it. */ + runId?: string + }, + ): Operation { + const busy = options?.busy ?? this.harness.busy ?? 'queue' + const inputId = options?.inputId ?? createInputId() + const overrides = options?.overrides + const ephemeral = options?.ephemeral + const principal = options?.principal ?? this.principal + const context = options?.context + const input: HarnessInput = { + op: 'prompt', + message, + busy, + ...(context !== undefined ? { context } : {}), + } + const known = this.knownTurn(inputId, input, principal) + if (known) return known + const runId = options?.runId + // A run id must not take the id of another operation. + if (runId !== undefined && this.operations.has(runId)) { + const refused = this.createTurnOperation() + this.refuse(refused, { inputId, status: 'rejected', reason: 'conflict' }) + return refused + } + const operation = this.createTurnOperation(runId) + this.bindTurn(inputId, operation) + void this.admitting(() => + this.accept(inputId, input, principal).then( + async (admission) => { + if (admission !== 'new') { + this.answerDuplicate(operation, inputId, admission) + return + } + // Awaited only for a client run id, so other prompts keep their order. + if (runId !== undefined && (await this.isRunTaken(runId))) { + this.reject(inputId, 'conflict') + this.refuse(operation, { + inputId, + status: 'rejected', + reason: 'conflict', + }) + return + } + if (this.activeTurn && busy === 'reject') { + this.reject(inputId, 'busy') + this.refuse(operation, { + inputId, + status: 'rejected', + reason: 'busy', + }) + return + } + const isWaiting = + this.activeTurn !== undefined || this.queue.length > 0 + // A steer joins the running turn, so the receipt names that turn, as + // `steer()` does. + const running = busy === 'steer' ? this.activeTurn : undefined + this.answerTurn(operation, { + inputId, + status: isWaiting && busy !== 'steer' ? 'queued' : 'accepted', + operationId: running?.id ?? operation.id, + }) + const sent = { principal, context } + if (running) { + this.steerQueue.push({ + inputId, + message, + operation, + overrides, + ephemeral, + ...sent, + }) + return + } + this.enqueueTurn({ + operation, + message, + inputId, + overrides, + ephemeral, + ...sent, + }) + }, + (error: unknown) => { + // The input was not stored. After a log failure, say why the + // session stopped. + const failure = this.logFailure ?? error + this.answerTurn(operation, { + inputId, + status: 'rejected', + reason: + failure instanceof Error ? failure.message : String(failure), + }) + operation.fail('failed', failure) + }, + ), + ) + return operation + } + + /** + * Add a message to the running turn at its next model call. `principal` + * and `context` mean the same as in `prompt`. + */ + async steer( + message: UserInput, + options?: { inputId?: string; principal?: Principal; context?: unknown }, + ): Promise { + return this.admitting(async () => { + const inputId = options?.inputId ?? createInputId() + const principal = options?.principal ?? this.principal + const context = options?.context + const admission = await this.accept( + inputId, + { + op: 'steer', + message, + ...(context !== undefined ? { context } : {}), + }, + principal, + ) + if (admission !== 'new') return this.duplicateReceipt(inputId, admission) + if (!this.activeTurn) { + const operation = this.createTurnOperation() + this.bindTurn(inputId, operation) + this.enqueueTurn({ operation, message, inputId, principal, context }) + return this.keep({ + inputId, + status: 'accepted', + operationId: operation.id, + }) + } + this.steerQueue.push({ inputId, message, principal, context }) + return this.keep({ + inputId, + status: 'accepted', + operationId: this.activeTurn.id, + }) + }) + } + + /** + * Run a turn after the current work settles. `overrides` changes the + * adapter, reasoning, prompt cache, or tools of this turn only. + * `principal` and `context` mean the same as in `prompt`. + */ + async followUp( + message: UserInput, + options?: { + inputId?: string + overrides?: TurnOverrides + principal?: Principal + context?: unknown + }, + ): Promise { + return this.queueFollowUp(message, options) + } + + /** + * `followUp`, plus `wake`: the harness wakes the session with `message` + * (a background job ended), so its user message is marked, not shown as + * something the user said. + */ + protected async queueFollowUp( + message: UserInput, + options?: { + inputId?: string + overrides?: TurnOverrides + principal?: Principal + context?: unknown + wake?: boolean + }, + ): Promise { + return this.admitting(async () => { + const inputId = options?.inputId ?? createInputId() + const principal = options?.principal ?? this.principal + const context = options?.context + const wake = options?.wake === true + const admission = await this.accept( + inputId, + { + op: 'followUp', + message, + ...(context !== undefined ? { context } : {}), + ...(wake ? { wake } : {}), + }, + principal, + ) + if (admission !== 'new') return this.duplicateReceipt(inputId, admission) + const operation = this.createTurnOperation() + this.bindTurn(inputId, operation) + const status = + this.activeTurn || this.queue.length > 0 ? 'queued' : 'accepted' + this.enqueueTurn({ + operation, + message, + inputId, + overrides: options?.overrides, + principal, + context, + ...(wake ? { wake } : {}), + }) + return this.keep({ inputId, status, operationId: operation.id }) + }) + } + + /** + * Start a chat turn from the stored transcript, with no new message. While + * a turn runs, it queues. When it starts, the transcript must end with a + * user or a tool message. Else the input is rejected with + * `'nothing_to_continue'`. `inputId`, `overrides`, `principal`, + * `context`, and `ephemeral` mean the same as in `prompt`. Each model call + * gets the `ephemeral` messages after the transcript. + * + * @example + * ```ts + * await session.append([{ type: 'app.signal', text: 'The build failed.' }]) + * await session.continue({ inputId: 'after-signal' }) + * ``` + */ + continue(options?: { + inputId?: string + overrides?: TurnOverrides + principal?: Principal + context?: unknown + ephemeral?: ReadonlyArray + }): Operation { + const inputId = options?.inputId ?? createInputId() + const principal = options?.principal ?? this.principal + const context = options?.context + const input: HarnessInput = { + op: 'continue', + ...(context !== undefined ? { context } : {}), + } + const known = this.knownTurn(inputId, input, principal) + if (known) return known + const operation = this.createTurnOperation() + this.bindTurn(inputId, operation) + void this.admitting(() => + this.accept(inputId, input, principal).then( + (admission) => { + if (admission !== 'new') { + this.answerDuplicate(operation, inputId, admission) + return + } + const isWaiting = + this.activeTurn !== undefined || this.queue.length > 0 + this.answerTurn(operation, { + inputId, + status: isWaiting ? 'queued' : 'accepted', + operationId: operation.id, + }) + this.enqueueTurn({ + operation, + inputId, + overrides: options?.overrides, + ephemeral: options?.ephemeral, + principal, + context, + isContinue: true, + }) + }, + (error: unknown) => { + // As in `prompt`: the input was not stored. + const failure = this.logFailure ?? error + this.answerTurn(operation, { + inputId, + status: 'rejected', + reason: errorText(failure), + }) + operation.fail('failed', failure) + }, + ), + ) + return operation + } + + /** + * Start a fresh model context. From the next turn, the model sees only + * `note` (when you give one) and what comes after the reset. + * `transcript()` keeps every message: the reset adds a user message with + * `metadata.harness.reset`, whose content is the note. + * + * While a turn runs, the reset waits and applies when that turn ends, + * before the next queued turn. A thread that waits for interrupts is + * refused with `'pending_interrupts'`: resolve them first. `principal` and + * `inputId` mean the same as in `prompt`. + */ + async reset( + note?: string, + options?: { principal?: Principal; inputId?: string }, + ): Promise { + return this.admitting(async () => { + const inputId = options?.inputId ?? createInputId() + const principal = options?.principal ?? this.principal + // An empty note is no note: the model never gets an empty message. + const text = note?.trim() ? note : undefined + const admission = await this.accept( + inputId, + { op: 'reset', ...(text !== undefined ? { note: text } : {}) }, + principal, + ) + if (admission !== 'new') return this.duplicateReceipt(inputId, admission) + const isRunning = this.activeTurn !== undefined + if (!isRunning && this.queue.length === 0 && this.interrupted) { + this.reject(inputId, 'pending_interrupts') + return { inputId, status: 'rejected', reason: 'pending_interrupts' } + } + const operation = this.queueReset(inputId, text, principal, 'front') + return this.keep({ + inputId, + status: isRunning ? 'queued' : 'accepted', + operationId: operation.id, + }) + }) + } + + /** + * Answer the interrupts of the last turn. One resume must answer every open + * interrupt of that turn (the AG-UI rule). `principal` means the same as in + * `prompt`: the turn that continues runs with its credentials. + */ + async resolve( + resume: Array, + options?: { + inputId?: string + principal?: Principal + /** @internal The AG-UI `runId` of a `POST run` request. The turn runs as it. */ + runId?: string + }, + ): Promise { + return this.admitting(async () => { + const inputId = options?.inputId ?? createInputId() + const principal = options?.principal ?? this.principal + const admission = await this.accept( + inputId, + { op: 'resolve', resume }, + principal, + ) + if (admission !== 'new') return this.duplicateReceipt(inputId, admission) + const runId = options?.runId + // A run id must not take the id of another operation or stored run. + const isTaken = + runId !== undefined && + (this.operations.has(runId) || (await this.isRunTaken(runId))) + if (isTaken) { + this.reject(inputId, 'conflict') + return { inputId, status: 'rejected', reason: 'conflict' } + } + const answers = this.interrupted + if (!answers) { + this.reject(inputId, 'no_pending_interrupts') + return { inputId, status: 'rejected', reason: 'no_pending_interrupts' } + } + // A client answers on RUN_FINISHED, while the interrupted turn still + // ends. The resolve runs right after it. + const isEnding = this.activeTurn?.id === answers.runId + if ( + (this.activeTurn && !isEnding) || + this.queue.some((turn) => turn.answers?.runId === answers.runId) + ) { + this.reject(inputId, 'busy') + return { inputId, status: 'rejected', reason: 'busy' } + } + // A cancelled sign-in fails its tool when the turn runs it again. + for (const entry of resume) { + const connector = signInConnector( + answers.interrupts.find((item) => item.id === entry.interruptId), + ) + if (connector && entry.status === 'cancelled') + this.declinedSignIns.add(connector) + } + const operation = this.createTurnOperation(runId) + this.bindTurn(inputId, operation) + // The turn goes on with the context its input sent. + const turn: QueuedTurn = { + operation, + resume, + inputId, + answers, + principal, + context: answers.context, + } + if (isEnding) this.queue.unshift(turn) + else this.enqueueTurn(turn) + return this.keep({ + inputId, + status: 'accepted', + operationId: operation.id, + }) + }) + } + + async cancel(operationId?: string): Promise { + const inputId = createInputId() + await this.accept(inputId, { op: 'cancel', operationId }) + const target = operationId + ? this.operations.get(operationId) + : this.activeTurn + if (!target || target.isSettled()) { + return this.reject(inputId, 'not_running') + } + const waitingSteer = this.steerQueue.find( + (steer) => steer.operation === target, + ) + // A join took it, so it ends with the running turn, as a joined steer. + if (waitingSteer && this.joining.has(waitingSteer.inputId)) { + return this.reject(inputId, 'not_running') + } + // It never joins now, and settles when the running turn ends. + if (waitingSteer) this.abortedSteers.add(waitingSteer.inputId) + const targetInput = this.operationInputs.get(target.id) + if (this.writer && targetInput) { + await this.writer.append([ + { type: 'harness.input.abort', inputId: targetInput }, + ]) + } + if (waitingSteer) { + await this.markApplied(inputId, target.id) + return { inputId, status: 'accepted', operationId: target.id } + } + const chain = this.chainOf.get(target.id) + const waiting = + chain?.followUps.findIndex((item) => item.operation === target) ?? -1 + if (chain && waiting >= 0) { + // A follow-up that did not start never runs. + const [item] = chain.followUps.splice(waiting, 1) + if (item) { + await this.settle({ + inputId: item.inputId, + outcome: 'aborted', + operationId: target.id, + }) + } + target.fail('cancelled', new Error('Cancelled before it started.')) + this.publishFinished(target) + await this.markApplied(inputId, target.id) + this.checkIdle() + return { inputId, status: 'accepted', operationId: target.id } + } + // A run's children stop first, bottom-up. + for (const child of chain?.children ?? []) { + if (!child.current.operation.isSettled()) { + await this.cancel(child.current.operation.id) + } + } + const queued = this.queue.findIndex((turn) => turn.operation === target) + if (queued >= 0) { + this.queue.splice(queued, 1) + if (targetInput) { + await this.settle({ + inputId: targetInput, + outcome: 'aborted', + operationId: target.id, + }) + } + await this.refreshTurnInterrupts(target) + target.fail('cancelled', new Error('Cancelled before it started.')) + this.publishFinished(target) + } else { + target.abortController.abort(RUN_CANCEL_REASON) + } + await this.markApplied(inputId, target.id) + return { inputId, status: 'accepted', operationId: target.id } + } + + /** + * Cancel an input that waits, so it never runs, also after a restart. It + * settles `aborted`. An input that started, or that joins the running + * turn now, is refused with `not_waiting`: stop it with + * `cancel(operationId)`. + * + * @example + * ```ts + * await session.followUp('Then write the tests.', { inputId: 'tests' }) + * await session.cancelInput('tests') + * ``` + */ + async cancelInput(target: string): Promise { + const inputId = createInputId() + await this.accept(inputId, { op: 'cancelInput', inputId: target }) + const taken = this.takeWaiting(target) + if (!taken) return this.reject(inputId, 'not_waiting') + await this.applyControl(inputId, { + type: 'harness.input.abort', + inputId: target, + }) + // Without a log, the inbox entry must not stay pending: a restart runs + // pending entries. + if (!this.writer) await this.inbox.markRejected(target, 'cancelled') + const { operation } = taken + await this.settle({ + inputId: target, + outcome: 'aborted', + ...(operation ? { operationId: operation.id } : {}), + }) + if (operation) { + operation.fail('cancelled', new Error('Cancelled before it started.')) + this.publishFinished(operation) + } + return { inputId, status: 'accepted' } + } + + /** + * Move an input that waits. `'steer'` joins it to the running turn at the + * next model call. `'queue'` runs it as its own turn, after the queued + * turns. An input that started, or that joins the running turn now, is + * refused with `not_waiting`. On a durable host the log keeps the + * delivery, and a restart honors it. Without a log, it is kept in memory. + * + * @example + * ```ts + * await session.followUp('Use tabs, not spaces.', { inputId: 'style' }) + * await session.setDelivery('style', 'steer') + * ``` + */ + async setDelivery( + target: string, + delivery: WaitingInput['delivery'], + ): Promise { + const inputId = createInputId() + await this.accept(inputId, { op: 'setDelivery', inputId: target, delivery }) + const current = this.inputs().find((input) => input.inputId === target) + if (current?.delivery !== delivery) { + const taken = this.takeWaiting(target) + if (!taken) return this.reject(inputId, 'not_waiting') + if (delivery === 'steer') this.steerQueue.push(taken) + else this.enqueueTurn(this.steerTurn(taken)) + } + await this.applyControl(inputId, { + type: 'harness.input.delivery', + inputId: target, + delivery, + }) + this.feed.publish( + 'session', + customEvent(HARNESS_EVENTS.inputDelivery, { inputId: target, delivery }), + ) + return { inputId, status: 'accepted' } + } + + /** + * A saved credential can answer the sign-ins the last turn waits for. When + * every open interrupt is a sign-in and each of their connectors has a + * credential of the turn's sender now, the turn goes on as that sender and + * its tools run again. The input id comes from the interrupts, so a second + * save does not resolve them twice. + */ + protected override async resumeSignIns(): Promise { + const stopped = this.interrupted + const waiting = stopped?.interrupts ?? [] + const connectors = waiting.map(signInConnector) + if (waiting.length === 0) return + // Another user's credential never answers this sign-in. + const principal = stopped?.principal ?? this.principal + const credentials = this.credentialsOf(principal) + for (const connector of connectors) { + if (!connector || !(await credentials.get(connector))) return + } + if (this.interrupted !== stopped) return + await this.resolve( + waiting.map((interrupt, index) => ({ + interruptId: interrupt.id, + status: 'resolved', + payload: { connector: connectors[index] }, + })), + { + inputId: `sign-in:${waiting.map((item) => item.id).join(',')}`, + ...(principal ? { principal } : {}), + }, + ) + } + + /** + * The session API of plugins. With `principal` (a command), it acts for + * that user. Without it, it follows the sender of the running turn. + */ + protected override pluginApi( + principal?: Principal, + metadataSchema?: SchemaInput, + ): PluginSessionApi { + const sender = () => principal ?? this.sender() + return { + threadId: this.threadId, + get principal() { + return sender() + }, + snapshot: () => this.snapshot(), + // A plugin turn always waits its turn. It cannot see a busy rejection. + prompt: (text) => + this.prompt(text, { busy: 'queue', principal: sender() }), + note: async (text, options) => { + await this.addNote(text) + if (options?.wake) + await this.queueFollowUp(text, { principal: sender(), wake: true }) + }, + transcript: () => this.transcript(), + replaceTranscript: (messages) => + this.messages.saveThread(this.threadId, messages), + entry: () => this.index.get(this.threadId), + updateEntry: (update) => { + // Only these fields: a plugin cannot change the owner, the parent, + // or the pin. + const pick = ( + change: + | Pick + | undefined, + ) => + change && { + title: change.title, + metadata: change.metadata, + usage: change.usage, + } + return this.index.update(this.threadId, (entry) => + changedEntry( + entry, + typeof update === 'function' + ? (current) => pick(update(current)) + : pick(update), + this.harness.metadataSchema, + metadataSchema, + ), + ) + }, + // The public type narrows the answer from the schema. + ask: ((question: Question) => + this.ask(question)) as PluginSessionApi['ask'], + authRequired: (info) => + this.feed.publish( + 'session', + customEvent(HARNESS_EVENTS.authRequired, { ...info }), + ), + setConfig: (key, value) => this.setConfig(key, value), + settings: () => this.settings(), + } + } + + /** + * Stop every running operation, wait for them, then dispose session plugins. + * Safe to call twice. + * + * 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 }) + * ``` + */ + close(options?: { recoverable?: boolean }): Promise { + this.closing ??= (async () => { + // A thread with work keeps its claim, so a sweep finds that work after + // the claim expires. An idle thread gives it back. + const wasIdle = this.isIdle() + this.stopClaimTimer() + if (wasIdle) await this.releaseClaim() + // Before the stop, so no write of the stopped work lands. + if (options?.recoverable) this.writer?.close() + for (const turn of this.queue.splice(0)) { + await this.refreshTurnInterrupts(turn.operation) + turn.operation.fail('cancelled', new Error('Session closed.')) + } + const running = [...this.operations.values()].filter( + (operation) => !operation.isSettled(), + ) + for (const operation of running) { + operation.abortController.abort( + options?.recoverable ? SHUTDOWN_REASON : RUN_CANCEL_REASON, + ) + } + for (const question of this.questions.values()) { + question.reject(new Error('Session closed.')) + } + this.questions.clear() + await Promise.allSettled( + running.map((operation) => Promise.resolve(operation)), + ) + // The events from before the close land first, so the next host has + // the text of a cut answer. + if (options?.recoverable) await this.writer?.flush().catch(() => {}) + this.lifetime.abort() + try { + await this.sessionPlugins?.dispose() + } finally { + this.feed.close() + this.onClose() + } + })() + return this.closing + } + + // =========================== + // Chat turns + // =========================== + + protected createTurnOperation(id?: string): OperationImpl { + const operation = new OperationImpl( + 'chat', + this.feed, + (target) => this.cancel(target.id), + undefined, + id, + ) + this.operations.set(operation.id, operation as OperationImpl) + return operation + } + + /** + * Steers a turn never reached run next, before other queued turns (after + * a resolve of that turn). A steer with an abort request settles `aborted` + * instead. + */ + protected requeueWaitingSteers(): void { + this.joining.clear() + const waiting = this.steerQueue.splice(0) + const aborted = waiting.filter((steer) => + this.isAbortRequested(steer.inputId), + ) + const rest = waiting.filter((steer) => !aborted.includes(steer)) + for (const steer of aborted) { + this.abortedSteers.delete(steer.inputId) + const { operation } = steer + // Not awaited: the session goes on to the next turn at once. The + // settlement lands before the operation ends. + void this.settle({ + inputId: steer.inputId, + outcome: 'aborted', + ...(operation ? { operationId: operation.id } : {}), + }) + .catch(() => {}) + .then(() => { + if (!operation) return + operation.fail('cancelled', new Error('Cancelled before it started.')) + this.publishFinished(operation) + }) + } + if (this.closing) { + // close() ended the queued turns already. These end the same way. + for (const steer of rest) { + steer.operation?.fail('cancelled', new Error('Session closed.')) + } + return + } + this.queue.splice( + this.frontOfQueue(), + 0, + ...rest.map((steer) => this.steerTurn(steer)), + ) + } + + /** A waiting steer as its own turn, with its own operation. */ + protected steerTurn(steer: WaitingSteer) { + const operation = steer.operation ?? this.createTurnOperation() + this.bindTurn(steer.inputId, operation) + const turn: QueuedTurn = { ...steer, operation } + return turn + } + + /** + * The operation for a prompt whose id is known: the live one, one that + * settles from the log, or a refused one for another payload. + */ + protected knownTurn( + inputId: string, + input: HarnessInput, + principal: Principal | undefined, + ) { + const payload = inputKey(input, principal) + const local = this.admitted.get(inputId) + const logged = this.writer?.state.inputs.get(inputId) + const stored = + local ?? (logged ? inputKey(logged.input, logged.principal) : undefined) + if (stored === undefined) return undefined + if (stored !== payload) { + const refused = this.createTurnOperation() + this.refuse(refused, { inputId, status: 'rejected', reason: 'conflict' }) + return refused + } + const live = this.turnOperations.get(inputId) + if (live) return live + const replay = this.createTurnOperation(logged?.operationId) + this.answerDuplicate(replay, inputId, 'duplicate') + return replay + } + + /** Settle `operation` like the earlier input with the same id. */ + protected answerDuplicate( + operation: OperationImpl, + inputId: string, + admission: 'duplicate' | 'conflict', + ) { + const receipt = this.duplicateReceipt(inputId, admission) + if (receipt.status === 'rejected') { + this.refuse(operation, receipt) + return + } + operation.resolveReceipt(receipt) + if (!this.writer) { + // Without a log, the result of an input from before a restart is gone. + operation.fail( + 'failed', + new Error( + 'This input ran before a restart. Keep its result with a durable host (stores.log).', + ), + ) + return + } + void this.settled(inputId).then( + (settlement) => this.settleOperation(operation, settlement), + (error: unknown) => operation.fail('failed', error), + ) + } + + /** + * Queue a reset: at the front for a new one (it applies when the running + * turn ends), at the end for one from before a restart (in its admission + * order). Its operation is a command, so clients do not see a chat turn. + */ + protected queueReset( + inputId: string, + note: string | undefined, + principal: Principal | undefined, + at: 'front' | 'end', + ) { + const operation = new OperationImpl( + 'command', + this.feed, + (target) => this.cancel(target.id), + ) + this.operations.set(operation.id, operation as OperationImpl) + this.bindTurn(inputId, operation) + const turn: QueuedTurn = { + operation, + inputId, + principal, + reset: note !== undefined ? { note } : {}, + } + this.queue.splice( + at === 'front' ? this.frontOfQueue() : this.queue.length, + 0, + turn, + ) + this.drain() + return operation + } +} diff --git a/packages/ai-harness/src/session/state.ts b/packages/ai-harness/src/session/state.ts new file mode 100644 index 0000000000..714eaa3717 --- /dev/null +++ b/packages/ai-harness/src/session/state.ts @@ -0,0 +1,966 @@ +import { + EventType, + convertSchemaToJsonSchema, + validateWithStandardSchema, +} from '@tanstack/ai' +import { checkConfigValue } from '../config' +import { isRecord } from '../utils' +import { changedEntry } from '../session-entry' +import { mediaIdOf, mediaOfMessage, mediaPart } from '../media-ref' +import { + BackgroundJobs, + CapabilityValues, + RevertFiles, + RevertStanding, + SessionMetadata, + SessionSignal, +} from '../plugins' +import { HARNESS_EVENTS } from '../types' +import { isSessionUsage } from '../usage' +import { + INTERRUPTED, + USAGE, + isInterruptedTurn, + createInputId, + customEvent, + tanstackOf, + toolCallOf, + SETTINGS, + changeSettings, +} from './helpers' +import { SessionInputs } from './inputs' +import type { OperationImpl } from '../operation' +import type { + Interrupt, + ModelMessage, + RunAgentResumeItem, + SchemaInput, +} from '@tanstack/ai' +import type { + BlobBody, + BlobRange, + SessionIndexEntry, +} from '@tanstack/ai-persistence' +import type { RevertState } from '../log' +import type { Question } from '../commands' +import type { ConfigOption } from '../config' +import type { AnyHarness, SessionMetadataOf } from '../define' +import type { SessionEntryChange, SessionEntryUpdate } from '../session-entry' +import type { + ChatTurnResult, + MediaRecord, + Principal, + Receipt, + ThreadSettings, + ThreadSettingsChange, +} from '../types' +import type { + SessionSnapshot, + SessionDescription, + InterruptedTurn, +} from './types' +import type { SettingNames } from './helpers' + +/** + * A layer of {@link HarnessSession}: config, thread settings, revert, + * interrupts, media, and the snapshot. + */ +export abstract class SessionState< + THarness extends AnyHarness = AnyHarness, +> extends SessionInputs { + /** The stores that first-party plugins read as capabilities. */ + protected storeCapabilities() { + const values = new CapabilityValues() + const { metadata } = this.persistence.stores + if (metadata) values.provide(SessionMetadata, metadata) + values.provide(RevertStanding, () => this.reverted !== undefined) + values.provide(SessionSignal, this.lifetime.signal) + values.provide(BackgroundJobs, { + started: (jobId) => this.logJob(jobId, false), + ended: (jobId) => this.logJob(jobId, true), + }) + return values + } + + /** + * Write the index entry of this thread on open: a new entry the first + * time, else a new `updatedAt`. The owner is who opened the thread first, + * so an open by another user does not change it. + */ + protected async openEntry() { + const now = Date.now() + await this.index.update(this.threadId, (entry) => ({ + threadId: this.threadId, + createdAt: now, + ...entry, + harness: entry?.harness ?? this.harness.name, + updatedAt: now, + ...storedPrincipal(entry?.principal ?? this.principal), + })) + } + + snapshot(): SessionSnapshot { + const active = [...this.operations.values()].filter( + (operation) => !operation.isSettled(), + ) + return { + threadId: this.threadId, + status: this.activeTurn + ? 'running' + : this.interrupted + ? 'requires_action' + : 'idle', + activeOperations: active.map((operation) => ({ + id: operation.id, + kind: operation.kind, + ...(operation.agent ? { agent: operation.agent } : {}), + startedCursor: operation.startedCursor, + })), + // A steer that waits to join the running turn waits too. + queuedTurns: this.queue.length + this.steerQueue.length, + waitingInputs: this.inputs(), + pendingInterrupts: + this.activeResume && !this.activeTurn?.isSettled() + ? [] + : (this.interrupted?.interrupts ?? []), + pendingQuestions: [...this.questions.entries()].map( + ([questionId, question]) => ({ + questionId, + message: question.message, + ...(question.schema + ? { schema: convertSchemaToJsonSchema(question.schema) } + : {}), + ...(question.secret ? { secret: true } : {}), + ...(question.url ? { url: question.url } : {}), + }), + ), + plugins: { ...this.stateDoc }, + usage: this.usage(), + cursor: this.feed.head(), + } + } + + // =========================== + // Media + // =========================== + + /** + * Store a file in the media store of this thread. Send it to a turn with + * `mediaPart(record)`. Throws a `MediaError`: 413 when the file is bigger + * than `media.maxBytes`, 415 for a type or kind the harness does not take. + * + * @example + * const record = await session.putMedia(bytes, { mimeType: 'image/png', name: 'cat.png' }) + * session.prompt([{ type: 'text', content: 'What is this?' }, mediaPart(record)]) + */ + putMedia(body: BlobBody, info: { mimeType: string; name: string }) { + return this.mediaStore.put(body, info) + } + + /** The record of a media file of this thread, or `null` when it is not found. */ + getMedia(id: string) { + return this.mediaStore.get(id) + } + + /** + * The bytes of a media file, or of one `range` of them. Throws a + * `MediaError` with 404 when it is not found. + */ + loadMedia(id: string, range?: BlobRange) { + return this.mediaStore.load(id, range) + } + + /** + * A URL for a media file, for ``, `