diff --git a/.changeset/ai-harness-cli.md b/.changeset/ai-harness-cli.md index 122cb06865..8569cf66b8 100644 --- a/.changeset/ai-harness-cli.md +++ b/.changeset/ai-harness-cli.md @@ -8,4 +8,4 @@ Add `@tanstack/ai-harness-cli`: `runCli(harness)` runs a harness from the termin - `-p "prompt"` runs one prompt for scripts and CI, and `--output ndjson` prints every AG-UI event. - `--serve` serves the session protocol over HTTP behind a bearer token. - `--acp`, `--mcp`, and `--dashboard` load `@tanstack/ai-acp`, `@tanstack/ai-mcp`, and `@tanstack/ai-dashboard` only when you use the flag. A missing package, or one without harness support, fails that flag with a clear message. -- `host` and `principal` options let the CLI use your durable host and run as your user. +- `host` and `principal` options let the CLI use your durable host and run as your user. Every mode except `--dashboard` opens its sessions as that principal. On a host with `access`, `--serve` reads the user of each request with `identify`, and `--mcp` does not start, because stdio has no login. diff --git a/.changeset/ai-harness-core.md b/.changeset/ai-harness-core.md index 4519636ee5..7eaa5b371c 100644 --- a/.changeset/ai-harness-core.md +++ b/.changeset/ai-harness-core.md @@ -6,6 +6,10 @@ Add `@tanstack/ai-harness`: build an agent harness on TanStack AI. A harness kee - `defineHarness()` names an agent: its adapter, tools, system prompts, middleware, agents, subagents, routing, plugins, turn hooks, and durability options. - `createHarnessHost()` opens sessions. A session takes `send` (a prompt, a steer into the running turn, a follow-up, or a message to an agent run with `to`), `continue`, `reset`, `resolve`, `cancel`, and `background` inputs. Each returns an operation with a receipt, and the session streams typed events. `host.isOpen(harness, { threadId })` says whether this host holds that session open. +- `createHarnessHost({ access })` decides who can use each session. `identify({ request })` reads the user of a request, and `authorize({ principal, action, session, run })` decides each action. The principal type comes from what `identify` returns. The host keeps each tenant apart, needs a `principal` on each call, and `host.open` returns a handle that acts only as its principal. A refused action throws `AccessDeniedError`. Without `access`, one user runs the host, with no checks. With `access`, keyed adapters and provider keys read env vars only when `access.envKeys` is `true`. +- `membersOnly()` is a ready `authorize` with the roles owner, member, and viewer. `host.sessions.share` and `host.sessions.unshare` change the members, and `host.sessions.fork` takes `members: 'self' | 'all'`. For apps with their own access tables, `host.revoke` takes the access of a user to a thread. `host.identify` reads the user of a request for your own routes. Each `host.sessions` method takes one object. +- A principal with a `tenantId` keeps its sessions in its tenant: one thread id in two tenants is two sessions, with two index entries and two logs. `logIdOf({ tenantId, logId })` gives the id of a log in the store. +- Every run has a run context with its user, and that user's credentials and provider keys: `ctx.context.run` in tools and middleware (type them with `HarnessRunContext`), `ctx.turn.run` in a turn plugin, and `run` in a command. The `ctx.credentials` and `ctx.keys` of a session plugin have no user, so each call throws. `ctx.agents.start(agent, input, { principal })` starts an agent for a user. The log keeps the sender of every input. - `definePlugin()` adds tools, prompts, middleware, agents, commands, config options, state, and extension points to a session or to one turn. - Agents run from code with `session.agents..run()` and `.start()`. In agent code, `ctx.harness` gives `step.do()` (a finished step returns its stored value when the run runs again) and `agents.start()` (a child that counts against the subagent budget). - With `stores.log` from `@tanstack/ai-persistence`, a session is durable: inputs, tool results, and steps are in the log, and the next host recovers a turn that a crash stopped. `durableTool()` runs side effects once. `durability.truncatedToolResult` sets the text of a tool call cut at the output limit. Without it, a cut answer ends the turn. diff --git a/.changeset/ai-harness-transport.md b/.changeset/ai-harness-transport.md index 7115dcda27..93f73d6597 100644 --- a/.changeset/ai-harness-transport.md +++ b/.changeset/ai-harness-transport.md @@ -4,8 +4,8 @@ Serve a harness session to clients. -- `createHarnessHandler()` from `@tanstack/ai-harness/server` serves a session over HTTP. `POST run` is standard AG-UI, so `useChat` with `fetchServerSentEvents` works on a harness, with approvals, client tools, and a reload that joins the running turn. `POST control`, `GET events`, `GET transcript`, `GET agent-transcript` (the messages of a child agent run), `GET snapshot`, `GET describe`, `sessions`, and the media routes cover the rest. -- `handleHarnessSocket()`, from the same entry, serves the same session over a WebSocket. +- `createHarnessHandler()` from `@tanstack/ai-harness/server` serves a session over HTTP, as the user that the host `access` reads from each request. `POST run` is standard AG-UI, so `useChat` with `fetchServerSentEvents` works on a harness, with approvals, client tools, and a reload that joins the running turn. `POST control`, `GET events`, `GET transcript`, `GET agent-transcript` (the messages of a child agent run), `GET snapshot`, and `GET describe` cover one session. `sessions` (with the `share` and `unshare` ops), `GET host-events`, and the media routes cover the rest. A refused action answers 403. A read of a thread that does not exist answers 404 and starts no session. +- `handleHarnessSocket()`, from the same entry, serves the same session over a WebSocket. It takes the upgrade `request` and reads the user from it. - `expose` on `defineHarness()` names the agents, settings, config keys, and commands that a client can use. Any other input of those kinds gets `rejected` with `not_exposed`. - `createHarnessClient()` from `@tanstack/ai-harness/client` is a typed, browser-safe client. It sends inputs, reads the session, and streams events that resume from a cursor. -- `oauthConnector()` and the OAuth helpers from `@tanstack/ai-harness/oauth` (`loopbackLogin`, `deviceLogin`, `exchangeCode`, `refreshCredential`, and more) keep a credential per sender. `token({ wait: true })` stops the turn with an `auth_required` interrupt until the user signs in. +- `oauthConnector()` and the OAuth helpers from `@tanstack/ai-harness/oauth` (`loopbackLogin`, `deviceLogin`, `exchangeCode`, `refreshCredential`, and more) keep a credential per user. A tool calls `token({ run: ctx.context.run })` for the token of the user of its run. `token({ run, wait: true })` stops the turn with an `auth_required` interrupt until the user signs in. diff --git a/.changeset/ai-mcp-harness.md b/.changeset/ai-mcp-harness.md index e3f5f39956..fa23fd493d 100644 --- a/.changeset/ai-mcp-harness.md +++ b/.changeset/ai-mcp-harness.md @@ -5,6 +5,7 @@ Use MCP in a TanStack AI harness. - `@tanstack/ai-mcp/harness` adds `createHarnessMcpServer`. It serves a harness as an MCP server, so Claude Code, Claude Desktop, Cursor, or another agent can use it. The tools are `chat`, `steer`, `cancel`, `approve`, `reject`, `resolve`, `answer`, and `status`, plus `agent_` for each exposed agent and `command_` for each exposed command. Every tool takes an optional `threadId`. With `approvals: 'ask'` (the default), the client asks the user about each approval when it supports elicitation. Otherwise the approvals come back in the result. `approvals: 'auto'` approves every tool call. `resolve` answers every kind of interrupt in one call. +- On a host with `access`, `createHarnessMcpServer` reads the user of each HTTP request with `identify`. The host checks each tool call of that user. A host without `access` needs `openAccess: true` (no user) or `openAccess: { principal }` (one local user), else the server throws at the start. - `@tanstack/ai-mcp/harness` also adds the `mcp()` harness plugin. It connects stdio and HTTP servers by name, gives the model their tools as `_`, and adds a `/mcp` status command. `oauth: true` signs in to a server with `/connect `. - `@tanstack/ai-mcp/connector` adds `mcpConnector`: one OAuth MCP server as a harness plugin, with `needsApproval`, `prefix`, `scopes`, and `clientName`. The sign-in passes the `iss` of the callback on (RFC 9207). - `mcp()` and `mcpConnector` answer the `elicitation/create` requests of MCP spec 2025 servers, and the input requests of spec 2026-07-28 servers, as session questions. One tool call can ask up to 5 times. diff --git a/.changeset/harness-acp-agent.md b/.changeset/harness-acp-agent.md index 8034ce646c..5797f4917b 100644 --- a/.changeset/harness-acp-agent.md +++ b/.changeset/harness-acp-agent.md @@ -2,4 +2,4 @@ '@tanstack/ai-acp': minor --- -`@tanstack/ai-acp/agent` adds `createAcpAgent` and `serveAcp`, which serve a harness as an ACP v2 agent (experimental, like the draft protocol). Tool approvals become permission requests, cancel cancels the running turn, and image, audio, and embedded file blocks go to the media store of the session. `session/list` lists the sessions of the harness from the session index of the host. `@tanstack/ai-acp` now uses `@agentclientprotocol/sdk` 1.8, pinned to `~1.8.0` because the agent uses its experimental v2 entry point. The existing ACP client adapters keep working. +`@tanstack/ai-acp/agent` adds `createAcpAgent` and `serveAcp`, which serve a harness as an ACP v2 agent (experimental, like the draft protocol). Tool approvals become permission requests, cancel cancels the running turn, and image, audio, and embedded file blocks go to the media store of the session. `session/list` lists the sessions of the harness from the session index of the host. The `principal` option is the user of the editor: each session opens as that user, and the list shows the sessions of that user. It is required on a host with `access`. `@tanstack/ai-acp` now uses `@agentclientprotocol/sdk` 1.8, pinned to `~1.8.0` because the agent uses its experimental v2 entry point. The existing ACP client adapters keep working. diff --git a/.changeset/harness-cli-handler-options.md b/.changeset/harness-cli-handler-options.md index fa70df563f..a53abd427c 100644 --- a/.changeset/harness-cli-handler-options.md +++ b/.changeset/harness-cli-handler-options.md @@ -2,4 +2,4 @@ '@tanstack/ai-harness-cli': patch --- -`--serve` passes `canAccess` and `media: { secret }` to the harness handler, which needs both. +`--serve` passes `media: { secret }` to the harness handler. On a host without `access`, it also passes `openAccess`: the `principal` of `runCli`, or `true` without one. So each request with the token runs as the user of the CLI. On a host with `access`, `identify` reads the user of each request. diff --git a/.changeset/harness-secure-defaults.md b/.changeset/harness-secure-defaults.md index 1cd5b98d5f..2bf23a5f05 100644 --- a/.changeset/harness-secure-defaults.md +++ b/.changeset/harness-secure-defaults.md @@ -4,6 +4,6 @@ Make the transport defaults safe for a server with many instances. -- `createHarnessHandler` and `handleHarnessSocket` need `canAccess`. To let every principal use every thread, pass `canAccess: () => true`. -- On a durable host (`stores.log`), `createHarnessHandler` throws without `media: { secret }`, so a media URL that one instance signs works on every instance. +- `createHarnessHandler` and `handleHarnessSocket` read the user of each request with the host `access`. On a host without `access`, pass `openAccess: true` (no user) or `openAccess: { principal }` (one local user), else they throw at the start. So a server never serves every thread to everyone by mistake. +- On a durable host (`stores.log`), `createHarnessHandler` throws without `media: { secret }`, so a media URL that one instance signs works on every instance. A signed URL names its user, and the handler checks that user again when it serves the file. - `oauthConnector` needs `login`. Use `'device'` on a server, and `'loopback'` only for a harness on the user's own machine, such as a CLI. diff --git a/.changeset/persistence-session-stores.md b/.changeset/persistence-session-stores.md index dc6eddfaf1..889ddce4c6 100644 --- a/.changeset/persistence-session-stores.md +++ b/.changeset/persistence-session-stores.md @@ -8,10 +8,10 @@ Add opt-in session stores. The chat and generation middleware work as before wit - `log` (`LogStore`): an append-only session log with 14 methods, each with one object argument. Records: `append`, `read`, and `subscribe`. Ownership: `acquire`, `renew`, `release`, `setAffinity`, and `listOrphaned`. The input lane: `offer` and `offered`, consumed in the same write with `append({ consumes })`. Snapshots and lifecycle: `getSnapshot`, `putSnapshot`, `truncate`, and `deleteLog`. Each fenced write carries the epoch from `acquire`, and a write with an old epoch rejects with `OwnershipLostError`. A writer that never calls `acquire` uses `UNOWNED_EPOCH`. `append` rejects with `LogConflictError` when another writer took the position, and `read` rejects with `LogTruncatedError` before the truncation point. `memoryLogStore()` is the reference store. It is not part of `memoryPersistence()`. `memoryLogBackend()` lets two store instances share one log. The conformance kit needs `twin`, a second store set on the same backend, for the `log` cases across instances. - `leases` (`LeaseStore`, `TurnLease`, `TurnLeaseKey`): turn leases from a system that already has them. `defineLeaseStore()` types a custom one, `memoryLeaseStore()` is the reference store, and the conformance testkit covers it. It is not part of `memoryPersistence()`. - `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. +- `sessions` (`SessionIndexStore`): one index entry per tenant and thread. The key is `{ tenantId?, threadId }` (`SessionIndexKey`), so one thread id in two tenants is two entries. An entry has an `ownerId` and `members` (`SessionMember`, with the role `owner`, `member`, or `viewer`). The store has `get(key)`, `getVersioned(key)`, `upsert(entry)`, `upsertIf({ entry, revision })` for compare-and-set writes, `list`, and `delete(key)`, a hard delete. `list` reads one tenant, filters by member, parent thread, title `search`, `harness` name, and exact `metadata` values, and pages with a cursor. `defineInboxStore`, `defineLogStore`, `defineWorkClaimStore`, and `defineSessionIndexStore` type an adapter. `memoryPersistence()` now also has `inbox`, `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. +With `stores.sessions`, `withPersistence` writes one index entry for each subagent child (`subagent:`), in the tenant of its `tenantId` option. The entry has the parent thread, the tool call, and the `ownerId` and `harness` name of the parent entry. It has no members: access to a child goes through its top-level thread. `loadChildTranscript(messages, subagentRunId, threadId)` reads the stored messages of a child agent run for a UI. It returns `undefined` for a child of another thread. diff --git a/docs/config.json b/docs/config.json index 57dfafcd41..6bfa335aee 100644 --- a/docs/config.json +++ b/docs/config.json @@ -761,7 +761,7 @@ "label": "Store Reference", "to": "persistence/store-reference", "addedAt": "2026-08-04", - "updatedAt": "2026-10-16" + "updatedAt": "2026-10-17" }, { "label": "How Persistence Works", @@ -927,7 +927,7 @@ "label": "Build your first harness", "to": "harness/overview", "addedAt": "2026-10-13", - "updatedAt": "2026-10-16" + "updatedAt": "2026-10-17" }, { "label": "Send and show media", @@ -939,23 +939,30 @@ "label": "Connect clients", "to": "harness/connect", "addedAt": "2026-10-13", - "updatedAt": "2026-10-16" + "updatedAt": "2026-10-17" + }, + { + "label": "Access and tenants", + "to": "harness/access", + "addedAt": "2026-10-17" }, { "label": "Run in the terminal", "to": "harness/cli", - "addedAt": "2026-10-13" + "addedAt": "2026-10-13", + "updatedAt": "2026-10-17" }, { "label": "Deploy a harness", "to": "harness/deploy", - "addedAt": "2026-10-13" + "addedAt": "2026-10-13", + "updatedAt": "2026-10-17" }, { "label": "Send inputs safely", "to": "harness/inputs", "addedAt": "2026-10-13", - "updatedAt": "2026-10-16" + "updatedAt": "2026-10-17" }, { "label": "Run a tool without a model", @@ -972,7 +979,7 @@ "label": "List, rename, and fork sessions", "to": "harness/sessions", "addedAt": "2026-10-13", - "updatedAt": "2026-10-16" + "updatedAt": "2026-10-17" }, { "label": "Run side effects once", @@ -983,7 +990,7 @@ "label": "Build on the session log", "to": "harness/session-log", "addedAt": "2026-10-13", - "updatedAt": "2026-10-16" + "updatedAt": "2026-10-17" }, { "label": "Control how a turn ends", @@ -1005,7 +1012,7 @@ "label": "Fork and reset a thread", "to": "harness/fork-and-reset", "addedAt": "2026-10-13", - "updatedAt": "2026-10-16" + "updatedAt": "2026-10-17" }, { "label": "Share one log between sessions", @@ -1017,13 +1024,13 @@ "label": "Write a plugin", "to": "harness/plugins", "addedAt": "2026-10-13", - "updatedAt": "2026-10-16" + "updatedAt": "2026-10-17" }, { "label": "Track usage and cost", "to": "harness/usage", "addedAt": "2026-10-13", - "updatedAt": "2026-10-16" + "updatedAt": "2026-10-17" }, { "label": "Ask before risky tool calls", @@ -1070,25 +1077,25 @@ "label": "Auth and connectors", "to": "harness/auth", "addedAt": "2026-10-13", - "updatedAt": "2026-10-16" + "updatedAt": "2026-10-17" }, { "label": "Share a thread between people", "to": "harness/shared-threads", "addedAt": "2026-10-13", - "updatedAt": "2026-10-16" + "updatedAt": "2026-10-17" }, { "label": "Connect model providers", "to": "harness/provider-keys", "addedAt": "2026-10-13", - "updatedAt": "2026-10-16" + "updatedAt": "2026-10-17" }, { "label": "Build your own UI", "to": "harness/custom-ui", "addedAt": "2026-10-13", - "updatedAt": "2026-10-16" + "updatedAt": "2026-10-17" }, { "label": "Code mode in a harness", @@ -1104,7 +1111,7 @@ "label": "Run agents from a harness", "to": "harness/subagents", "addedAt": "2026-10-13", - "updatedAt": "2026-10-16" + "updatedAt": "2026-10-17" }, { "label": "Add skills", @@ -1121,12 +1128,13 @@ "label": "Use from any MCP client", "to": "harness/mcp-server", "addedAt": "2026-10-13", - "updatedAt": "2026-10-16" + "updatedAt": "2026-10-17" }, { "label": "Use a harness from an editor (ACP)", "to": "harness/acp", - "addedAt": "2026-10-13" + "addedAt": "2026-10-13", + "updatedAt": "2026-10-17" } ] }, diff --git a/docs/harness/access.md b/docs/harness/access.md new file mode 100644 index 0000000000..5d212d92a1 --- /dev/null +++ b/docs/harness/access.md @@ -0,0 +1,238 @@ +--- +title: Access and tenants +id: harness-access +order: 3 +description: "Decide who can use each session of a harness. Read the user of each request, check each action, and keep the threads of each organization apart." +keywords: + - tanstack ai + - harness + - access + - identify + - authorize + - tenantId + - membersOnly +--- + +Your harness serves many users. Ada must not read the threads of Bob, and the company Acme must never see a thread of the company Globex. Give the host `access`. Then the host reads the user of each request, asks you about each action, and keeps each organization apart. + +## 1. Turn on access + +Give `createHarnessHost` an `access` object and a `sessions` store. The store keeps the owner and the members of each session: + +```ts group=harness-access +import { createHarnessHost, defineHarness, membersOnly } from '@tanstack/ai-harness' +import { createHarnessHandler } from '@tanstack/ai-harness/server' +import { memoryPersistence } from '@tanstack/ai-persistence' +import { openaiText } from '@tanstack/ai-openai' + +export const assistant = defineHarness({ + name: 'acme/assistant', + adapter: openaiText('gpt-6.1-sol'), +}) + +// Your own login check: a session token to a user. +const users = new Map([ + ['Bearer token-ada', { id: 'ada', tenantId: 'acme' }], + ['Bearer token-gus', { id: 'gus', tenantId: 'globex' }], +]) + +export const host = createHarnessHost({ + persistence: memoryPersistence(), + access: { + identify: ({ request }) => users.get(request.headers.get('authorization') ?? '') ?? null, + authorize: membersOnly(), + }, +}) + +export const handler = createHarnessHandler({ host, harness: assistant }) +``` + +- `identify` reads the user of a request with your own login check. Return `null` to refuse the request with `401`. +- `authorize` decides each action. `membersOnly()` lets only the members of a session use it. See [Share a thread between people](./shared-threads). +- `memoryPersistence()` has a `sessions` store. Without `stores.sessions`, `createHarnessHost` throws. + +The handler gets the user from `identify`, so it needs no login option of its own. A new thread belongs to the user who opens it first. + +## 2. Write your own rule + +`authorize` gets the user, the action, and the session. Return `false` to refuse. The user has the type that `identify` returns, so your own fields are typed too: + +```ts group=harness-access-rule +import { createHarnessHost } from '@tanstack/ai-harness' +import { memoryPersistence } from '@tanstack/ai-persistence' + +interface User { + id: string + tenantId: string + isAdmin: boolean +} + +const users = new Map([ + ['Bearer token-ada', { id: 'ada', tenantId: 'acme', isAdmin: false }], + ['Bearer token-eve', { id: 'eve', tenantId: 'acme', isAdmin: true }], +]) + +export const host = createHarnessHost({ + persistence: memoryPersistence(), + access: { + identify: ({ request }) => users.get(request.headers.get('authorization') ?? '') ?? null, + authorize: ({ principal, action, session }) => { + const isOwner = session.ownerId === principal.id + if (isOwner || principal.isAdmin) return true + // Everyone in the organization can read every thread of it. + return action === 'read' || action === 'list' + }, + }, +}) +``` + +Each action covers these calls: + +| Action | What it covers | +| --- | --- | +| `read` | Open the session, and read its stream, snapshot, transcript, media, and child runs. | +| `list` | See the session in `host.sessions.list`. | +| `send` | Prompt, steer, follow up, run a command or a tool, start an agent, and upload media. | +| `answer` | Answer a question or an approval. | +| `cancel` | Stop a run, or drop an input that waits. | +| `configure` | Change the thread settings, the plugin config, the title, and the metadata. | +| `share` | Add or remove members. | +| `delete` | Remove the session from the index. | +| `fork` | Copy the session into a new session. | +| `transfer` | Give the session to a new owner. | + +- For `answer` and `cancel`, `authorize` also gets `run`: the `runId` and the `principalId` of the user who started that run. +- A handle checks `read` once, when it opens. Each action that changes the session asks `authorize` again. + +## What the host checks for you + +`authorize` decides only the rule of your app. The host does these checks before it calls it: + +- The tenant. A user reaches only the sessions of their own tenant. `authorize` never gets a session of another tenant. +- A user on every call. With access, `host.open`, `host.fork`, and each `host.sessions` method need a `principal`. A call without one throws. +- One user per handle. `host.open` returns a session that acts only as its `principal`. A call on it that passes another `principal` throws. +- Child threads. The host checks a subagent child as its root thread, so `session` is the root session. +- Missing threads. A read of a thread with no entry answers `404` and starts no session. `GET run?threadId=` answers an empty hydration instead, because `useChat` reads it before the first message. + +A refused action throws `AccessDeniedError`, with `code: 'forbidden'`, the `action`, and the `threadId`. The handler answers it with `403`. + +The host keeps only the `id` and the `tenantId` of a user, in the log and in signed media URLs. So `authorize` can get a user with only these two fields, for example when the handler serves a signed media URL. A rule that reads another field, such as `isAdmin` above, must refuse when the field is missing. + +## Keep tenants apart + +Give each user a `tenantId` in `identify`. The tenant keeps each organization apart: + +- One `threadId` in two tenants is two sessions, with two index entries and two logs. +- `host.sessions.list` reads one tenant: the tenant of `principal`. A user with no `tenantId` is in the default tenant. +- The stores keep each thread under `logIdOf({ tenantId, logId })`, for example `acme:thread-1`. The default tenant gives `:thread-1`. Use `logIdOf` when you read a log yourself. +- The credentials of a user are kept per tenant. See [Keep organizations apart](./shared-threads#7-keep-organizations-apart). + +## Call the host from code with no request + +A cron job, an admin tool, or a queue worker has no request, so `identify` cannot run. Pass the user yourself. The host checks the action for that user, the same as for a request: + +```ts group=harness-access +// Your database knows the owner of each thread. +export async function nightlyReport(threadId: string, owner: { id: string; tenantId: string }) { + const session = await host.open(assistant, { threadId, principal: owner }) + await session.send('Write the report of today.') +} +``` + +Your own routes and an MCP server have a request. Read its user with `host.identify`: + +```ts group=harness-access +export async function myThreads(request: Request) { + const principal = await host.identify({ request }) + if (!principal) return new Response('Unauthorized', { status: 401 }) + const page = await host.sessions.list({ principal, limit: 20 }) + return Response.json(page.entries.map((entry) => entry.title ?? 'New chat')) +} +``` + +`host.identify` throws on a host with no `access`. + +## Use your own access tables + +Your app already keeps who can use each thread, for example a table of teams. Read it in `authorize`. When you take a user out of the table, call `host.revoke`: + +```ts group=harness-access-own-tables +import { createHarnessHost } from '@tanstack/ai-harness' +import { memoryPersistence } from '@tanstack/ai-persistence' + +const users = new Map([ + ['Bearer token-ada', { id: 'ada', tenantId: 'acme' }], + ['Bearer token-bob', { id: 'bob', tenantId: 'acme' }], +]) + +// Your own table: the team of each thread. +const teams = new Map([['plan-q3', new Set(['ada', 'bob'])]]) + +export const host = createHarnessHost({ + persistence: memoryPersistence(), + access: { + identify: ({ request }) => users.get(request.headers.get('authorization') ?? '') ?? null, + authorize: ({ principal, session }) => + session.ownerId === principal.id || (teams.get(session.threadId)?.has(principal.id) ?? false), + }, +}) + +export function removeFromTeam(threadId: string, user: { id: string; tenantId: string }) { + teams.get(threadId)?.delete(user.id) + host.revoke({ threadId, principal: user }) +} +``` + +After `host.revoke`: + +- The streams of the user on the thread close. A socket gets a `forbidden` error frame. +- The open handles of the user on the thread throw `AccessDeniedError`. A new `host.open` asks `authorize` again. +- The signed media URLs of the user stop working. A signed URL names its user, and the handler checks that user again when it serves the file. +- `host.events()` yields a `revoked` event. + +`host.sessions.share` and `unshare` work only with `membersOnly()`. With your own tables, change the tables, then call `host.revoke`. + +## Serve with no login + +A demo, or a tool on your own machine, has no login. Give the handler `openAccess` on a host with no `access`: + +```ts group=harness-access-open +import { createHarnessHost, defineHarness } from '@tanstack/ai-harness' +import { createHarnessHandler } from '@tanstack/ai-harness/server' +import { memoryPersistence } from '@tanstack/ai-persistence' +import { openaiText } from '@tanstack/ai-openai' + +const assistant = defineHarness({ + name: 'acme/assistant', + adapter: openaiText('gpt-6.1-sol'), +}) + +const host = createHarnessHost({ persistence: memoryPersistence() }) + +// One user on this machine: every request runs as Ada. +export const handler = createHarnessHandler({ + host, + harness: assistant, + openAccess: { principal: { id: 'ada' } }, +}) +``` + +- `openAccess: true`: every request runs with no user, in the default tenant. +- `openAccess: { principal }`: every request runs as that one user. Their credentials and keys do the work. +- Without host `access` and without `openAccess`, `createHarnessHandler` throws at the start. So a server never serves every thread to everyone by mistake. +- On a host with `access`, `openAccess` throws. The host reads the user itself. +- `handleHarnessSocket` takes the same option. + +CAUTION: Do not use `openAccess` on a server that the internet can reach. Every caller can then use every thread. + +## Use the provider keys of the server + +With access, a keyed adapter and `run.keys` read only the keys that each user saved. They do not read the env vars of the server, so a user never spends the key of the server. To let users fall back to the server key, set `envKeys: true` in `access`, next to `identify` and `authorize`. See [Connect model providers](./provider-keys). + +## What you have now + +- A host that reads the user of each request, and checks each action. +- A rule of your own, typed from what `identify` returns. +- Organizations that never see each other's threads. +- Cron jobs and admin tools that act as a user they name. +- A way to take access away from a user at once. diff --git a/docs/harness/acp.md b/docs/harness/acp.md index e87958c313..82e20882ad 100644 --- a/docs/harness/acp.md +++ b/docs/harness/acp.md @@ -49,6 +49,7 @@ serveAcp({ host, harness: assistant, version: '1.0.0' }) ``` - Each ACP session is one harness session. The ACP session id is the thread id. +- `principal` is the user of the editor, for example `serveAcp({ host, harness: assistant, principal: { id: 'ada' } })`. Every session opens as that user, so `/connect` saves keys for them. On a host with [`access`](./access), `principal` is required. - The prompts of a session run one after another. The harness `busy` setting does not apply. - `version` is the version that `initialize` reports. Default: `'0.0.0'`. - `cwd` is the working directory that the session list shows for each session. Default: `process.cwd()`. @@ -70,9 +71,9 @@ Give the editor the command that runs the script, for example `node agent.js`. T The editor lists the sessions of the harness, newest first. The user picks one, and the next prompt continues that thread. -- The list shows the top-level sessions of the harness. Sessions of child agents stay out. +- The list shows the top-level sessions of the harness that `principal` can list. Sessions of child agents stay out. - The list comes from the session index of the host. Use a persistent store, so that the list keeps the sessions after the process stops. -- A session id that the session index does not know fails with a not found error. +- A session id that the session index does not know, or that `principal` cannot read, fails with a not found error. ## Known limits diff --git a/docs/harness/auth.md b/docs/harness/auth.md index 45c2da95c4..7efe944084 100644 --- a/docs/harness/auth.md +++ b/docs/harness/auth.md @@ -19,6 +19,7 @@ Your agent needs to open a pull request, so it needs the user's GitHub token. Th import { toolDefinition } from '@tanstack/ai' import { defineHarness, oauthConnector } from '@tanstack/ai-harness' import { openaiText } from '@tanstack/ai-openai' +import type { HarnessRunContext } from '@tanstack/ai-harness' const github = oauthConnector({ id: 'github', @@ -33,10 +34,10 @@ const github = oauthConnector({ scopes: ['repo'], }, tools: (token) => [ - toolDefinition({ name: 'list_issues', description: 'List my open issues' }).server( - async () => { + toolDefinition({ name: 'list_issues', description: 'List my open issues' }).server( + async (_args, ctx) => { const response = await fetch('https://api.github.com/issues', { - headers: { authorization: `Bearer ${await token()}` }, + headers: { authorization: `Bearer ${await token({ run: ctx.context.run })}` }, }) return response.json() }, @@ -51,7 +52,7 @@ export const assistant = defineHarness({ }) ``` -`token()` returns a fresh access token. It refreshes an expired token when the service gave a refresh token. +`token({ run })` returns a fresh access token of the user of that run. It refreshes an expired token when the service gave a refresh token. Pass the `run` of the tool call, `ctx.context.run`, so each user gets their own token. See [Act as the user of a run](./plugins#act-as-the-user-of-a-run). ## 2. Sign in @@ -77,12 +78,13 @@ export function showSignIns(session: HarnessSession) { ## Wait for the sign-in -A turn from a background agent or an email runs while nobody watches it. That turn can need a sign-in. Without a sign-in, the tool fails, and the work of that turn stops. Pass `wait: true` to `token()`. Then the turn waits for the sign-in and continues after it. +A turn from a background agent or an email runs while nobody watches it. That turn can need a sign-in. Without a sign-in, the tool fails, and the work of that turn stops. Pass `wait: true` to `token`. Then the turn waits for the sign-in and continues after it. ```ts group=harness-auth-wait import { toolDefinition } from '@tanstack/ai' import { defineHarness, oauthConnector } from '@tanstack/ai-harness' import { openaiText } from '@tanstack/ai-openai' +import type { HarnessRunContext } from '@tanstack/ai-harness' const github = oauthConnector({ id: 'github', @@ -97,10 +99,10 @@ const github = oauthConnector({ scopes: ['repo'], }, tools: (token) => [ - toolDefinition({ name: 'list_issues', description: 'List my open issues' }).server( - async () => { + toolDefinition({ name: 'list_issues', description: 'List my open issues' }).server( + async (_args, ctx) => { const response = await fetch('https://api.github.com/issues', { - headers: { authorization: `Bearer ${await token({ wait: true })}` }, + headers: { authorization: `Bearer ${await token({ run: ctx.context.run, wait: true })}` }, }) return response.json() }, @@ -117,17 +119,17 @@ export const assistant = defineHarness({ }) ``` -In your own plugin, use `ctx.credentials.require('github', { wait: true })`. +In a tool of your own plugin, use `ctx.context.run.credentials.require('github', { wait: true })`. When the credential is missing, the turn stops with an interrupt: - The interrupt has `reason: 'auth_required'`. Its `metadata['tanstack:interruptPayload'].request.connector` is `'github'`. -- The sign-in belongs to the sender of the stopped turn. When that person saves the credential, the same turn continues as that person, and the tool runs again. The save can be `/connect github`, the `connect:github` command, or `ctx.credentials.set`. +- The sign-in belongs to the sender of the stopped turn. When that person saves the credential, the same turn continues as that person, and the tool runs again. The save can be `/connect github`, the `connect:github` command, or `run.credentials.set` in a command of your own. - A credential that another person saves never continues the turn. - To stop the wait, resolve the interrupt with `status: 'cancelled'`. Then the tool fails, the same as without `wait`. - On a durable host, the wait continues after a restart. - `wait` is for a tool of the running chat turn. Do not pass it in a command, a background agent, or a hook. The credentials of a command never wait. -- Provider keys (`ctx.keys.require`) do not wait. +- Provider keys (`run.keys.require`) do not wait. In a web app, the wait shows in `state.signIns` of the session view, also after a reload. When the user presses your connect button, run the connector command. Each sign-in then has the `url` and the `userCode` to show. A client can run only the commands in `expose.commands`, so the harness above exposes `connect:github` (see [Choose what clients can change](./connect#choose-what-clients-can-change)): @@ -150,7 +152,7 @@ After the sign-in, the turn finishes its work with no new message. ## Keep credentials -Pass a credential store to the host. The host reads credentials from it, keyed by the sender of the running turn. Without one, credentials live in memory until the process stops. When several people write in one thread, see [Share a thread between people](./shared-threads). +Pass a credential store to the host. The host reads credentials from it, keyed by the user of each run. Without one, credentials live in memory until the process stops. When several people write in one thread, see [Share a thread between people](./shared-threads). ```ts group=harness-auth import { defineCredentialStore } from '@tanstack/ai/byok' @@ -188,7 +190,13 @@ The store can be a table in your own database or a secrets manager. See the rule ## Read credentials in your own plugin -`ctx.credentials.require('github')` returns the credential, or stops with `auth_required` when the user has not signed in. Plugins read only the credentials of the sender of the running turn. Outside a turn, they read the credentials of the person who opened the session. +Your plugin reads the credentials of the user of each run: + +- In a tool or a middleware hook, `ctx.context.run.credentials.require('github')` returns the credential. Before the user signs in, it stops with `auth_required`. +- In a command, use `run.credentials` of the command context. +- The `ctx.credentials` of a session plugin have no user, so each call throws. + +See [Act as the user of a run](./plugins#act-as-the-user-of-a-run). ## What you have now diff --git a/docs/harness/cli.md b/docs/harness/cli.md index 92eb7a56c6..3e89a48f77 100644 --- a/docs/harness/cli.md +++ b/docs/harness/cli.md @@ -79,7 +79,7 @@ Line mode reads one message or command per line and waits for each turn. It work - It binds to `127.0.0.1:8787`, reachable only from the same machine. Change the bind with `--host` and `--port`. `--host 0.0.0.0` puts it on the network, so do that only behind TLS and a proxy you trust. - The CLI signs media links with the token, so rotating the token stops any link a turn already handed out. -Need per-user access, your own sign-in check, or a rule for which threads each user can open? Use [`createHarnessHandler`](./connect) instead. It requires `authorize` and `canAccess`, so no thread is open by accident. +Need per-user access, your own sign-in check, or a rule for which threads each user can open? Use [`createHarnessHandler`](./connect) on a host with [`access`](./access) instead. The host reads the user of each request and checks each action, so no thread is open by accident. ## Use your own host and user @@ -109,9 +109,9 @@ await host.close() - You own the host, so the CLI does not close it. Close it yourself when `runCli` returns. Without `host`, the CLI builds a host and closes it. - Give `host` or `persistence`, not both. With both, `runCli` writes `Give runCli a host or persistence, not both.` and returns code 1. -- Line mode, `-p`, and your own `ui` open the session as `principal`, so `/connect` saves keys for that user. -- `--serve` gives the same principal to each request with the token. -- `--acp`, `--mcp`, and `--dashboard` open their sessions without the principal. +- Line mode, `-p`, your own `ui`, `--acp`, and `--mcp` open their sessions as `principal`, so `/connect` saves keys for that user. +- `--serve` checks the token first. On a host without `access`, each request then runs as `principal`. On a host with [`access`](./access), `identify` reads the user of each request. +- On a host with `access`, give `principal` too. `--mcp` does not start on such a host, because stdio has no login. `--dashboard` opens its sessions without the principal. ## Commands in line mode diff --git a/docs/harness/connect.md b/docs/harness/connect.md index 93482b85c0..9701c97d3a 100644 --- a/docs/harness/connect.md +++ b/docs/harness/connect.md @@ -15,10 +15,10 @@ Your harness runs on a server, but the people who use it are in a browser or ano ## Serve the session over HTTP -Mount one fetch handler on a route. `authorize` is required, so no endpoint is open by accident. +Mount one fetch handler on a route. The host reads the user of each request with its `access`, so no endpoint is open by accident. ```ts group=harness-connect -import { createHarnessHost, defineHarness } from '@tanstack/ai-harness' +import { createHarnessHost, defineHarness, membersOnly } from '@tanstack/ai-harness' import { createHarnessHandler } from '@tanstack/ai-harness/server' import { memoryPersistence } from '@tanstack/ai-persistence' import { openaiText } from '@tanstack/ai-openai' @@ -28,33 +28,51 @@ export const assistant = defineHarness({ adapter: openaiText('gpt-6.1-sol'), }) -const host = createHarnessHost({ persistence: memoryPersistence() }) - -export const handler = createHarnessHandler({ - host, - harness: assistant, - authorize: (request) => - request.headers.get('authorization') === `Bearer ${process.env.HARNESS_TOKEN}` - ? { id: 'user-1' } - : null, - canAccess: (principal, threadId) => threadId.startsWith(`${principal.id}-`), +const host = createHarnessHost({ + persistence: memoryPersistence(), + access: { + identify: ({ request }) => + request.headers.get('authorization') === `Bearer ${process.env.HARNESS_TOKEN}` + ? { id: 'user-1' } + : null, + authorize: membersOnly(), + }, }) + +export const handler = createHarnessHandler({ host, harness: assistant }) ``` -`canAccess` is required. It decides which threads each principal can use. To let every principal use every thread, pass `canAccess: () => true`. +- `identify` reads the user of each request. `null` answers `401`. +- `authorize` decides each action. A refused action answers `403`. `membersOnly()` lets a user use the threads they opened and the threads shared with them. +- A host with no `access` needs `openAccess` on the handler, else `createHarnessHandler` throws at the start. See [Serve with no login](./access#serve-with-no-login). + +[Access and tenants](./access) has the full rules. On a durable host (`stores.log`), also set `media: { secret }`. The secret signs media URLs, so every instance must use the same value. Without a log store, the default is a random key, and signed URLs stop working after a restart. -The handler answers these paths under your route: +The handler answers these paths under your route. For chat and events: - `GET capabilities`: the AG-UI capabilities, with the agents in `expose.agents`. - `POST run`: standard AG-UI. One request runs one prompt and streams it as SSE. The prompt keeps every content part of the last user message, for example an image. `useChat` works with it. See [Use it from useChat](#use-it-from-usechat). - `GET run?threadId=`: the saved messages, the running turn, and the waiting approvals, for a `useChat` that loads the thread after a reload. - `GET run?runId=`: the turn with that run id as SSE, from its first event. A reloaded `useChat` joins a running turn with it. - `GET events?threadId=`: every event of the session as SSE. Each event id is a cursor, so a reconnect with `Last-Event-ID` continues where it stopped. + +For one session: + - `POST control`: send `{ threadId, input }`, for example `{ op: 'prompt', message }`. You get a receipt back. Agents, settings, config keys, and commands need `expose`. See [Choose what clients can change](#choose-what-clients-can-change). - `GET snapshot?threadId=`: the status, running operations, and waiting approvals. +- `GET transcript?threadId=` and `GET describe?threadId=`: the saved messages, and the commands and settings that a client can use. - `GET agent-transcript?threadId=&subagentRunId=`: the saved messages of one child agent of the thread. See [Show what a child wrote after a reload](./subagents#show-what-a-child-wrote-after-a-reload). +- `POST media`, `GET media-url`, and `GET media`: upload a file, get a signed URL, and read the bytes. A signed URL works for 1 hour, and only for the user it was signed for. See [Send and show media](./media). + +For a list of sessions: + +- `GET sessions`: one page of the sessions that this user can list, newest first. +- `POST sessions`: change a session with `{ op, threadId }`. The ops are `rename`, `update`, `delete`, `fork`, `share`, and `unshare`. See [List, rename, and fork sessions](./sessions) and [Share the thread](./shared-threads#2-share-the-thread). +- `GET host-events`: the status changes and index changes of the sessions that this user can read, as SSE. + +A read of a thread that has no session answers `404`, and it starts no session. Only `GET run?threadId=` answers an empty thread, because `useChat` reads it before the first message. ## Cursors @@ -128,7 +146,7 @@ A background agent runs for minutes, and the user wants to correct it from the p ```ts group=harness-connect-agents import { defineAgent } from '@tanstack/ai' -import { createHarnessHost, defineHarness } from '@tanstack/ai-harness' +import { createHarnessHost, defineHarness, membersOnly } from '@tanstack/ai-harness' import { createHarnessHandler } from '@tanstack/ai-harness/server' import { memoryPersistence } from '@tanstack/ai-persistence' import { openaiText } from '@tanstack/ai-openai' @@ -146,15 +164,18 @@ export const studio = defineHarness({ expose: { agents: ['drafter'] }, }) -export const handler = createHarnessHandler({ - host: createHarnessHost({ persistence: memoryPersistence() }), - harness: studio, - authorize: (request) => - request.headers.get('authorization') === `Bearer ${process.env.HARNESS_TOKEN}` - ? { id: 'user-1' } - : null, - canAccess: (principal, threadId) => threadId.startsWith(`${principal.id}-`), +const host = createHarnessHost({ + persistence: memoryPersistence(), + access: { + identify: ({ request }) => + request.headers.get('authorization') === `Bearer ${process.env.HARNESS_TOKEN}` + ? { id: 'user-1' } + : null, + authorize: membersOnly(), + }, }) + +export const handler = createHarnessHandler({ host, harness: studio }) ``` 2. In the browser, send the message to the id of the run: @@ -188,7 +209,7 @@ The id of the run is the `operationId` in the receipt of `client.agents.drafter. - `mode: 'steer'` (the default): the message joins the next model call of the run. - `mode: 'followUp'`: the agent runs again after the run ends. `receipt.operationId` is the id of the new run. -`send` with `to` posts an `agentMessage` input to `POST control`. It runs as the user that `authorize` returned. The input looks like this: +`send` with `to` posts an `agentMessage` input to `POST control`. It runs as the user that `identify` returned. The input looks like this: ```json { "op": "agentMessage", "operationId": "op-agent-m1x2-3", "message": "Add a title.", "mode": "followUp" } @@ -301,32 +322,31 @@ export function ReloadableChat() { The join has these rules: -- `canAccess` decides who can join, the same as for every other path. +- The host checks that the user can read the thread, the same as for every other path. - On a durable host, any server can answer the join. It reads the turn from the shared log. - Without a log store, the turn must run on the host that answers the request. On another host, the join gets `404`. - A turn that ended gives the events that the session still keeps, and then the stream ends. ## Use a WebSocket -For one connection that carries events and inputs, authorize the upgrade, then hand the socket to `handleHarnessSocket`: +For one connection that carries events and inputs, hand the socket and its upgrade request to `handleHarnessSocket`: ```ts group=harness-connect import { handleHarnessSocket } from '@tanstack/ai-harness/server' import type { WebSocketLike } from '@tanstack/ai' -export function onUpgrade(socket: WebSocketLike) { - handleHarnessSocket({ - host, - harness: assistant, - socket, - principal: { id: 'user-1' }, - canAccess: (principal, threadId) => threadId.startsWith(`${principal.id}-`), - }) +export function onUpgrade(socket: WebSocketLike, request: Request) { + handleHarnessSocket({ host, harness: assistant, socket, request }) } ``` The client sends `{ type: 'harness.subscribe', threadId }` first. After that, each `harness.input` frame gets a `harness.receipt`, and every event arrives as a `harness.event` frame with its cursor. +- The host reads the user from the upgrade request with `identify`, the same as for HTTP. Every input on the socket runs as that user. +- A user that `identify` refuses gets an `unauthorized` error frame, and the socket closes. +- A user who cannot read the thread, or who loses access later, gets a `forbidden` error frame, and the socket closes. +- On a host with no `access`, pass `openAccess`, the same as for `createHarnessHandler`. + ## Use it as the model of another chat `harnessText` turns a harness into a text adapter. The outer chat sends a message, and the harness runs a full turn with its own tools, plugins, and agents: @@ -348,7 +368,7 @@ If a plugin picks the model of the harness, the harness has no adapter to read i ## What you have now -- One handler that serves standard AG-UI and the session stream. +- One handler that serves standard AG-UI and the session stream, as the user of each request. - Clients that can change only what you expose. - A `useChat` screen on the harness, with approvals, tools that run in the browser, and a reload that joins the running turn. - A typed web client that reconnects and resumes from a cursor. diff --git a/docs/harness/custom-ui.md b/docs/harness/custom-ui.md index ecf160444c..fb49b21a3d 100644 --- a/docs/harness/custom-ui.md +++ b/docs/harness/custom-ui.md @@ -459,7 +459,7 @@ Your harness runs on a server, and your UI is a web page. Give `createSessionVie Serve the session with `createHarnessHandler`, and mount the handler on `/api/harness`: ```ts group=harness-custom-ui-server -import { createHarnessHost, defineHarness } from '@tanstack/ai-harness' +import { createHarnessHost, defineHarness, membersOnly } from '@tanstack/ai-harness' import { createHarnessHandler } from '@tanstack/ai-harness/server' import { memoryPersistence } from '@tanstack/ai-persistence' import { openaiText } from '@tanstack/ai-openai' @@ -469,19 +469,22 @@ export const assistant = defineHarness({ adapter: openaiText('gpt-6.1-sol'), }) -const host = createHarnessHost({ persistence: memoryPersistence() }) - -export const handler = createHarnessHandler({ - host, - harness: assistant, - authorize: (request) => - request.headers.get('authorization') === `Bearer ${process.env.HARNESS_TOKEN}` - ? { id: 'user-1' } - : null, - canAccess: (principal, threadId) => threadId.startsWith(principal.id), +const host = createHarnessHost({ + persistence: memoryPersistence(), + access: { + identify: ({ request }) => + request.headers.get('authorization') === `Bearer ${process.env.HARNESS_TOKEN}` + ? { id: 'user-1' } + : null, + authorize: membersOnly(), + }, }) + +export const handler = createHarnessHandler({ host, harness: assistant }) ``` +`membersOnly()` lets each user use the threads that they opened, and the threads shared with them. See [Access and tenants](./access). + The handler also serves `GET transcript` and `GET describe`. The view reads them when it starts. For the other routes, see [Connect clients](./connect). ### In the browser diff --git a/docs/harness/deploy.md b/docs/harness/deploy.md index 5ac168fe12..91dddf4892 100644 --- a/docs/harness/deploy.md +++ b/docs/harness/deploy.md @@ -61,7 +61,7 @@ The worker keeps the conversation of its thread between calls. Workers get the e ## Use a harness on another machine -On the other machine, serve the harness with `createHarnessHandler`, and check a bearer token in `authorize`. See [Serve the session over HTTP](./connect#serve-the-session-over-http). +On the other machine, serve the harness with `createHarnessHandler`. Check the bearer token in `identify` of the host `access`. See [Serve the session over HTTP](./connect#serve-the-session-over-http). Here, use it as a model: diff --git a/docs/harness/fork-and-reset.md b/docs/harness/fork-and-reset.md index 11c834eb00..6c64562118 100644 --- a/docs/harness/fork-and-reset.md +++ b/docs/harness/fork-and-reset.md @@ -53,6 +53,8 @@ The fork copies: The fork does not copy plugin state, pending interrupts, queued inputs, running work, or [usage totals](./usage). If `newThreadId` already has messages, `host.fork` throws. It also throws for an `at` that the thread does not have. +On a host with [`access`](./access), pass `principal` too. The user needs the `fork` action on the old thread, and owns the new thread. + If a thread waits for an approval, fork it at a message before the tool call. A fork of the whole transcript copies the tool call, but the fork has no interrupt to answer it. ## Fork before or through a message @@ -74,7 +76,9 @@ export async function editAndResend(threadId: string, messageId: string, text: s - `{ before: id }` copies the messages before `id`. Before the first message, the fork has no messages, but it keeps the settings. - `{ through: id }` copies the messages up to and including `id`, as `at` does. -- The host picks the new thread id. The new index entry has the title of the old one plus ` (fork)`, and the same owner. +- The host picks the new thread id. The new index entry has the title of the old one plus ` (fork)`. +- `members` picks who can use the fork. `'self'`: only the user who forks, as the owner. `'all'`: the user who forks is the owner, and the other members keep their roles. See [Fork a shared thread](./shared-threads#6-fork-a-shared-thread). +- Without host `access` and without `members`, the fork keeps the owner and the members of the old thread. - `host.sessions.fork` calls `host.fork`, so it copies the same things as `host.fork`. ## Reset with a handoff note diff --git a/docs/harness/inputs.md b/docs/harness/inputs.md index a003b4a82e..3438e1b17c 100644 --- a/docs/harness/inputs.md +++ b/docs/harness/inputs.md @@ -112,7 +112,7 @@ const continuedReceipt = await client.continue({ inputId: 'req-43' }) ``` - `inputId` works as for `send`: a retry with the same id runs once. -- `continue` takes `overrides`, `principal`, and `context` too, as `send` does. +- `continue` takes `overrides`, `principal`, and `context` too, as `send` does. On a host with [`access`](./access), a session acts only as the user that opened it, so leave out `principal`. - When the turn starts, the last message of the transcript must be a user message or a tool result. Else the input is rejected with the reason `nothing_to_continue`, and `await continued` rejects with `InputRejectedError`. - A `continue` that a host left after a crash runs on the next host, as a prompt does. @@ -336,8 +336,9 @@ On `POST run`, the AG-UI `forwardedProps` of the request are the `context`, for - The `context` is part of the input for the `inputId` check. An input with the same id and another `context` gets `'conflict'`. - When `HarnessConfig.context` and the input `context` are both plain objects, tools get them merged, and the harness value wins for a key in both. In all other cases, tools get the harness value. If the harness has no value, tools get the input value. +- The harness adds the run of the turn to a plain-object context as `run`, and it replaces a `run` key that a client sent. A context that is not a plain object gets no `run`, so tools that read `ctx.context.run` fail. See [Act as the user of a run](./plugins#act-as-the-user-of-a-run). - Plugins and the router get it as `turn.context`. See [Pick the model of a turn](./plugins#pick-the-model-of-a-turn). -- The `context` is client data, so do not trust it. Who sent the input is the principal from `authorize`. See [Share a thread](./shared-threads). +- The `context` is client data, so do not trust it. Who sent the input comes from the server: the user of the request, or of the session handle. See [Access and tenants](./access). ## What you have now diff --git a/docs/harness/mcp-server.md b/docs/harness/mcp-server.md index 97053a0c3b..93128223b1 100644 --- a/docs/harness/mcp-server.md +++ b/docs/harness/mcp-server.md @@ -50,11 +50,15 @@ const host = createHarnessHost({ persistence: memoryPersistence() }) export const server = await createHarnessMcpServer({ host, harness: assistant, + // One user on this machine, so no login. + openAccess: true, }) serveMCPStdio(server) ``` +The host has no `access`, so the server needs `openAccess`. `true` opens each session with no user. `{ principal }` opens each session as that one user. + For Claude Code, run this command one time: ```bash @@ -211,19 +215,41 @@ When the work makes media, the results of `chat`, `approve`, `reject`, `resolve` ## Serve it over HTTP -The server also has a `fetch` handler. Mount it on a route of your app, for example `/mcp`, and check a token first: +Several people use the server over HTTP, so each request needs its user. Give the host [`access`](./access), and leave out `openAccess`. The server then reads the user of each HTTP request with `identify`: -```ts group=harness-mcp-server -export async function handleMcp(request: Request) { - const token = request.headers.get('authorization') - if (token !== `Bearer ${process.env.MCP_TOKEN}`) { - return new Response('Unauthorized', { status: 401 }) - } - return server.fetch(request) -} +```ts group=harness-mcp-server-http +import { createHarnessHost, defineHarness, membersOnly } from '@tanstack/ai-harness' +import { createHarnessMcpServer } from '@tanstack/ai-mcp/harness' +import { memoryPersistence } from '@tanstack/ai-persistence' +import { openaiText } from '@tanstack/ai-openai' + +const assistant = defineHarness({ + name: 'acme/assistant', + adapter: openaiText('gpt-6.1-sol'), +}) + +// Your own token check: a bearer token to a user. +const users = new Map([['Bearer token-ada', { id: 'ada', tenantId: 'acme' }]]) + +const host = createHarnessHost({ + persistence: memoryPersistence(), + access: { + identify: ({ request }) => users.get(request.headers.get('authorization') ?? '') ?? null, + authorize: membersOnly(), + }, +}) + +const server = await createHarnessMcpServer({ host, harness: assistant }) + +// Mount it on a route of your app, for example `/mcp`. +export const handleMcp = (request: Request) => server.fetch(request) ``` -The server opens sessions with no principal. Every caller that passes your check can use every thread and its media, so give each person their own server and token. +Each request runs as its user, and the host checks each tool call. A user reaches only the threads and the media that `authorize` allows. + +- A request that `identify` refuses gets `401`. A refused action comes back as a tool error. +- A server on a host with `access` works over HTTP only. Over stdio, it throws. +- It opens no session at the start, so it lists only the agents of the harness definition. The commands and the agents that plugins add get no tool on this server. Then add it to Claude Code: @@ -236,6 +262,7 @@ claude mcp add --transport http assistant https://your-app.example.com/mcp --hea - `threadId`: the conversation of a tool call without a `threadId`. Default `main`. - `approvals`: `'ask'` (default) asks in the client, or returns the approvals when the client cannot ask. `'auto'` approves every tool call. Both answer approvals only. - `filePaths`: the folders that a `path` attachment can read from. The real path of the file, after every symlink, must be inside one of them. Default: none, so the server refuses every `path` attachment. +- `openAccess`: for a host with no `access` only. `true` opens each session with no user, and `{ principal }` opens each session as that user. - `name` and `version`: the MCP server name and version. Default: the harness name and `1.0.0`. ## What you have now diff --git a/docs/harness/overview.md b/docs/harness/overview.md index 29c796ef1a..84883dc9ec 100644 --- a/docs/harness/overview.md +++ b/docs/harness/overview.md @@ -104,6 +104,8 @@ const session = await host.open(assistant, { threadId: 'thread-1' }) Opening the same `threadId` again returns the live session. +Does the host serve many users? Give it `access`, so each user reaches only their own sessions. See [Access and tenants](./access). + ## 3. Send a message `send` returns an operation. Await it for the text of the turn. diff --git a/docs/harness/plugins.md b/docs/harness/plugins.md index b5f1fe22b5..73123d25a0 100644 --- a/docs/harness/plugins.md +++ b/docs/harness/plugins.md @@ -225,6 +225,51 @@ export const usageByAgent = definePlugin({ - `ctx.subagentRunId` is different for each run of an agent. For a nested child, `ctx.parentSubagentRunId` is the id of the child that started it. - Each model call goes to `onUsage` one time, in the run that made the call. The lead turn does not count the usage of a child again. +## Act as the user of a run + +Several people can write in one thread, and a background agent can run after its turn ended. A tool must use the GitHub token of the user who started its own run, not of the user who sent the latest message. Each run carries its user, and that user's credentials and keys, as the run context: + +```ts group=harness-plugins +import { toolDefinition } from '@tanstack/ai' +import type { HarnessRunContext } from '@tanstack/ai-harness' + +export const github = definePlugin({ + name: 'acme/github', + setup: () => ({ + tools: [ + toolDefinition({ name: 'my_repos', description: 'List my GitHub repos' }).server( + async (_args, ctx) => { + const credential = await ctx.context.run.credentials.require('github') + const token = credential.type === 'api_key' ? credential.value : credential.accessToken + const response = await fetch('https://api.github.com/user/repos', { + headers: { authorization: `Bearer ${token}` }, + }) + return response.json() + }, + ), + ], + }), +}) +``` + +Where each part of a plugin reads the run: + +- A tool or a middleware hook: `ctx.context.run`. Type it with `HarnessRunContext`, for example `.server(...)` or `ChatMiddleware`. +- A command: `run` of the command context. Its `credentials` is the same as `run.credentials`. +- A plugin with `lifetime: 'turn'`: `ctx.turn.run`. Its `ctx.credentials` and `ctx.keys` belong to the same user. +- A session plugin (the default lifetime): `ctx.credentials` and `ctx.keys` have no user, so each call throws. `ctx.session.principal` is the owner of the session. + +The run has these fields: + +| Field | What it is | +| --- | --- | +| `principal` | The user who started the run: `{ id, tenantId }`. `undefined` in a session with no user. | +| `tenantId` | The tenant of the session. `undefined` in the default tenant. | +| `threadId`, `runId`, `parentRunId` | The thread, this run, and the run that started this one. | +| `logId` | The id of the log of the session in the store, with its tenant. | +| `signal` | Aborts when the run stops. | +| `credentials`, `keys` | The credentials and the model provider keys of the user of the run. | + ## Pick the model of a turn Some turns need another model: a short question does not need your strongest model, and an organization can allow only some models. Return an `adapter` picker from `setup`. When each turn starts, the session calls it with the message, context, and sender of that turn: @@ -255,6 +300,7 @@ The `turn` has these fields: | `context` | The [context of the input](./inputs#send-context-with-a-message). It is client data, so do not trust it. | | `principal` | Who sent the input. Default: the principal that opened the session. See [Share a thread](./shared-threads). | | `overrides` | The [overrides of the prompt](./turn-control#give-one-prompt-its-own-settings), when it has some. | +| `run` | The run of the turn: its user, and that user's credentials and keys. See [Act as the user of a run](#act-as-the-user-of-a-run). | The session uses the picks like this: @@ -398,6 +444,7 @@ If a `setup` throws, the session keeps working with no plugins. The promise reje - A plugin that adds tools, prompts, commands, settings, and state to any harness. - Middleware that sees every model call, in the lead turn and in every agent. +- Tools and commands that act as the user of their own run. - A picker that chooses the model of each turn from its message, context, and sender. - Plugins that share services, lists, and events without knowing each other. - Commands, agents, and whole plugins that change while the session runs. diff --git a/docs/harness/provider-keys.md b/docs/harness/provider-keys.md index 628728227e..dcef1dbcfc 100644 --- a/docs/harness/provider-keys.md +++ b/docs/harness/provider-keys.md @@ -37,7 +37,7 @@ const claude = keyedAdapter(anthropicByok, (key) => - Each provider package exports one from its `/byok` subpath, for example `anthropicByok`, `grokByok`, `falByok`, and `openrouterByok`. - `keyedAdapter` works for every adapter kind: text, image, speech, audio, and video. -Before each turn, the session calls your function with the key of the user of that session. +Before each turn, the session calls your function with the key of the user who sent that turn. ## 2. Add the providerKeys plugin @@ -137,7 +137,7 @@ After `/connect:openai`, send the message again. ## Where the keys are stored -`/connect:` saves the key in the credential store of the host, as an `api_key` credential with the provider id. The store keeps it for the principal of the session: +`/connect:` saves the key in the credential store of the host, as an `api_key` credential with the provider id. The store keeps it for the user who runs the command: - CLI: pass the store to `runCli(harness, { credentials })`. - Hosted app: your database or your secrets manager. Each signed-in user is a principal, so each user has their own keys. Pass the store to `createHarnessHost({ credentials })`. @@ -153,9 +153,11 @@ The key that a user pastes can be a Claude OAuth token (`sk-ant-oat...`). `creat If the user has no saved key, the session reads the env var of the provider, for example `OPENAI_API_KEY`. A saved key wins over the env var. So you can develop with a `.env` file, and your users connect in the app. +On a host with [`access`](./access), the session does not read the env vars. Each user must connect their own key, so no user spends the key of your server. To let users fall back to the env var, set `envKeys: true` in `access`. + ## Use the key of the user in an agent -Your agent makes images with fal, so it needs the fal key of the user. Every agent gets `ctx.keys`. In a harness, `ctx.keys` reads the same keys as the turn: +Your agent makes images with fal, so it needs the fal key of the user. Every agent gets `ctx.keys`. In a harness, `ctx.keys` reads the keys of the user who started the agent: ```ts group=harness-provider-keys import { defineAgent } from '@tanstack/ai' @@ -182,7 +184,7 @@ Put the agent in `subagents.agents` of the harness, and the model calls it as a - `ctx.keys.get(provider)`: the key, or `null`. - `ctx.keys.require(provider)`: the key. If it is missing, the run stops with `harness.auth_required`, the same as a turn. - The children of the agent get the same keys. -- Plugins get the same keys as `ctx.keys` in `setup`. +- In a plugin, a tool reads the keys of its run as `ctx.context.run.keys`, and a command reads `run.keys`. See [Act as the user of a run](./plugins#act-as-the-user-of-a-run). Outside a harness, `ctx.keys` reads the env var. See [Keys inside agents](../advanced/byok#keys-inside-agents). @@ -197,5 +199,5 @@ Outside a harness, `ctx.keys` reads the env var. See [Keys inside agents](../adv - `/connect:` for OpenAI, Anthropic, xAI, fal, and OpenRouter in the harness that you ship, with no `.env` file. - One key for each user, in your credential store. -- The env var as a fallback while you develop. -- Agents and plugins that use the key of the user through `ctx.keys`. +- The env var as a fallback while you develop, and no fallback on a host with access unless you set `envKeys`. +- Agents and plugins that use the key of the user of each run. diff --git a/docs/harness/session-log.md b/docs/harness/session-log.md index 4169368759..c2f29e097e 100644 --- a/docs/harness/session-log.md +++ b/docs/harness/session-log.md @@ -87,18 +87,26 @@ export const watcher = defineHarness({ durability: { project }, }) +const { runs, sessions } = memoryPersistence().stores +const watcherHost = createHarnessHost({ + persistence: { stores: { log: memoryLogStore(), runs, sessions } }, + access: { + identify: ({ request }) => + request.headers.get('authorization') === `Bearer ${process.env.HARNESS_TOKEN}` + ? { id: 'ci-bot' } + : null, + // The CI bot is the only user, so it can use every thread. + authorize: ({ principal }) => principal.id === 'ci-bot', + }, +}) + // A durable host signs media URLs with the same secret on every instance. const mediaSecret = process.env.HARNESS_MEDIA_SECRET if (!mediaSecret) throw new Error('Set HARNESS_MEDIA_SECRET.') export const handler = createHarnessHandler({ - host, + host: watcherHost, harness: watcher, - authorize: (request) => - request.headers.get('authorization') === `Bearer ${process.env.HARNESS_TOKEN}` - ? { id: 'ci-bot' } - : null, - canAccess: () => true, media: { secret: mediaSecret }, }) ``` @@ -126,15 +134,20 @@ await client.continue() ## Read the log -The store is yours, so read your records back from it, for example to rebuild your state when a session opens: +The store is yours, so read your records back from it, for example to rebuild your state when a session opens. The store keeps each log under an id with its tenant, so build the id with `logIdOf`: ```ts group=harness-session-log -const entries = await log.read({ logId: 'thread-1' }) +import { logIdOf } from '@tanstack/ai-harness' + +// The default tenant. For a tenant, pass `tenantId` too: 'acme:thread-1'. +const entries = await log.read({ logId: logIdOf({ logId: 'thread-1' }) }) const signals = entries.filter((entry) => entry.record.type === 'app.signal') console.log(signals.map((entry) => entry.seq)) ``` -A reader outside a session that needs a `MessageStore`, for example `reconstructChat`, uses `logMessageStore`. It folds the log with the same `project`: +Each `harness.input` record has the `principal` of its sender: its `id` and `tenantId`. Every kind of input has it when the sender is known, also a cancel, an answer, a delivery change, and a change of settings. + +A reader outside a session that needs a `MessageStore`, for example `reconstructChat`, uses `logMessageStore`. It folds the log with the same `project`. Pass `tenantId` to read the threads of a tenant: ```ts group=harness-session-log import { logMessageStore } from '@tanstack/ai-harness' diff --git a/docs/harness/sessions.md b/docs/harness/sessions.md index 5b5b014010..4d4b912b90 100644 --- a/docs/harness/sessions.md +++ b/docs/harness/sessions.md @@ -40,16 +40,16 @@ const session = await host.open(assistant, { principal: { id: 'ada' }, }) await session.send('Plan a trip to Lisbon.') -console.log((await host.sessions.get('chat-1'))?.principal) // { id: 'ada' } +console.log((await host.sessions.get({ threadId: 'chat-1' }))?.ownerId) // 'ada' ``` What an entry holds: - `title`: the name of the chat. Set it with `host.sessions.rename`, or from a plugin with `ctx.session.updateEntry({ title })`. - `updatedAt`: the time of the last change. The list sorts by it. -- `principal`: the owner. This is the user who opened the thread first. -- `usage`: the token totals of the thread. -- `parentThreadId`: the parent thread of a child session, for example a subagent run. +- `ownerId` and `members`: the user who opened the thread first, and the users who can open it, each with a role. The owner is a member with the role `owner`. +- `tenantId`: the tenant of the thread. Without it, the thread is in the default tenant. See [Access and tenants](./access). +- `usage` and `parentThreadId`: the token totals of the thread, and the parent thread of a child session, for example a subagent run. A durable host takes `sessions` next to `log`: @@ -62,7 +62,11 @@ const durableHost = createHarnessHost({ }) ``` -Without `stores.sessions`, `list` is empty, `get` returns `undefined`, and `rename`, `update`, and `delete` write nothing. To write your own store, read the [`SessionIndexStore` contract](../persistence/store-reference#sessionindexstore). +- Without `stores.sessions`, `list` is empty, `get` returns `undefined`, and `rename`, `update`, and `delete` write nothing. +- The store needs `getVersioned` and `upsertIf`, else `createHarnessHost` throws. The host changes an entry with compare-and-set, so two hosts never lose each other's changes. +- To write your own store, read the [`SessionIndexStore` contract](../persistence/store-reference#sessionindexstore). + +Each method of `host.sessions` takes one object with a `threadId` or filters, and a `principal`. The tenant of `principal` picks the entries. On a host with `access`, the `principal` is required, and the host checks each call. See [Access and tenants](./access). ## List the sessions of a user @@ -85,7 +89,8 @@ if (page.truncated) { - The newest `updatedAt` comes first. - By default, the list has only top-level sessions. Pass `parentThreadId: 'chat-1'` to get the child sessions of `chat-1`. - With `limit`, a page that has more entries after it has `truncated: true` and a `cursor`. Pass the `cursor` to get the next page. -- If `principal` has a `tenantId`, the tenant of the entry must match too. +- A list reads one tenant: the tenant of `principal`. Without a `tenantId`, the default tenant. +- Without host `access`, the list has every session of that tenant. With `membersOnly()`, it has only the sessions where `principal` is a member. With your own `authorize`, it has the sessions that `authorize` lets `principal` list. To find sessions, add filters. All filters must match, and `limit` and `cursor` page the filtered list: @@ -121,7 +126,8 @@ await fork.send('Plan a trip to Porto.') - `at: { through: id }` copies the messages up to and including `id`. - `at: { before: id }` copies the messages before `id`. A fork before the first message has no messages, but it keeps the settings. -- The host picks the new thread id. The new entry has the title of the old one plus ` (fork)`, and the same owner. A thread without a title gives a fork without a title. +- The host picks the new thread id. The new entry has the title of the old one plus ` (fork)`. A thread without a title gives a fork without a title. +- With a `principal` and host `access`, the user who forks owns the fork. `members` picks who else is in it. See [Fork a shared thread](./shared-threads#6-fork-a-shared-thread). Without access and without `members`, the fork keeps the owner and the members of the old thread. - If the transcript has no message with that id, `fork` throws. - If the index entry of the thread names another harness, `fork` throws. One host can run more than one harness. @@ -150,12 +156,14 @@ await session.unrevert() ```ts group=harness-sessions await host.sessions.rename({ threadId: forked.threadId, title: 'Porto trip' }) -await host.sessions.delete('chat-1') +await host.sessions.delete({ threadId: 'chat-1' }) ``` - `rename` resolves to the changed entry. For a thread with no entry, it resolves to `undefined`. -- `delete` removes the index entry only. The transcript, the log, the media, and the settings of the thread stay in their stores. -- After `delete`, the turns of the open session do not write the entry again. If the thread opens again after a restart, the host writes a new entry for it. +- `delete` removes the index entry for good. The transcript, the log, the media, and the settings of the thread stay in their stores. +- After `delete`, the turns of the open session do not write the entry again. If the thread opens again after a restart, the host writes a new entry for it, with the user who opens it as the owner. + +`delete` closes the open session and removes the entry and the transcript, so the next user who opens the same `threadId` starts empty. Uploaded media and plugin data of the thread stay in their stores. ## Keep app data on a session @@ -192,7 +200,7 @@ await host.sessions.update({ ### Writes from several servers -A turn on one server writes the usage of the thread. At the same time, the app on another server can write the pins. Your store keeps both writes when it has `getVersioned` and `upsertIf`. Then the host runs your change function again on the new entry. `memoryPersistence()` has both. To add them to your own store, read [Keep the writes of several hosts](../persistence/store-reference#keep-the-writes-of-several-hosts). +A turn on one server writes the usage of the thread. At the same time, the app on another server can write the pins. The store keeps both writes with `getVersioned` and `upsertIf`: when another server wrote first, the host runs your change function again on the new entry. For the rules of your own store, read [Keep the writes of several hosts](../persistence/store-reference#keep-the-writes-of-several-hosts). ### Read and write the metadata in a plugin @@ -245,6 +253,7 @@ void watch() - `status`: `running` (an operation runs), `waiting` (a question or an approval waits for an answer), or `idle`. `at` is the time of the change. - `session`: the host wrote an entry, for a new session, the end of a turn, a rename, an update, or a fork. `entry` is the new entry. - `session-deleted`: the host removed an entry. `entry` is the removed entry. +- `revoked`: `host.revoke` or `unshare` took the access of the user `principalId` to `threadId`. See [Use your own access tables](./access#use-your-own-access-tables). Each host reports only the sessions that are open in it. A session with no status event is not open, so show it as idle. diff --git a/docs/harness/shared-threads.md b/docs/harness/shared-threads.md index 305ec28fb2..6e573d8df7 100644 --- a/docs/harness/shared-threads.md +++ b/docs/harness/shared-threads.md @@ -2,24 +2,25 @@ title: Share a thread between people id: harness-shared-threads order: 7 -description: "Let several people write in one thread. Each message runs as its sender, with the credentials and the provider key of that person." +description: "Let a team use one thread. Share it with members and viewers, and run each message as its sender, with the credentials and the provider key of that person." keywords: - tanstack ai - harness - shared thread - - principal + - membersOnly + - share - tenantId - canJoin --- -Several people write in one thread: a team in one chat, or a customer and a support agent. A message from Bob must run as Bob. His GitHub token and his provider key must do the work, not those of the person who opened the thread first. The harness runs each message as the person who sent it. +Several people work in one thread: a team in one chat, or a customer and a support agent. Only the people you invite can see it. A message from Bob must run as Bob: his GitHub token and his provider key do the work, not those of the person who opened the thread. `membersOnly()` gives each thread an owner and a list of members, and the harness runs each message as the person who sent it. -## 1. Return the sender from `authorize` +## 1. Turn on members -`authorize` decides who sends each request. Each message runs as the principal that it returns. If your users belong to organizations, add `tenantId`: +Use `membersOnly()` as the `authorize` of the host. [Access and tenants](./access) explains `identify`: ```ts group=harness-shared-threads -import { createHarnessHost, defineHarness } from '@tanstack/ai-harness' +import { createHarnessHost, defineHarness, membersOnly } from '@tanstack/ai-harness' import { createHarnessHandler } from '@tanstack/ai-harness/server' import { memoryPersistence } from '@tanstack/ai-persistence' import { openaiText } from '@tanstack/ai-openai' @@ -29,35 +30,78 @@ export const assistant = defineHarness({ adapter: openaiText('gpt-6.1-sol'), }) -const host = createHarnessHost({ persistence: memoryPersistence() }) - -// Your own lookups: a session token to a user, and a thread to its members. +// Your own login check: a session token to a user. const users = new Map([ ['Bearer token-ada', { id: 'ada', tenantId: 'acme' }], ['Bearer token-bob', { id: 'bob', tenantId: 'acme' }], ]) -const members = new Map([['team-thread', ['ada', 'bob']]]) -export const handler = createHarnessHandler({ - host, - harness: assistant, - authorize: (request) => users.get(request.headers.get('authorization') ?? '') ?? null, - canAccess: (principal, threadId) => - members.get(threadId)?.includes(principal.id) ?? false, +export const host = createHarnessHost({ + persistence: memoryPersistence(), + access: { + identify: ({ request }) => users.get(request.headers.get('authorization') ?? '') ?? null, + authorize: membersOnly(), + }, +}) + +export const handler = createHarnessHandler({ host, harness: assistant }) +``` + +The user who opens a new thread is its owner. Nobody else can open it until the owner shares it. What each role can do: + +| Action | Owner | Member | Viewer | +| --- | --- | --- | --- | +| Read the thread and see it in a list | Yes | Yes | Yes | +| Send messages, run commands, start agents | Yes | Yes | No | +| Answer and cancel | Every run | Only the runs they started | No | +| Fork | Yes | Yes | No | +| Configure, share, delete, transfer | Yes | No | No | + +A user who is not a member can do nothing. A child thread of a subagent uses the members of its root thread. + +## 2. Share the thread + +The owner adds a member from your server code: + +```ts group=harness-shared-threads +const ada = { id: 'ada', tenantId: 'acme' } + +await host.open(assistant, { threadId: 'team-thread', principal: ada }) +await host.sessions.share({ + threadId: 'team-thread', + userId: 'bob', + role: 'member', + principal: ada, }) ``` -You pass nothing else. The sender of each request is used everywhere: +From the browser, the owner sends the same change to `POST sessions`: -- The log keeps the sender of each message. After a restart, a message runs again as its sender. -- The tools of a turn get the credentials of the sender, in `ctx.credentials`. -- The model call uses the provider key of the sender, in `ctx.keys`. So the sender pays for the call. -- In a plugin, `ctx.session.principal` is the sender of the running turn. Outside a turn, it is the person who opened the session first. -- Commands, for example `connect:github`, run as their sender. So a sign-in belongs to the person who signed in. A turn that waits for a sign-in continues only when its own sender signs in. +```ts group=harness-shared-threads-invite +export async function invite(userId: string, role: 'member' | 'viewer') { + const response = await fetch('/api/harness/sessions', { + method: 'POST', + headers: { 'content-type': 'application/json', authorization: 'Bearer token-ada' }, + body: JSON.stringify({ op: 'share', threadId: 'team-thread', userId, role }), + }) + if (!response.ok) throw new Error(`The share failed with ${response.status}.`) + return response.json() +} +``` + +- `share` adds a member, or changes the role of a member. It resolves to the changed entry. +- `unshare` removes a member: `host.sessions.unshare({ threadId, userId, principal })` on the server, or `{ op: 'unshare', threadId, userId }` from the browser. +- Only the owner can share. Another user gets `AccessDeniedError`, and the handler answers `403`. +- The owner cannot be removed. -`handleHarnessSocket` works the same way. Every input on the socket runs as the `principal` that you give it. Context that a client sends with a message is client data, and it never changes the sender. See [Send context with a message](./inputs#send-context-with-a-message). +When `unshare` removes a user: -## 2. Connect each person +- The queued and running work that the user started on this host stops. The log keeps who cancelled it. +- The streams of the user on the thread close. A socket gets a `forbidden` error frame. +- The open handles of the user throw `AccessDeniedError`. +- The signed media URLs of the user stop working. + +## 3. Connect each person Each person opens the same thread with their own token. This snippet shows two clients. In a real app, each person has one client in their own browser: @@ -82,19 +126,47 @@ await bob.send('List my open pull requests.') Bob sends his message while the turn of Ada runs. His message waits. Then it runs as its own turn, as Bob. -On the server, give the sender to `send` yourself. `resolve` and `command` take `principal` too. Without it, the message runs as the person who opened the session: +On the server, open one handle for each person. A handle acts only as its user: ```ts group=harness-shared-threads -const session = await host.open(assistant, { threadId: 'team-thread' }) -const ada = { id: 'ada', tenantId: 'acme' } +const bob = { id: 'bob', tenantId: 'acme' } +const bobSession = await host.open(assistant, { threadId: 'team-thread', principal: bob }) -session.send('Summarize the open issues.', { principal: ada }) -session.command({ name: 'connect:github', principal: ada }) +await bobSession.send('Summarize the open issues.') ``` -A command acts for the person who runs it. In its `run` context, `session.principal` is that person, and `session.send` runs as that person. +## 4. Each message runs as its sender -If you write a command that saves a sign-in or a key, save it with the `credentials` of the second `run` argument. Those credentials belong to the person who runs the command. The `ctx.credentials` of `setup` belong to the sender of the running turn: +Each message, command, and agent runs as the person who started it: + +- The log keeps the sender of each input, also of a cancel, an answer, and a change of settings. After a restart, a message runs again as its sender. +- Tools and middleware read the run as `ctx.context.run`. It has the user, the credentials, and the provider keys of the sender. +- The model call uses the provider key of the sender. So the sender pays for the call. +- A command runs as the person who runs it. So a sign-in, for example `connect:github`, belongs to the person who signed in. +- A background agent runs as the person who started it, also after the turn ends. The wake turn after the agent runs as that person too. + +A tool reads the credentials of the user of its own run: + +```ts group=harness-shared-threads +import { toolDefinition } from '@tanstack/ai' +import type { HarnessRunContext } from '@tanstack/ai-harness' + +export const myIssues = toolDefinition({ + name: 'my_linear_issues', + description: 'List the Linear issues of the user', +}).server(async (_args, ctx) => { + const credential = await ctx.context.run.credentials.require('linear') + const key = credential.type === 'api_key' ? credential.value : credential.accessToken + const response = await fetch('https://api.linear.app/graphql', { + method: 'POST', + headers: { authorization: key, 'content-type': 'application/json' }, + body: JSON.stringify({ query: '{ viewer { assignedIssues { nodes { title } } } }' }), + }) + return response.json() +}) +``` + +A command that saves a sign-in or a key saves it with the `credentials` of its run context. Those credentials belong to the person who runs the command: ```ts group=harness-shared-threads import { defineCommand, definePlugin } from '@tanstack/ai-harness' @@ -120,7 +192,11 @@ export const linearKey = definePlugin({ }) ``` -## 3. Keep other people out of a running turn +The `ctx.credentials` of a plugin `setup` have no user, so each call throws. See [Act as the user of a run](./plugins#act-as-the-user-of-a-run). + +Context that a client sends with a message is client data, and it never changes the sender. See [Send context with a message](./inputs#send-context-with-a-message). + +## 5. Keep other people out of a running turn A message with `busy: 'steer'` can join the turn that runs now. A joined message runs with the credentials of the sender of that turn. So without `turn.canJoin`, a waiting message joins the running turn only when the same person sent both: @@ -142,13 +218,36 @@ export const team = defineHarness({ }) ``` -A background agent that a turn starts keeps the provider key of the sender of that turn. It keeps the key also after the turn ends. The wake turn after the agent runs as the person who started the agent. +## 6. Fork a shared thread + +Bob wants to try another path without changing the thread of the team. He forks it. The user who forks owns the fork: + +```ts group=harness-shared-threads +const [first] = await bobSession.transcript() +if (!first?.id) throw new Error('The thread has no messages.') + +const fork = await host.sessions.fork({ + harness: assistant, + threadId: 'team-thread', + at: { through: first.id }, + members: 'all', + principal: bob, +}) +console.log(fork.ownerId) // 'bob' +``` + +`members` picks who else is in the fork: + +- `'self'`: only the user who forks. +- `'all'`: the members of the old thread keep their roles. The old owner becomes a `member`. + +Without `members`, the fork uses `fork.members` of `membersOnly()`, and the default is `'self'`. To change it for every fork, pass `membersOnly({ fork: { members: 'all' } })`. -One limit applies to code that runs after its turn ended, for example a tool of a background agent. That code reads `ctx.credentials` and `ctx.session.principal` of the turn that runs at that time. In an agent, use `ctx.keys`. In a command, use the `credentials` of its `run` context. +## 7. Keep organizations apart -## 4. Keep organizations apart +The tenant of a user picks the thread. One `threadId` in two tenants is two threads, so Acme and Globex never see each other's work. See [Keep tenants apart](./access#keep-tenants-apart). -The credential scope of a turn has the `userId` and the `tenantId` of the sender. So one person in two organizations has two sets of credentials. The memory store of `memoryPersistence()` keeps them apart. In your own credential store, put `scope.tenantId` in the key: +The credential scope of a run has the `userId` and the `tenantId` of its user. So one person in two organizations has two sets of credentials. The memory store of `memoryPersistence()` keeps them apart. In your own credential store, put `scope.tenantId` in the key: ```ts group=harness-shared-threads import type { Scope } from '@tanstack/ai' @@ -180,15 +279,16 @@ export async function shareKey( } ``` -- A turn reads the credential of its sender first. If the sender has none, the turn reads the credential of the tenant. -- A save from a turn goes to the sender only. This includes `connect:` and `ctx.credentials.set`. The credential of the tenant does not change. -- `ctx.credentials.list()` shows the credentials of the sender and of the tenant. For the same id, the credential of the sender wins. -- Use an API key for a credential of the tenant. If the turn refreshes an OAuth token of the tenant, the new token goes to the sender only. +- A run reads the credential of its user first. If the user has none, the run reads the credential of the tenant. +- A save from a run goes to its user only. This includes `connect:` and `run.credentials.set`. The credential of the tenant does not change. +- `run.credentials.list()` shows the credentials of the user and of the tenant. For the same id, the credential of the user wins. +- Use an API key for a credential of the tenant. If a run refreshes an OAuth token of the tenant, the new token goes to its user only. ## What you have now +- Threads that only their members can see, with an owner, members, and viewers. - Each message runs as the person who sent it, also after a restart. - Tools and model calls use the credentials and the provider key of the sender. - A message from another person waits for its own turn. -- Sign-ins and keys belong to the person who saved them. -- The credentials of two organizations stay apart, and one key can serve a whole organization. +- Forks that the forker owns, with the members you pick. +- The threads and credentials of two organizations stay apart. diff --git a/docs/harness/snapshots.md b/docs/harness/snapshots.md index 790719cb9d..9a8fb5e6be 100644 --- a/docs/harness/snapshots.md +++ b/docs/harness/snapshots.md @@ -73,9 +73,11 @@ Each model step that changed files leaves a step: `{ from, to, files, toolCallId On a [durable host](./durable-sessions), like the host above, the steps are `tanstack/snapshots:step` records in the session log: ```ts group=harness-snapshots +import { logIdOf } from '@tanstack/ai-harness' + await session.send('Fix the login form.') -for (const { record } of await log.read({ logId: 'fix-login' })) { +for (const { record } of await log.read({ logId: logIdOf({ logId: 'fix-login' }) })) { if (record.type !== 'tanstack/snapshots:step') continue if (typeof record.from !== 'string' || typeof record.to !== 'string') continue const { files, patch } = await history.diff(record.from, record.to) diff --git a/docs/harness/subagents.md b/docs/harness/subagents.md index 90b0e37ca4..b5002780cb 100644 --- a/docs/harness/subagents.md +++ b/docs/harness/subagents.md @@ -61,7 +61,8 @@ export const review = definePlugin({ - `ctx.agents.run(agent, input)` runs one agent and resolves with its result. - `ctx.agents.start(agent, input, { wake: true })` runs it in the background and starts a turn when it finishes or fails. -- `ctx.agents.start(agent, input, { principal })` runs it as that user. Without `principal`, it runs as the user of the current turn. The wake turn runs as the same user. +- `ctx.agents.start(agent, input, { principal })` runs it as that user. The wake turn runs as the same user. +- Without `principal`, a plugin with `lifetime: 'turn'` starts it as the user of its turn. A session plugin starts it as the user who opened the session. To start it for the user of a tool call, pass `principal: ctx.context.run.principal`. - `ctx.agents.group(options, body)` runs several. With `onFailure: 'cancel-siblings'`, one failure cancels the others. With `'collect'`, use `group.runSettled` to get every result or error. Every child settles before `group` returns. The turn that `wake: true` starts sends the model a user message from the harness, for example `Background agent pricer finished: ...`. The user did not write it. So the transcript marks it with `metadata.harness.wake: true`, and a view from `createSessionView` does not show it. In your own chat UI, skip the user messages with this mark. diff --git a/docs/harness/turn-control.md b/docs/harness/turn-control.md index d3c36357e1..bbed5a6696 100644 --- a/docs/harness/turn-control.md +++ b/docs/harness/turn-control.md @@ -191,7 +191,7 @@ The join stops at the first message that: That message and the ones after it wait. `canJoin` runs again before each model call and at each final answer, so it can run more than once for one message. Keep it a quick check with no side effects. What still waits when the turn ends runs as its own turn, before other queued turns. A cancelled one settles `aborted` and does not run. -A joined message runs with the credentials of the sender of the turn. So without `turn.canJoin`, a waiting message joins the running turn only when the same person sent both. A message from another person waits and runs as its own turn after the running one, as its sender. Joins are a prefix, so it also holds back the messages that wait after it. Set `canJoin` to choose your own rule. Compare `principal` with `turnPrincipal` to keep one person's message out of another person's turn. See [Keep other people out of a running turn](./shared-threads#3-keep-other-people-out-of-a-running-turn). +A joined message runs with the credentials of the sender of the turn. So without `turn.canJoin`, a waiting message joins the running turn only when the same person sent both. A message from another person waits and runs as its own turn after the running one, as its sender. Joins are a prefix, so it also holds back the messages that wait after it. Set `canJoin` to choose your own rule. Compare `principal` with `turnPrincipal` to keep one person's message out of another person's turn. See [Keep other people out of a running turn](./shared-threads#5-keep-other-people-out-of-a-running-turn). `onJoin` runs before the model call that the joined messages reach. It can return messages, which come after the joined ones, and records, which land in the same append as the join. Records need a durable host. diff --git a/docs/harness/usage.md b/docs/harness/usage.md index 338198c363..ed9411612d 100644 --- a/docs/harness/usage.md +++ b/docs/harness/usage.md @@ -90,12 +90,12 @@ Read the totals of the sender before you send a prompt. This route refuses a pro const BUDGET_USD = 5 export async function send(threadId: string, userId: string, text: string) { - const thread = await host.open(assistant, { threadId }) + const thread = await host.open(assistant, { threadId, principal: { id: userId } }) const spent = thread.usage().bySender[userId]?.cost ?? 0 if (spent >= BUDGET_USD) { return { error: 'You used your budget for this thread.' } } - return thread.send(text, { principal: { id: userId } }).receipt + return thread.send(text).receipt } ``` @@ -106,7 +106,7 @@ Many providers report no cost, so `cost` can stay empty. For those, compare the With `stores.sessions`, each [session index](./sessions) entry has a `usage` field. The host writes the token counts of `session.usage()` at the end of each turn. `turns` is the number of model calls. ```ts group=harness-usage -const page = await host.sessions.list({ principal: { id: 'ada' }, limit: 20 }) +const page = await host.sessions.list({ limit: 20 }) for (const entry of page.entries) { console.log(entry.title ?? 'New chat', entry.usage?.totalTokens ?? 0) } diff --git a/docs/persistence/store-reference.md b/docs/persistence/store-reference.md index c0de9fc442..8432c35b13 100644 --- a/docs/persistence/store-reference.md +++ b/docs/persistence/store-reference.md @@ -602,6 +602,13 @@ list of sessions to show. The session stores hold that data. All of them are opt-in. `withPersistence` uses only `sessions`, to index subagent children. The other stores are for your host code. +A harness host keeps the data of each tenant apart. In these stores, and in +`messages`, `runs`, `interrupts`, `metadata`, and the media artifacts, it writes the +thread id as +`logIdOf({ tenantId, logId: threadId })` from `@tanstack/ai-harness`: `acme:thread-1`, +or `:thread-1` in the default tenant. Keep it as an opaque string. The session index +is the one store that keeps `tenantId` and `threadId` as two fields. + ## InboxStore A client sends an input, and the host must not lose it in a crash or run it twice. @@ -1138,12 +1145,25 @@ restart. 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. +`SessionIndexStore` keeps one small entry per tenant and thread. The index holds no +messages. The data of each thread stays in the stores that hold it. ```ts +interface SessionMember { + userId: string + role: 'owner' | 'member' | 'viewer' +} + +interface SessionIndexKey { + tenantId?: string // without it, the default tenant + threadId: string +} + interface SessionIndexEntry { + tenantId?: string // with `threadId`, the key of the entry threadId: string + ownerId?: string // the user id of the owner. Without it, the thread has no owner + members: Array // the users who can open the thread. Empty for a subagent child harness?: string // the name of the setup that runs the thread title?: string parentThreadId?: string // the parent thread. A subagent child's own threadId is `subagent:` @@ -1151,7 +1171,6 @@ interface SessionIndexEntry { 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 @@ -1165,10 +1184,11 @@ interface SessionIndexEntry { } interface SessionIndexListOptions { + tenantId?: string // one tenant per list. Without it, the default tenant + memberId?: string // only the entries with this user in `members` limit?: number cursor?: string parentThreadId?: string | null - principal?: { id: string; tenantId?: string } search?: string harness?: string metadata?: Record @@ -1181,81 +1201,95 @@ interface SessionIndexPage { } 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, + get(key: SessionIndexKey): Promise + getVersioned( + key: SessionIndexKey, ): Promise<{ entry: SessionIndexEntry; revision: string } | undefined> - upsertIf?(write: { + upsert(entry: SessionIndexEntry): Promise + upsertIf(input: { entry: SessionIndexEntry - expectedRevision: string | null + revision: string | null }): Promise<{ ok: true; revision: string } | { ok: false; reason: 'conflict' }> + list(options?: SessionIndexListOptions): Promise + delete(key: SessionIndexKey): Promise } ``` -Rules for `upsert`, `get`, and `delete`: +The key of an entry is its `tenantId` and its `threadId` together. A key without +`tenantId` is in the default tenant. So one thread id in two tenants is two entries. +In SQL, store the default tenant as `''`, and make `(tenant_id, thread_id)` the +primary key. + +Rules for `get`, `upsert`, 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. +- `get` returns `undefined` for an unknown key, not `null`. +- `upsert` replaces the whole entry for the key of `entry`. 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 with `upsertIf`. +- Keep `ownerId` and `members` as they are written. The host decides what each role + can do. +- `delete` is a hard delete. After it, `get` and `getVersioned` return `undefined`, + and `upsertIf` with revision `null` writes a new entry. The messages, the log, and + the other data of the thread stay. For an unknown key, `delete` does nothing. Rules for `list`: +- Return the entries of one tenant only: the `tenantId` of the options, or the + default tenant without it. Never return an entry of another tenant. - Sort by `updatedAt`, the newest first. Entries with the same `updatedAt` go in `threadId` order. `pinned` gets no special place. +- `memberId` returns only the entries whose `members` have that user id, with any + role. A subagent child has no members, so it never matches. - `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. + only the entries with no parent. `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, in +the tenant of its `tenantId` option. The entry has the parent thread, the tool call +that started the child, and the `ownerId` and `harness` of the parent entry. It has no +members: access to a child goes through its top-level thread. ### 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. +usage of a turn while another host writes the pinned files of the app. The host +changes an entry with `getVersioned` and `upsertIf`, so the second write never removes +the change of the first. Both methods are required: without them, `createHarnessHost` +throws. 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. + key. The format of the revision is yours. - A write that changes the entry gives it a new revision: `upsert` and `upsertIf` both. A revision can be a hash of the content (an ETag of S3 or Netlify Blobs). Then the same entry, written again after a `delete`, gets the same revision back. -- `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. +- `upsertIf` writes only when the current revision of the key is `revision`, and + returns `{ ok: true, revision }` with the new revision. Pass `null` to write only + when the key 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 a revision, write with + `UPDATE ... WHERE tenant_id = ? AND 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 +`memoryPersistence()` has every method. Its cursor is `:` of the last entry of a page. ### Check the session stores with the conformance suite `runPersistenceConformance` runs the cases of `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`. +needs no `skip` entry for it. The `sessions` cases also check the tenants, the +`memberId` filter, the hard delete, and compare-and-set. To check one session store alone, skip the stores that the suite reads by default: @@ -1399,4 +1433,4 @@ To annotate the value instead, use a named shape: - [Build a chat adapter](./build-your-own-chat-adapter): these contracts implemented against SQLite. - [Build a generation adapter](./build-your-own-generation-adapter): the generation half. -- [Build your own adapter](./build-your-own-adapter#verify-with-the-conformance-suite): check an implementation against the conformance suite. +- [Build your own adapter](./build-your-own-adapter#prove-it-with-the-conformance-suite): check an implementation against the conformance suite. diff --git a/packages/ai-acp/src/agent/index.ts b/packages/ai-acp/src/agent/index.ts index 2aedd26b39..41a534de02 100644 --- a/packages/ai-acp/src/agent/index.ts +++ b/packages/ai-acp/src/agent/index.ts @@ -6,7 +6,7 @@ import { agent, ndJsonStream, } from '@agentclientprotocol/sdk/experimental/v2' -import { InputRejectedError } from '@tanstack/ai-harness' +import { AccessDeniedError, InputRejectedError } from '@tanstack/ai-harness' import { MediaError, kindOf, mediaPart } from '@tanstack/ai-harness/media' import type { AgentApp, @@ -26,11 +26,19 @@ import type { HarnessHost, HarnessSession, Operation, + Principal, } from '@tanstack/ai-harness' -export interface AcpAgentOptions { - host: HarnessHost +export interface AcpAgentOptions

{ + host: HarnessHost

harness: AnyHarness + /** + * The user that the agent serves: an editor starts the agent for one + * local user. Each session opens as this principal, and `session/list` + * shows the sessions of this principal. Required when the host has + * `access`. + */ + principal?: P /** Reported in `initialize`. Default `'0.0.0'`. */ version?: string /** @@ -237,30 +245,57 @@ const APPROVAL_OPTIONS = [ * store of the session. A file the harness refuses (too big, or a type it * does not take) fails `session/prompt` with an invalid params error. * + * Each session opens as `principal`. On a host with `access`, the host + * checks each action, and a session that the principal may not read is not + * found. Without `principal`, a host with access throws at the start. + * * ACP v2 is a draft, so this API is experimental too. */ -export function createAcpAgent(options: AcpAgentOptions): AgentApp { - const { host, harness } = options +export function createAcpAgent

( + options: AcpAgentOptions

, +): AgentApp { + const { host, harness, principal } = options + if (principal === undefined && host.hasAccess) { + throw new Error( + 'createAcpAgent needs a principal: the host has access. Pass the local user as principal.', + ) + } /** The ids of the sessions that this agent opened. */ const opened = new Set() /** The last pump of each session. The pumps of one session run in order. */ const tails = new Map>() - // The host gives back the live session of a thread, or opens it again after - // it closed, so the agent keeps no session objects. + // The host gives back a handle to the live session of a thread, or opens + // it again after it closed, so the agent keeps no session objects. const live = (sessionId: string) => - host.open(harness, { threadId: sessionId }) + host.open(harness, { threadId: sessionId, principal }) - /** The session of an id that this agent or the session index knows. */ + /** + * The session of an id that this agent or the session index knows. A + * subagent child is not a session of its own: ACP reaches it through its + * parent. A session that the principal may not read gets the same answer + * as an unknown id. + */ const known = async (sessionId: string) => { - if ( - !opened.has(sessionId) && - (await host.sessions.get(sessionId))?.harness !== harness.name - ) { - throw RequestError.resourceNotFound(sessionId) + try { + if (!opened.has(sessionId)) { + const entry = await host.sessions.get({ + threadId: sessionId, + principal, + }) + const isOwnSession = + entry?.harness === harness.name && entry.parentThreadId === undefined + if (!isOwnSession) throw RequestError.resourceNotFound(sessionId) + } + const session = await live(sessionId) + opened.add(sessionId) + return session + } catch (error) { + if (error instanceof AccessDeniedError) { + throw RequestError.resourceNotFound(sessionId) + } + throw error } - opened.add(sessionId) - return live(sessionId) } const notify = ( @@ -451,6 +486,7 @@ export function createAcpAgent(options: AcpAgentOptions): AgentApp { const cwd = options.cwd ?? process.cwd() if (params.cwd && params.cwd !== cwd) return { sessions: [] } const page = await host.sessions.list({ + principal, harness: harness.name, limit: LIST_PAGE_SIZE, ...(params.cursor ? { cursor: params.cursor } : {}), @@ -505,6 +541,8 @@ export function createAcpAgent(options: AcpAgentOptions): AgentApp { }) .onRequest('session/close', async ({ params }) => { if (!opened.delete(params.sessionId)) return + // ponytail: a handle has no close of its own yet, so this closes the + // live session of the thread on this host, for every handle. await (await live(params.sessionId)).close() }) .onNotification('session/cancel', async ({ params }) => { @@ -518,10 +556,16 @@ export function createAcpAgent(options: AcpAgentOptions): AgentApp { /** * Serve a harness as an ACP v2 agent over a stream. With no stream, it uses * stdin and stdout as newline-delimited JSON (for editors that start the - * agent as a process). + * agent as a process). The sessions open as `principal`, as in + * `createAcpAgent`. + * + * @example + * ```ts + * serveAcp({ host, harness: assistant, principal: { id: 'local-user' } }) + * ``` */ -export function serveAcp( - options: AcpAgentOptions & { stream?: Stream }, +export function serveAcp

( + options: AcpAgentOptions

& { stream?: Stream }, ): AgentConnection { return createAcpAgent(options).connect(options.stream ?? stdioStream()) } diff --git a/packages/ai-acp/tests/agent-handlers.test.ts b/packages/ai-acp/tests/agent-handlers.test.ts index 3c30efb84b..2f34acc0fc 100644 --- a/packages/ai-acp/tests/agent-handlers.test.ts +++ b/packages/ai-acp/tests/agent-handlers.test.ts @@ -2,9 +2,14 @@ import { afterEach, describe, expect, it, vi } from 'vitest' import { EventType, toolDefinition } from '@tanstack/ai' import { client, ndJsonStream } from '@agentclientprotocol/sdk/experimental/v2' import { memoryPersistence } from '@tanstack/ai-persistence' -import { createHarnessHost, defineHarness } from '@tanstack/ai-harness' +import { + createHarnessHost, + defineHarness, + membersOnly, +} from '@tanstack/ai-harness' import { createAcpAgent, serveAcp, toSessionUpdate } from '../src/agent' import type { AnyTextAdapter, ModelMessage, StreamChunk } from '@tanstack/ai' +import type { HarnessAccess } from '@tanstack/ai-harness' import type { ClientApp, RequestPermissionOutcome, @@ -130,7 +135,8 @@ function model(): AnyTextAdapter { } } -function setup() { +/** With `access`, the host checks each action of each principal. */ +function setup({ access }: { access?: HarnessAccess } = {}) { const remove = vi.fn(async () => ({ removed: true })) const harness = defineHarness({ name: 'test/acp-handlers', @@ -154,7 +160,8 @@ function setup() { }), ], }) - const host = createHarnessHost({ persistence: memoryPersistence() }) + const persistence = memoryPersistence() + const host = createHarnessHost({ persistence, access }) const updates: Array<{ sessionId: string; update: SessionUpdate }> = [] const permissions: Array<{ title: string; toolCallId: string }> = [] // The answer of the editor to each permission request. Default: Reject. @@ -216,6 +223,7 @@ function setup() { ) return { harness, + persistence, host, editor, updates, @@ -696,6 +704,31 @@ describe('ACP agent handlers', () => { }) }) + it('refuses a subagent child: the editor reaches it through its parent', async () => { + const { harness, persistence, host, editor } = setup() + close = () => host.close() + await persistence.stores.sessions.upsert({ + threadId: 'subagent:run-1', + parentThreadId: 'lead', + harness: harness.name, + members: [], + createdAt: 1, + updatedAt: 1, + }) + await editor.connectWith(createAcpAgent({ host, harness }), async (ctx) => { + await ctx.request('initialize', { + protocolVersion: 2, + info: { name: 'editor', version: '1.0.0' }, + }) + await expect( + ctx.request('session/resume', { + sessionId: 'subagent:run-1', + cwd: '/tmp', + }), + ).rejects.toMatchObject({ code: -32002 }) + }) + }) + it('resumes a session by id, cancels a running turn, and closes it', async () => { const { harness, host, editor, idle, running, textOf } = setup() close = () => host.close() @@ -733,6 +766,63 @@ describe('ACP agent handlers', () => { await ctx.notify('session/cancel', { sessionId: 'never-opened' }) }) }) + + describe('on a host with access', () => { + const olga = { id: 'olga', tenantId: 'acme' } + const bea = { id: 'bea', tenantId: 'acme' } + // The agent runs for one local user, so no request is identified. + const access: HarnessAccess = { + identify: () => null, + authorize: membersOnly(), + } + + it('throws at the start without a principal', () => { + const { harness, host } = setup({ access }) + close = () => host.close() + + expect(() => createAcpAgent({ host, harness })).toThrow( + 'createAcpAgent needs a principal: the host has access. Pass the local user as principal.', + ) + }) + + it('opens and lists only the sessions of its principal', async () => { + const { harness, host, editor, idle, textOf } = setup({ access }) + close = () => host.close() + await host.open(harness, { threadId: 'bea-thread', principal: bea }) + const agentApp = createAcpAgent({ host, harness, principal: olga }) + await editor.connectWith(agentApp, async (ctx) => { + await ctx.request('initialize', { + protocolVersion: 2, + info: { name: 'editor', version: '1.0.0' }, + }) + const { sessionId } = await ctx.request('session/new', { + cwd: '/tmp', + mcpServers: [], + }) + await ctx.request('session/prompt', { + sessionId, + prompt: [{ type: 'text', text: 'hello' }], + }) + await vi.waitFor(() => expect(idle(sessionId)).toBeDefined()) + + expect(textOf(sessionId)).toBe('echo: hello') + expect( + await host.sessions.get({ threadId: sessionId, principal: olga }), + ).toMatchObject({ ownerId: 'olga' }) + const listed = await ctx.request('session/list', {}) + expect(listed.sessions.map((session) => session.sessionId)).toEqual([ + sessionId, + ]) + // A session of another user gets the answer of an unknown id. + await expect( + ctx.request('session/resume', { + sessionId: 'bea-thread', + cwd: '/tmp', + }), + ).rejects.toMatchObject({ code: -32002 }) + }) + }) + }) }) describe('serveAcp', () => { diff --git a/packages/ai-harness-cli/package.json b/packages/ai-harness-cli/package.json index 95a749a9b1..76c4920ab6 100644 --- a/packages/ai-harness-cli/package.json +++ b/packages/ai-harness-cli/package.json @@ -83,6 +83,7 @@ "devDependencies": { "@tanstack/ai": "workspace:*", "@tanstack/ai-harness": "workspace:*", + "@tanstack/ai-mcp": "workspace:*", "@tanstack/ai-persistence": "workspace:*", "@vitest/coverage-v8": "4.1.10", "zod": "^4.2.0" diff --git a/packages/ai-harness-cli/src/index.ts b/packages/ai-harness-cli/src/index.ts index 0bb0170dd0..1f07d97974 100644 --- a/packages/ai-harness-cli/src/index.ts +++ b/packages/ai-harness-cli/src/index.ts @@ -35,7 +35,7 @@ function needs(flag: string, name: string): Error { ) } -export interface RunCliOptions { +export interface RunCliOptions

{ /** Where sessions keep state. Default: in memory. Do not set it with `host`. */ persistence?: HarnessPersistence /** @@ -46,14 +46,17 @@ export interface RunCliOptions { /** * Your own host, so the CLI keeps its hooks: leases, recovery, and stores. * You own it: the CLI does not close it. Do not set it with `persistence` - * or `credentials`. + * or `credentials`. With host `access`, give `principal` too. `--serve` + * then reads the user of each request with `access.identify`, and `--mcp` + * refuses to start, as stdio has no login. */ - host?: HarnessHost + host?: HarnessHost

/** - * Who uses the CLI. Sessions open as this principal, so credentials and - * provider keys are theirs. `--serve` authorizes each request as them. + * Who uses the CLI. Sessions open as this principal in every mode except + * `--dashboard`, so credentials and provider keys are theirs. On a host + * without access, `--serve` runs each request with the token as them. */ - principal?: Principal + principal?: P /** Default: `process.argv.slice(2)`. */ argv?: ReadonlyArray stdin?: NodeJS.ReadStream @@ -117,9 +120,9 @@ function attachingSource(session: HarnessSession) { * process.exitCode = await runCli(studio) * ``` */ -export async function runCli( +export async function runCli

( harness: AnyHarness, - options: RunCliOptions = {}, + options: RunCliOptions

= {}, ): Promise { const stdout = options.stdout ?? process.stdout const stderr = options.stderr ?? process.stderr @@ -146,7 +149,7 @@ export async function runCli( ) return EXIT.failed } - const host = + const host: HarnessHost

= options.host ?? createHarnessHost({ ...(options.persistence ? { persistence: options.persistence } : {}), @@ -161,12 +164,18 @@ export async function runCli( 'serveAcp', ) if (!acp) throw needs('--acp', '@tanstack/ai-acp') - const connection = acp.serveAcp({ host, harness }) + const connection = acp.serveAcp({ host, harness, principal }) await connection.closed return EXIT.ok } if (args.mcp) { + // A host with access reads the user of each HTTP request. + if (host.hasAccess) { + throw new Error( + '--mcp serves over stdio, which has no login, so it cannot use a host with access. Give runCli a host without access, and the local user as principal.', + ) + } const [mcp, stdio] = await Promise.all([ loadOptional( '@tanstack/ai-mcp/harness', @@ -181,6 +190,8 @@ export async function runCli( const server = await mcp.createHarnessMcpServer({ host, harness, + // One local user: every call runs as them. + openAccess: principal ? { principal } : true, threadId: args.thread, approvals, // A local client can attach the files of the working folder by path. @@ -203,6 +214,8 @@ export async function runCli( ) if (!dashboard) throw needs('--dashboard', '@tanstack/ai-dashboard') const savedToken = env.HARNESS_DASHBOARD_TOKEN + // ponytail: the dashboard takes no principal, so its sessions open + // without one. const connection = await dashboard.connectDashboard({ host, harness, diff --git a/packages/ai-harness-cli/src/optional.ts b/packages/ai-harness-cli/src/optional.ts index 2fc807be6a..5c4584ba97 100644 --- a/packages/ai-harness-cli/src/optional.ts +++ b/packages/ai-harness-cli/src/optional.ts @@ -1,4 +1,5 @@ -import type { AnyHarness, HarnessHost } from '@tanstack/ai-harness' +import type { AnyHarness, HarnessHost, Principal } from '@tanstack/ai-harness' +import type { OpenAccess } from '@tanstack/ai-harness/server' // The packages behind `--acp`, `--mcp`, `/mcp`, and `--dashboard` are // optional. The specifier is a variable, so a build or a type check never @@ -37,9 +38,11 @@ export interface McpFetchServer { /** `@tanstack/ai-mcp/harness` */ export interface McpHarnessModule { - createHarnessMcpServer: (options: { - host: HarnessHost + createHarnessMcpServer:

(options: { + host: HarnessHost

harness: AnyHarness + /** For a host without access: who the calls run as. */ + openAccess?: OpenAccess

threadId?: string approvals?: 'ask' | 'auto' filePaths?: Array @@ -53,15 +56,23 @@ export interface McpStdioModule { /** `@tanstack/ai-acp/agent` */ export interface AcpAgentModule { - serveAcp: (options: { host: HarnessHost; harness: AnyHarness }) => { + serveAcp:

(options: { + host: HarnessHost

+ harness: AnyHarness + /** The local user. Required on a host with access. */ + principal?: P + }) => { closed: Promise } } -/** `@tanstack/ai-dashboard/connect` */ +/** + * `@tanstack/ai-dashboard/connect`. It takes no principal: its sessions + * open without one. + */ export interface DashboardConnectModule { - connectDashboard: (options: { - host: HarnessHost + connectDashboard:

(options: { + host: HarnessHost

harness: AnyHarness url: string threads: Array diff --git a/packages/ai-harness-cli/src/serve.ts b/packages/ai-harness-cli/src/serve.ts index 0a709c9f49..a695efc495 100644 --- a/packages/ai-harness-cli/src/serve.ts +++ b/packages/ai-harness-cli/src/serve.ts @@ -9,36 +9,69 @@ import type { ReadableStream as NodeReadableStream } from 'node:stream/web' import type { AnyHarness, HarnessHost, Principal } from '@tanstack/ai-harness' import type { McpHarnessModule } from './optional' -export interface ServeOptions { - host: HarnessHost +export interface ServeOptions

{ + /** + * The host. With `access`, `access.identify` reads the user of each + * request that has the token. + */ + host: HarnessHost

harness: AnyHarness port: number hostname: string - /** Every request must send `Authorization: Bearer `. */ + /** + * Every request must send `Authorization: Bearer `. Only a signed + * media URL needs no token: its signature is its login. + */ token: string /** The conversation of an MCP tool call that names none. Default `'main'`. */ threadId?: string /** How `/mcp` handles tool calls that need approval. Default `'ask'`. */ approvals?: 'ask' | 'auto' - /** Who a request with the token is. Default `{ id: 'cli' }`. */ - principal?: Principal + /** + * On a host without access: who a request with the token is. Without it, + * each request runs with no principal (one user). + */ + principal?: P } -/** The MCP server for `/mcp`, or `undefined` without `@tanstack/ai-mcp`. */ -async function mcpServer(options: ServeOptions) { +/** + * The MCP server for `/mcp`, or `undefined` without `@tanstack/ai-mcp`. It + * runs each call as the same user as the session protocol routes. + */ +async function mcpServer

( + options: ServeOptions

, + withAccess: boolean, +) { const mcp = await loadOptional( '@tanstack/ai-mcp/harness', 'createHarnessMcpServer', ) if (mcp === undefined) return undefined - return mcp.createHarnessMcpServer({ - host: options.host, - harness: options.harness, + const { host, harness, principal } = options + const shared = { + host, + harness, threadId: options.threadId, approvals: options.approvals, + } + // On a host with access, `identify` reads the user of each request. + if (withAccess) return mcp.createHarnessMcpServer(shared) + return mcp.createHarnessMcpServer({ + ...shared, + openAccess: principal ? { principal } : true, }) } +/** A signed media URL is its own login: the handler checks its signature. */ +function isSignedMedia(request: Request) { + const url = new URL(request.url) + return ( + request.method === 'GET' && + url.pathname.split('/').at(-1) === 'media' && + url.searchParams.has('sig') + ) +} + /** A random token for `--serve` when none is given. */ export function createToken(): string { return randomBytes(24).toString('base64url') @@ -89,31 +122,36 @@ async function send(res: ServerResponse, response: Response): Promise { /** * Serve the harness session protocol over HTTP (capabilities, run, events, * control, snapshot). With `@tanstack/ai-mcp` installed, `/mcp` serves the - * harness as an MCP server. Every request needs the bearer token. + * harness as an MCP server. Every request needs the bearer token, except a + * signed media URL. Then on a host with access, `identify` reads the user of + * the request, and the host checks each action. Without access, each + * request runs as `principal`. */ -export async function serve( - options: ServeOptions, +export async function serve

( + options: ServeOptions

, ): Promise<{ url: string; close: () => Promise }> { + const { host, harness, principal } = options const isAllowed = (request: Request) => sameToken(request.headers.get('authorization'), options.token) - const principal = options.principal ?? { id: 'cli' } - const handler = createHarnessHandler({ - host: options.host, - harness: options.harness, - authorize: (request) => (isAllowed(request) ? principal : null), - // One principal holds the token, so it can use every thread. - canAccess: () => true, - // ponytail: the token signs media URLs, so they work as long as it does. - media: { secret: options.token }, - }) - const mcp = await mcpServer(options) + const withAccess = host.hasAccess + // ponytail: the token signs media URLs, so they work as long as it does. + const media = { secret: options.token } + const handler = withAccess + ? createHarnessHandler({ host, harness, media }) + : createHarnessHandler({ + host, + harness, + // One principal holds the token, so it can use every thread. + openAccess: principal ? { principal } : true, + media, + }) + const mcp = await mcpServer(options, withAccess) const route = (request: Request) => { - const isMcp = mcp !== undefined && new URL(request.url).pathname === '/mcp' - if (!isMcp) return handler(request) - if (!isAllowed(request)) { + if (!isAllowed(request) && !isSignedMedia(request)) { return Response.json({ error: 'unauthorized' }, { status: 401 }) } - return mcp.fetch(request) + const isMcp = mcp !== undefined && new URL(request.url).pathname === '/mcp' + return isMcp ? mcp.fetch(request) : handler(request) } const server = createServer((req, res) => { const aborter = new AbortController() diff --git a/packages/ai-harness-cli/tests/cli.test.ts b/packages/ai-harness-cli/tests/cli.test.ts index 1a3172997a..20aced812d 100644 --- a/packages/ai-harness-cli/tests/cli.test.ts +++ b/packages/ai-harness-cli/tests/cli.test.ts @@ -9,7 +9,9 @@ import { defineCommand, defineHarness, definePlugin, + membersOnly, } from '@tanstack/ai-harness' +import { createMCPClient } from '@tanstack/ai-mcp' import { memoryPersistence } from '@tanstack/ai-persistence' import { EXIT, parseCliArgs, runCli } from '../src' import { serve } from '../src/serve' @@ -132,6 +134,22 @@ function capture() { } } +const olga = { id: 'olga' } +const cid = { id: 'cid' } +const users = new Map([olga, cid].map((user) => [user.id, user])) + +/** A membersOnly() host whose users come from the `x-user` header. */ +function accessHost() { + return createHarnessHost({ + persistence: memoryPersistence(), + access: { + identify: ({ request }) => + users.get(request.headers.get('x-user') ?? '') ?? null, + authorize: membersOnly(), + }, + }) +} + /** A stdin that reads `lines`, as a terminal or as a pipe. */ function stdinFrom(lines: Array, { isTTY }: { isTTY: boolean }) { // NodeJS.ReadStream is a TTY socket. A test cannot make one without a real @@ -521,6 +539,82 @@ describe('a host and a principal', () => { await host.close() } }) + + it('serves /mcp as the principal', async () => { + const host = createHarnessHost({ persistence: memoryPersistence() }) + const server = await serve({ + host, + harness: defineHarness({ + name: 'test/serve-mcp', + adapter: scripted([textTurn('ok')]).adapter, + }), + port: 0, + hostname: '127.0.0.1', + token: 'secret', + principal: { id: 'ada' }, + }) + const client = await createMCPClient({ + transport: { + type: 'http', + url: `${server.url}/mcp`, + headers: { authorization: 'Bearer secret' }, + }, + }) + try { + await client.callTool('chat', { threadId: 't-mcp', message: 'hi' }) + + expect(await host.sessions.get({ threadId: 't-mcp' })).toMatchObject({ + ownerId: 'ada', + }) + } finally { + await client.close() + await server.close() + await host.close() + } + }) + + it('passes the principal to --acp, so a host with access starts', async () => { + const host = accessHost() + // The agent reads process.stdin. Ending it ends the agent. + const ends: Array<() => void> = [] + const on = vi + .spyOn(process.stdin, 'on') + .mockImplementation((event: string | symbol, listener: () => void) => { + if (event === 'end') ends.push(listener) + return process.stdin + }) + const stderr = capture() + try { + const running = runCli( + defineHarness({ name: 'test/acp', adapter: scripted([]).adapter }), + { argv: ['--acp'], stdout: capture(), stderr, host, principal: olga }, + ) + await vi.waitFor(() => expect(ends).toHaveLength(1)) + for (const end of ends) end() + + expect(await running).toBe(EXIT.ok) + expect(stderr.text).toBe('') + } finally { + on.mockRestore() + await host.close() + } + }) + + it('refuses --mcp on a host with access, as stdio has no login', async () => { + const host = accessHost() + const stderr = capture() + + const code = await runCli( + defineHarness({ name: 'test/mcp', adapter: scripted([]).adapter }), + { argv: ['--mcp'], stdout: capture(), stderr, host, principal: olga }, + ) + + expect(code).toBe(EXIT.failed) + expect(stderr.text).toBe( + '--mcp serves over stdio, which has no login, so it cannot use a host with access. Give runCli a host without access, and the local user as principal.\n', + ) + await host.close() + }) }) describe('serve mode', () => { @@ -563,6 +657,51 @@ describe('serve mode', () => { await host.close() } }) + + it('reads the user of each request with identify on a host with access, behind the token', async () => { + const host = accessHost() + const server = await serve({ + host, + harness: defineHarness({ + name: 'test/serve-access', + adapter: scripted([textTurn('ok')]).adapter, + }), + port: 0, + hostname: '127.0.0.1', + token: 'secret', + }) + const call = (path: string, headers: Record) => + fetch(`${server.url}/${path}`, { headers }) + const prompt = (headers: Record) => + fetch(`${server.url}/control`, { + method: 'POST', + headers: { 'content-type': 'application/json', ...headers }, + body: JSON.stringify({ + threadId: 't1', + input: { op: 'prompt', message: 'hi' }, + }), + }) + const token = { authorization: 'Bearer secret' } + try { + expect((await prompt({ ...token, 'x-user': 'olga' })).status).toBe(200) + + const noToken = await call('transcript?threadId=t1', { 'x-user': 'olga' }) + expect(noToken.status).toBe(401) + const noUser = await call('transcript?threadId=t1', token) + expect(noUser.status).toBe(401) + const stranger = await call('transcript?threadId=t1', { + ...token, + 'x-user': 'cid', + }) + expect(stranger.status).toBe(403) + expect( + await host.sessions.get({ threadId: 't1', principal: olga }), + ).toMatchObject({ ownerId: 'olga' }) + } finally { + await server.close() + await host.close() + } + }) }) describe('custom ui', () => { diff --git a/packages/ai-harness/src/access.ts b/packages/ai-harness/src/access.ts new file mode 100644 index 0000000000..45d746bc98 --- /dev/null +++ b/packages/ai-harness/src/access.ts @@ -0,0 +1,212 @@ +import type { SessionMember } from '@tanstack/ai-persistence' +import type { Principal } from './types' + +/** + * What a principal tries to do with a session. + * + * - `read`: stream, snapshot, transcript, media, and child runs. + * - `list`: see the session in a list of sessions. + * - `send`: prompt, steer, follow up, start agents, and run commands. + * - `answer`: answer a question or an approval of a run. + * - `cancel`: stop a run. + * - `configure`: change the settings of the thread. + * - `share`: add or remove members. + * - `delete`: delete the session. + * - `fork`: copy the session into a new session. + * - `transfer`: give the session to a new owner. + */ +export type AccessAction = + | 'read' + | 'send' + | 'answer' + | 'cancel' + | 'configure' + | 'share' + | 'delete' + | 'fork' + | 'transfer' + | 'list' + +/** The session that an action targets, as the index row has it. */ +export interface AccessSession { + threadId: string + tenantId?: string + /** The user id of the owner. Without it, the session has no owner. */ + ownerId?: string + /** The thread that started this one. Set on a subagent child. */ + parentThreadId?: string + members: ReadonlyArray +} + +/** The run that an `answer` or `cancel` targets: who started it. */ +export interface AccessRun { + runId: string + /** The id of the principal who started the run. */ + principalId?: string +} + +/** What `authorize` gets for one action. */ +export interface AuthorizeInput

{ + principal: P + action: AccessAction + session: AccessSession + /** The run that an `answer` or `cancel` targets, when the host knows it. */ + run?: AccessRun +} + +/** Decide if a principal may do one action. Return `false` to refuse. */ +export type Authorize

= ( + input: AuthorizeInput

, +) => boolean | Promise + +/** + * Who may use the sessions of a host, and what they may do. The type of the + * principal comes from the return value of `identify`. + */ +export interface HarnessAccess

{ + /** + * Who is calling: your own login check, once per request. Return null to + * refuse with 401. + */ + identify: (input: { request: Request }) => P | null | Promise

+ /** + * What they may do. Called for every action. The harness checks the + * tenant before it calls this. + */ + authorize: Authorize

+ /** + * Let keyed (BYOK) adapters and `ctx.keys` fall back to the server's env + * vars. Default `false` when access is on. + */ + envKeys?: boolean +} + +/** + * Thrown when the host refuses an action: `authorize` returned `false`, or + * `host.revoke` took the access of the handle's principal. The HTTP handler + * answers it with 403. Catch it with `instanceof AccessDeniedError`. + */ +export class AccessDeniedError extends Error { + readonly code = 'forbidden' + readonly action: AccessAction + readonly threadId: string + + constructor({ + action, + threadId, + }: { + action: AccessAction + threadId: string + }) { + super(`Access denied: ${action} on thread ${JSON.stringify(threadId)}.`) + this.name = 'AccessDeniedError' + this.action = action + this.threadId = threadId + } +} + +/** Options for {@link membersOnly}. */ +export interface MembersOnlyOptions { + fork?: { + /** + * Who is a member of the new session after a fork. + * + * - `'self'` (default): only the user who forked it, as the owner. + * - `'all'`: the members of the source session too. + */ + members?: 'self' | 'all' + } +} + +/** + * The options of a {@link membersOnly} authorize, with the defaults filled + * in. The host reads it with {@link membersOnlyPolicyOf}. + */ +export interface MembersOnlyPolicy { + readonly fork: { readonly members: 'self' | 'all' } +} + +type Role = SessionMember['role'] + +/** + * The role of a user in a session. The `ownerId` makes a user the owner, + * also when `members` has no entry for that user. + */ +function roleOf(session: AccessSession, userId: string) { + if (session.ownerId === userId) return 'owner' + return session.members.find((member) => member.userId === userId)?.role +} + +function allows(role: Role, { principal, action, run }: AuthorizeInput) { + switch (action) { + case 'read': + case 'list': + return true + case 'send': + case 'fork': + return role !== 'viewer' + case 'answer': + case 'cancel': { + const startedTheRun = run?.principalId === principal.id + return role === 'owner' || (role === 'member' && startedTheRun) + } + case 'configure': + case 'share': + case 'delete': + case 'transfer': + return role === 'owner' + } +} + +/** + * A ready-made `authorize`: only the members of a session may use it, each + * with the limits of the member's role. + * + * - Owner: every action. + * - Member: read, list, send, and fork. Answer and cancel only the runs that + * the member started. + * - Viewer: read and list. + * - Not a member: no action. + * + * The owner is the `ownerId` of the session, or a member with the role + * `owner`. This function does not check the tenant: the host does that + * before it calls `authorize`. For a child thread, the host passes the root + * session, so this function has no special rules for children. + * + * The function also has a `membersOnly` property with its options. The host + * reads it to list only the sessions of a member, and to get the default of + * a fork. + * + * @param options - `fork.members` sets who is a member of a forked session. + * The default is `'self'`. + * + * @example + * access: { + * identify: ({ request }) => getUser(request), + * authorize: membersOnly({ fork: { members: 'self' } }), + * } + */ +export function membersOnly(options: MembersOnlyOptions = {}) { + const marker: { readonly membersOnly: MembersOnlyPolicy } = { + membersOnly: { fork: { members: options.fork?.members ?? 'self' } }, + } + // A new function on each call, so each one keeps its own marker. + const authorize: Authorize = (input) => { + const role = roleOf(input.session, input.principal.id) + if (role === undefined) return false + return allows(role, input) + } + return Object.assign(authorize, marker) +} + +/** + * The options of an `authorize` that {@link membersOnly} made, or + * `undefined` for any other `authorize`. + * + * @param authorize - The `authorize` of the host's `access`. + */ +export function membersOnlyPolicyOf

( + authorize: Authorize

& { readonly membersOnly?: MembersOnlyPolicy }, +) { + return authorize.membersOnly +} diff --git a/packages/ai-harness/src/auth.ts b/packages/ai-harness/src/auth.ts index b1c5fb1362..07e84d30f3 100644 --- a/packages/ai-harness/src/auth.ts +++ b/packages/ai-harness/src/auth.ts @@ -33,8 +33,8 @@ export class AuthRequiredError extends Error { } /** - * Credentials of the running turn's sender (outside a turn, of the session's - * principal), as plugins see them. + * The credentials of one user: the user of a run (`ctx.context.run`), or of + * a command. Reads and saves go to that user only. */ export interface CredentialsAccess { get: (id: string) => Promise @@ -94,19 +94,17 @@ function secretValues(credential: Credential) { } /** - * 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. + * Scope credentials to one principal and thread. * * 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), + scope: Scope, onMissing: (error: AuthRequiredError) => void, hooks?: CredentialHooks, ): CredentialsAccess { - const scopeNow = typeof scope === 'function' ? scope : () => scope const secrets = hooks?.secrets const keep = (credential: Credential | null) => { for (const value of credential ? secretValues(credential) : []) { @@ -117,8 +115,7 @@ export function credentialsFor( // 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) => - keep(await readCredential(store, scopeNow(), id)) + const get = async (id: string) => keep(await readCredential(store, scope, id)) const access: CredentialsAccess = { get, require: async (id, options) => { @@ -131,14 +128,13 @@ export function credentialsFor( }, set: async ({ id, credential }) => { keep(credential) - await store.set({ scope: scopeNow(), id, credential }) + await store.set({ scope, id, credential }) await hooks?.onSet?.(id) }, - delete: (id) => store.delete({ scope: scopeNow(), id }), + delete: (id) => store.delete({ scope, id }), list: async () => { - const user = scopeNow() - const tenant = tenantScopeOf(user) - const own = await store.list(user) + const tenant = tenantScopeOf(scope) + const own = await store.list(scope) if (!tenant) return own const shared = await store.list(tenant) return [ @@ -151,22 +147,24 @@ export function credentialsFor( } /** - * The provider keys of the session's principal. A key is the `api_key` + * The provider keys of the user of `credentials`. 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 `. - * A key read from env goes in `secrets`, like a saved credential. + * the first `provider.env` name that is set. With `envFallback: false`, env + * vars are not read. A missing key calls `onMissing` and throws + * {@link AuthRequiredError}, so the user sees `/connect `. A key read + * from env goes in `secrets`, like a saved credential. */ export function providerKeysFor( credentials: CredentialsAccess, onMissing: (error: AuthRequiredError) => void, - secrets?: Set, + { secrets, envFallback }: { secrets?: Set; envFallback: boolean }, ) { 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 + if (!envFallback) return null const fromEnv = envKeyEntry(provider)?.value if (fromEnv) secrets?.add(fromEnv) return fromEnv ?? null diff --git a/packages/ai-harness/src/client.ts b/packages/ai-harness/src/client.ts index 4e973ffdb8..79a071c84d 100644 --- a/packages/ai-harness/src/client.ts +++ b/packages/ai-harness/src/client.ts @@ -522,7 +522,9 @@ export function createHarnessClient( } } catch (error) { if (signal?.aborted) return - if (error instanceof Error && /\((401|403|404)\)/.test(error.message)) + // A 404 means that the thread has no entry yet: a new chat, before + // its first message. Try again, the same as after a network error. + if (error instanceof Error && /\((401|403)\)/.test(error.message)) throw error } if (signal?.aborted) return diff --git a/packages/ai-harness/src/commands.ts b/packages/ai-harness/src/commands.ts index 3c1b62bb82..767bbe1cc1 100644 --- a/packages/ai-harness/src/commands.ts +++ b/packages/ai-harness/src/commands.ts @@ -8,6 +8,7 @@ import type { Operation, Principal, Receipt, + RunContext, ThreadSettings, } from './types' @@ -39,16 +40,18 @@ export type AnswerOf = TSchema extends SchemaInput 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. + * The user this API acts for. In a command: who runs the command. In a + * turn plugin: the sender of its turn. In a session plugin: the principal + * that opened the session, never the sender of a running turn. For the + * user of a tool call, read `ctx.context.run.principal`. */ 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. + * In a command or a turn plugin, the turn runs as `principal`. In a session + * plugin, it runs as the sender of the running turn, else as `principal`. */ send: (text: string) => Operation /** @@ -124,9 +127,14 @@ export interface CommandContext { principal?: Principal /** * The credentials of the user who runs the command. Save a sign-in here, - * so it belongs to that user. + * so it belongs to that user. The same as `run.credentials`. */ credentials: CredentialsAccess + /** + * The run of the command: the user who runs it, and that user's + * credentials and model provider keys (`run.keys`). + */ + run: RunContext } /** A user action: a slash command, a button, a dashboard action. Not a model tool. */ diff --git a/packages/ai-harness/src/connectors.ts b/packages/ai-harness/src/connectors.ts index bd5639d999..b70675e211 100644 --- a/packages/ai-harness/src/connectors.ts +++ b/packages/ai-harness/src/connectors.ts @@ -8,6 +8,7 @@ import { } from './oauth' import type { AnyTool } from '@tanstack/ai' import type { OAuthConfig } from './oauth' +import type { RunContext } from './types' export interface OAuthConnectorOptions { /** A short id, for example `'github'`. Commands are `connect:` and `disconnect:`. */ @@ -26,13 +27,14 @@ export interface OAuthConnectorOptions { */ login: 'loopback' | 'device' /** - * The tools of this service. Call `token()` inside a tool: it returns a - * fresh access token, or fails with `auth_required` before sign-in. With - * `token({ wait: true })`, the turn waits for the sign-in and the tool - * runs again after it (see `CredentialsAccess.require`). + * The tools of this service. Call `token({ run: ctx.context.run })` inside + * a tool (typed `.server`): it returns a fresh access + * token of the user of that call, or fails with `auth_required` before + * that user signs in. With `wait: true`, the turn waits for the sign-in and + * the tool runs again after it (see `CredentialsAccess.require`). */ tools?: ( - token: (options?: { wait?: boolean }) => Promise, + token: (options: { run: RunContext; wait?: boolean }) => Promise, ) => ReadonlyArray /** Test hook for the token endpoint. */ fetch?: typeof fetch @@ -50,7 +52,11 @@ export interface OAuthConnectorOptions { * label: 'GitHub', * login: 'device', * oauth: { authorizationUrl, tokenUrl, deviceUrl, clientId, scopes: ['repo'] }, - * tools: (token) => [listIssues(token)], + * tools: (token) => [ + * listIssues.server(async (_args, ctx) => + * fetchIssues(await token({ run: ctx.context.run })), + * ), + * ], * }) * ``` */ @@ -59,14 +65,20 @@ export function oauthConnector(options: OAuthConnectorOptions) { return definePlugin({ name: `connector/${id}`, setup: (ctx) => { - const token = async (tokenOptions?: { + // The token of the user of the run that calls, never of another run. + const token = async ({ + run, + wait, + }: { + run: RunContext wait?: boolean - }): Promise => { - let credential = await ctx.credentials.require(id, tokenOptions) + }) => { + const { credentials } = run + let credential = await credentials.require(id, { wait }) if (credential.type === 'api_key') return credential.value if (isExpired(credential) && credential.refreshToken) { credential = await refreshCredential(oauth, credential, options.fetch) - await ctx.credentials.set({ id, credential }) + await credentials.set({ id, credential }) } if (credential.type === 'api_key') return credential.value return credential.accessToken diff --git a/packages/ai-harness/src/first-party/coding/workspace.ts b/packages/ai-harness/src/first-party/coding/workspace.ts index 6e0396fa43..37dfbd8fbb 100644 --- a/packages/ai-harness/src/first-party/coding/workspace.ts +++ b/packages/ai-harness/src/first-party/coding/workspace.ts @@ -390,6 +390,8 @@ export function workspaceTools(options: WorkspaceToolsOptions) { const cwd = () => ctx.session.settings().cwd const { tools, prompt, resources } = createWorkspaceTools(options, { access: { + // ponytail: D9, PR 7 rebuilds ask. The question goes to the + // running turn, not to the run of the tool call that asks. ask: (message) => ctx.session.ask({ message }), mode: () => ctx.config.get('mode'), cwd, @@ -404,6 +406,8 @@ export function workspaceTools(options: WorkspaceToolsOptions) { }, }, hooks: () => hooks, + // ponytail: D9, PR 7 rebuilds ask and note. The wake turn runs as the + // sender of the running turn, not as the user of the job's turn. note: (text) => ctx.session.note(text, { wake: true }), // Kills the background `bash` jobs when the session closes. A reload // keeps them. Without a session, the plugin cleanup kills them. diff --git a/packages/ai-harness/src/first-party/files.ts b/packages/ai-harness/src/first-party/files.ts index 49623a4fbd..3cccec4766 100644 --- a/packages/ai-harness/src/first-party/files.ts +++ b/packages/ai-harness/src/first-party/files.ts @@ -291,7 +291,7 @@ export async function readMarkdownFiles(dir: string) { export function fileCommands(options: { dir: string }) { return definePlugin({ name: 'tanstack/file-commands', - setup: async (ctx) => { + setup: async () => { const files = await readMarkdownFiles(options.dir) const commands: Record = {} for (const { name, fields, body } of files) { @@ -301,14 +301,15 @@ export function fileCommands(options: { dir: string }) { typeof description === 'string' && description !== '' ? description : `Run ${name}.md`, - run: (input: unknown) => { + // The prompt runs as the user who runs the command. + run: (input: unknown, { session }) => { const args = typeof input === 'string' ? input : input === undefined ? '' : JSON.stringify(input) - ctx.session.send(body.split('$ARGUMENTS').join(args).trim()) + session.send(body.split('$ARGUMENTS').join(args).trim()) return `Sent /${name}.` }, }) diff --git a/packages/ai-harness/src/first-party/goal.ts b/packages/ai-harness/src/first-party/goal.ts index fc59812e59..a18c603af4 100644 --- a/packages/ai-harness/src/first-party/goal.ts +++ b/packages/ai-harness/src/first-party/goal.ts @@ -3,8 +3,9 @@ import { defineCommand } from '../commands' import { createPluginEvent } from '../extensions' import { definePlugin } from '../plugins' import { textOf, transcriptText } from './session-tools' -import type { AnyTextAdapter, KeyedAdapter } from '@tanstack/ai' -import type { ChatTurnResult, Operation } from '../types' +import type { AnyTextAdapter, ChatMiddleware, KeyedAdapter } from '@tanstack/ai' +import type { PluginSessionApi } from '../commands' +import type { ChatTurnResult, HarnessRunContext, Operation } from '../types' /** The goal of a session. It is plugin state, so it survives restarts. */ export interface Goal { @@ -120,8 +121,15 @@ export function goal(options: { // The turn the plugin queued, until it starts. let queuedTurn: Operation | undefined - const queue = (message: string) => { - queuedTurn = ctx.session.send(message) + /** + * Queue the next turn of the goal. `/goal` passes its own session, so + * the turn runs as the user who runs it. + */ + // ponytail: D9. From `onFinish`, `ctx.session.send` runs as the sender + // of the running turn: the goal's own turn there. PR 7 rebuilds send + // on the run that calls. + const queue = (message: string, session = ctx.session) => { + queuedTurn = session.send(message) } /** Cancel the queued turn if it has not started. */ @@ -144,7 +152,11 @@ export function goal(options: { } /** Make `text` the active goal with a new round count, and start a turn. */ - const begin = async (text: string, reason: string) => { + const begin = async ( + text: string, + reason: string, + session: PluginSessionApi, + ) => { const message = reason === '' ? `Work toward this goal: ${text}` @@ -160,7 +172,7 @@ export function goal(options: { note: '', queued: message, })) - queue(message) + queue(message, session) } /** Pause the goal when the plugin's own turn ends early. */ @@ -170,13 +182,14 @@ export function goal(options: { await change(current.text, { status: 'paused', note }) } - const judge = async (text: string) => { + /** Check the goal with the keys of the user of the turn. */ + const judge = async (text: string, { run }: HarnessRunContext) => { // ponytail: the last 10 messages, clipped. Summarize first if a judge needs more. const transcript = transcriptText( (await ctx.session.transcript()).slice(-10), ).slice(-8000) const verdict = await chat({ - adapter: await ctx.keys.adapter(options.judge), + adapter: await run.keys.adapter(options.judge), messages: [ { role: 'user', @@ -211,7 +224,7 @@ export function goal(options: { goal: defineCommand({ description: 'Work until a goal is met: /goal . Also /goal, /goal stop, /goal resume', - run: async (input: unknown) => { + run: async (input: unknown, { session }) => { const arg = typeof input === 'string' ? input.trim() : '' switch (arg) { case '': @@ -245,11 +258,11 @@ export function goal(options: { current?.status === 'stopped' if (!current || !canResume) return describe(current, maxRounds) - await begin(current.text, current.reason) + await begin(current.text, current.reason, session) return describe(current, maxRounds) } default: - await begin(arg, '') + await begin(arg, '', session) return describe(current, maxRounds) } }, @@ -301,8 +314,9 @@ export function goal(options: { if (run.runId !== ownRun || saved?.status !== 'active') return ownRun = undefined const { text, round } = saved - const verdict = await judge(text).catch((error: unknown) => - error instanceof Error ? error.message : String(error), + const verdict = await judge(text, run.context).catch( + (error: unknown) => + error instanceof Error ? error.message : String(error), ) if (typeof verdict === 'string') { await change(text, { @@ -338,7 +352,7 @@ export function goal(options: { if (await change(text, { reason, queued: message })) queue(message) }, - }, + } satisfies ChatMiddleware, ], } }, diff --git a/packages/ai-harness/src/first-party/provider-keys.ts b/packages/ai-harness/src/first-party/provider-keys.ts index 9e1fe4b569..3086b12541 100644 --- a/packages/ai-harness/src/first-party/provider-keys.ts +++ b/packages/ai-harness/src/first-party/provider-keys.ts @@ -50,14 +50,17 @@ function untilAborted(promise: Promise, signal: AbortSignal) { ]) } -/** The state of one provider, and its `/keys` line. The key shows masked only. */ +/** + * The state of one provider, and its `/keys` line. The key shows masked only. + * Without `credentials` (no user), only the env var counts. + */ async function describeKey( - credentials: CredentialsAccess, + credentials: CredentialsAccess | undefined, provider: ByokProvider, ) { const { id, label } = provider const status = (state: ProviderKeyStatus['state']) => ({ id, label, state }) - const saved = await credentials.get(id) + const saved = await credentials?.get(id) if (saved?.type === 'api_key') { return { status: status('connected'), @@ -83,7 +86,7 @@ async function describeKey( * `/connect ` and `/disconnect `, and `/keys` lists where each key * comes from. `/connect` runs the provider's `signIn` when it has one (the * host opens the browser), else it asks the user to paste the key. The key - * is saved in the credential store of the session's principal. A + * is saved in the credential store of the user who runs `/connect`. A * `keyedAdapter(...)` for that provider then uses it. The env var still works * when no key is saved. Keys show masked only (the last 4 characters). * @@ -116,9 +119,10 @@ export function providerKeys(options: { const state = ctx.state({ providers: [] }) /** * Check every provider, keep the result for UIs, and return the `/keys` - * lines. A command passes the credentials of the user who runs it. + * lines. A command passes the credentials of the user who runs it. At + * setup there is no user, so the state shows the env vars only. */ - const refresh = async (credentials = ctx.credentials) => { + const refresh = async (credentials?: CredentialsAccess) => { const described = await Promise.all( providers.map((provider) => describeKey(credentials, provider)), ) diff --git a/packages/ai-harness/src/first-party/session-tools.ts b/packages/ai-harness/src/first-party/session-tools.ts index f4de9ea03e..26e6e5aaaf 100644 --- a/packages/ai-harness/src/first-party/session-tools.ts +++ b/packages/ai-harness/src/first-party/session-tools.ts @@ -46,7 +46,8 @@ export function compact(options: { commands: { compact: defineCommand({ description: 'Summarize the conversation to save tokens', - run: async () => { + // The summary call uses the key of the user who runs `/compact`. + run: async (_input, { run }) => { const messages = await ctx.session.transcript() const keep = options.keepLast ?? 0 if (messages.length <= keep + 2) @@ -59,7 +60,7 @@ export function compact(options: { const older = messages.slice(0, cut) const transcript = transcriptText(older) const { text: summary } = await chat({ - adapter: await ctx.keys.adapter(options.adapter), + adapter: await run.keys.adapter(options.adapter), messages: [ { role: 'user', content: `${SUMMARY_PROMPT}\n\n${transcript}` }, ], diff --git a/packages/ai-harness/src/first-party/title.ts b/packages/ai-harness/src/first-party/title.ts index 929e97dad1..df2728b837 100644 --- a/packages/ai-harness/src/first-party/title.ts +++ b/packages/ai-harness/src/first-party/title.ts @@ -3,11 +3,12 @@ import { createPluginEvent } from '../extensions' import { definePlugin } from '../plugins' import { textOf } from './session-tools' import type { - AnyChatMiddleware, AnyTextAdapter, + ChatMiddleware, KeyedAdapter, ModelMessage, } from '@tanstack/ai' +import type { HarnessRunContext } from '../types' const TITLE_PROMPT = 'Write a short title for a conversation that starts with the message below. Use at most 60 characters on one line, with no quotes. Answer with the title only.' @@ -46,7 +47,11 @@ export function title(options: { name: 'tanstack/title', setup: (ctx) => { let isNaming = false - const name = async (messages: ReadonlyArray) => { + /** Name the session with the keys of the user of the turn. */ + const name = async ( + messages: ReadonlyArray, + { run }: HarnessRunContext, + ) => { const first = messages.find((message) => message.role === 'user') const entry = await ctx.session.entry() // No sessions store, or the session has a title (also after a restart). @@ -54,7 +59,7 @@ export function title(options: { // ponytail: the first 2000 characters are enough for a title. const message = textOf(first).slice(0, 2000) const answer = await chat({ - adapter: await ctx.keys.adapter(options.adapter), + adapter: await run.keys.adapter(options.adapter), messages: [ { role: 'user', content: `${TITLE_PROMPT}\n\n${message}` }, ], @@ -71,7 +76,7 @@ export function title(options: { if (run.phase !== 'init' || isNaming) return isNaming = true // Not awaited: the title call never slows or fails the turn. - void name(config.messages) + void name(config.messages, run.context) .catch((error: unknown) => ctx.emit(TitleFailed, { message: @@ -82,7 +87,7 @@ export function title(options: { isNaming = false }) }, - } satisfies AnyChatMiddleware, + } satisfies ChatMiddleware, ], } }, diff --git a/packages/ai-harness/src/harness-text.ts b/packages/ai-harness/src/harness-text.ts index f02517fe44..7b89194a9e 100644 --- a/packages/ai-harness/src/harness-text.ts +++ b/packages/ai-harness/src/harness-text.ts @@ -292,11 +292,9 @@ export function harnessText( const signal = signalOf(chatOptions.request) // A session this call opens closes with it. One the caller holds open // stays, as in host.fork. - const wasOpen = host.isOpen(harness, { threadId }) - const session = await host.open(harness, { - threadId, - ...(principal ? { principal } : {}), - }) + const user = principal ? { principal } : {} + const wasOpen = host.isOpen(harness, { threadId, ...user }) + const session = await host.open(harness, { threadId, ...user }) // The outer call's tools join the inner turn. Its system prompts do // not: the harness owns its prompt stack. const tools = chatOptions.tools?.length diff --git a/packages/ai-harness/src/host.ts b/packages/ai-harness/src/host.ts index 5aff3cfa97..62b69163e4 100644 --- a/packages/ai-harness/src/host.ts +++ b/packages/ai-harness/src/host.ts @@ -2,9 +2,11 @@ import { EventType } from '@tanstack/ai' import { memoryCredentialStore } from '@tanstack/ai/byok/server' import { memoryPersistence } from '@tanstack/ai-persistence' import type { CredentialStore } from '@tanstack/ai/byok' -import { SharedLog, loadLogState } from './log' +import { AccessDeniedError, membersOnlyPolicyOf } from './access' +import { SharedLog, loadLogState, logIdOf, parseLogId } from './log' import { LEASE } from './resume' import { SessionImpl } from './session/harness-session' +import { sessionHandle } from './session/handle' import { HARNESS_EVENTS } from './types' import { isRecord } from './utils' import { changedEntry } from './session-entry' @@ -22,10 +24,19 @@ import type { MetadataStore, WorkClaimStore, SessionIndexEntry, + SessionIndexKey, SessionIndexListOptions, SessionIndexPage, SessionIndexStore, + SessionMember, } from '@tanstack/ai-persistence' +import type { + AccessAction, + AccessRun, + AccessSession, + HarnessAccess, +} from './access' +import type { HandleGuard } from './session/handle' import type { AnyHarness, HarnessLogStateOf, SessionMetadataOf } from './define' import type { SessionEntryChange, SessionEntryUpdate } from './session-entry' import type { ReduceOptions } from './log' @@ -62,8 +73,9 @@ type DurableLeases = * in the inbox for the owner instead (receipt `queued`). * * Without `artifacts`, `blobs`, and `generationRuns`, media is kept in memory. - * With `sessions`, the host keeps one index entry per thread (see - * `host.sessions`). + * With `sessions`, the host keeps one index entry per tenant and thread (see + * `host.sessions`). The store needs `getVersioned` and `upsertIf`, or the host + * throws at the start. */ export type HarnessPersistence = | AIPersistence< @@ -85,9 +97,30 @@ export type HarnessPersistence = SharedStores > -export interface HarnessHostOptions { +export interface HarnessHostOptions

{ /** Where sessions keep their state. Default: in memory, lost on restart. */ persistence?: HarnessPersistence + /** + * Who may use the sessions, and what they may do: `identify` reads the + * user of a request, and `authorize` decides each action. The principal + * type of the host comes from what `identify` returns. With it, every + * public host call needs a `principal`, a session is reached only from + * its own tenant, and keyed adapters do not read env vars unless + * `envKeys` is `true`. It needs `stores.sessions`, where each session + * keeps its owner and members. Without it, there is no check: one user + * runs the host. + * + * @example + * ```ts + * const host = createHarnessHost({ + * access: { + * identify: async ({ request }) => getUser(request), + * authorize: membersOnly(), + * }, + * }) + * ``` + */ + access?: HarnessAccess

/** * Where users' API keys and OAuth tokens are saved: a table in your own * database, or a secrets manager. Encrypt the secrets at rest. Default: in @@ -114,11 +147,17 @@ export interface HarnessHostOptions { ownership?: LeaseOptions } -export interface OpenSessionOptions { +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 + /** + * Who opens the session. The session that `open` returns acts as this + * principal: a call that passes no `principal` stores it on the input. + * Its `tenantId` is the tenant of the session (without it, the default + * tenant): one `threadId` in two tenants is two sessions, with two index + * entries and two logs. Required when the host has `access`. + */ + principal?: P /** * 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 @@ -134,20 +173,48 @@ export interface OpenSessionOptions { } /** - * 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 + * The session index of a host: one entry per tenant and thread, for a list + * of sessions. The host writes the entry when a session opens and after each * turn. Without `stores.sessions`, `list` is empty, `get` is `undefined`, * and `rename` and `delete` write nothing. + * + * Each method takes a `principal`. Its `tenantId` picks the tenant (without + * it, the default tenant), as in `host.open`. A method never reads or + * changes the entries of another tenant. With host `access`, the principal + * is required, and each method asks `authorize` first: `get` reads, `list` + * lists, `rename` and `update` configure, `delete` deletes, `fork` forks, and + * `share` and `unshare` share. A refused action throws `AccessDeniedError`. + * A child thread is checked through its root thread. */ -export interface HostSessions { +export interface HostSessions

{ + /** + * The entries of the tenant of `principal`, newest `updatedAt` first. + * Default: only top-level sessions (`parentThreadId: null`). Pass a thread + * id to get its child sessions. With access, only the entries that + * `authorize` lets the principal list. + * + * @example + * ```ts + * const page = await host.sessions.list({ principal, limit: 20 }) + * ``` + */ + list: ( + options?: Omit & { + principal?: P + }, + ) => Promise /** - * 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. + * The entry of a thread, or `undefined`. + * + * @example + * ```ts + * const entry = await host.sessions.get({ threadId: 't-1', principal }) + * ``` */ - list: (options?: SessionIndexListOptions) => Promise - /** The entry of a thread, or `undefined`. */ - get: (threadId: string) => Promise + get: (options: { + threadId: string + principal?: P + }) => Promise /** * Set the title of a thread. Resolves to the changed entry, or to * `undefined` for a thread that has no entry. @@ -155,6 +222,7 @@ export interface HostSessions { rename: (change: { threadId: string title: string + principal?: P }) => Promise /** * Change the entry of a thread from app code, with no session open: its @@ -182,21 +250,36 @@ export interface HostSessions { harness: THarness threadId: string update: SessionEntryUpdate> + principal?: P }) => 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. + * Remove the index entry of a thread. After it, `get` is `undefined`, and + * the next open of the thread writes a new entry. Only the entry goes: the + * transcript, the log, and the other data of the thread stay in their + * stores. + * + * @example + * ```ts + * await host.sessions.delete({ threadId: 't-1', principal }) + * ``` */ - delete: (threadId: string) => Promise + delete: (options: { threadId: string; principal?: P }) => 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. + * thread gets its own entry in the same tenant, with the title of the old + * one plus ` (fork)`, and the harness. 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. + * + * `members` picks who can use the fork. `'self'`: the user who forks is + * the owner and the only member. `'all'`: that user is the owner, and the + * other members keep their roles (the old owner becomes a `member`). With + * access, the default is the `fork.members` of `membersOnly()`, else + * `'self'`. Without access and without `members`, the fork gets the owner + * and members of the old thread. * * @example * ```ts @@ -204,32 +287,72 @@ export interface HostSessions { * harness: assistant, * threadId: 't-1', * at: { before: messageId }, + * principal, * }) - * const fork = await host.open(assistant, { threadId: entry.threadId }) + * const fork = await host.open(assistant, { threadId: entry.threadId, principal }) * ``` */ fork: (options: { harness: AnyHarness threadId: string at: ForkPoint + members?: 'self' | 'all' + principal?: P }) => Promise + /** + * Add the user `userId` to the members of a thread with `role`, or change + * the role of a member. Needs `membersOnly()` as the host's `authorize`, + * and the `share` action (the owner). Resolves to the changed entry, or + * `undefined` for a thread that has no entry. + * + * @example + * ```ts + * await host.sessions.share({ threadId, userId: 'u3', role: 'member', principal }) + * ``` + */ + share: (options: { + threadId: string + userId: string + role: SessionMember['role'] + principal?: P + }) => Promise + /** + * Remove the user `userId` from the members of a thread. Needs + * `membersOnly()`, and the `share` action. The work that the user started + * on the session in this host stops, and `host.revoke` runs for the user. + * The owner cannot be removed. Resolves to the changed entry, or + * `undefined` for a thread that has no entry. + * + * @example + * ```ts + * await host.sessions.unshare({ threadId, userId: 'u3', principal }) + * ``` + */ + unshare: (options: { + threadId: string + userId: string + principal?: P + }) => 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. + * `tenantId` is the tenant of the session, when it is not the default + * tenant. */ export interface HostStatusEvent { type: 'status' threadId: string + tenantId?: 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. + * turn, a rename, or a fork. `entry` is the new entry, with its `tenantId`. */ export interface HostSessionEvent { type: 'session' @@ -244,13 +367,27 @@ export interface HostSessionDeletedEvent { entry: SessionIndexEntry } +/** + * `host.revoke` took the access of the user `principalId` to a thread, or + * `host.sessions.unshare` removed the user. The streams of that user on the + * thread close. + */ +export interface HostRevokedEvent { + type: 'revoked' + tenantId?: string + threadId: string + principalId: string + at: number +} + /** What `host.events()` yields. */ export type HostEvent = | HostStatusEvent | HostSessionEvent | HostSessionDeletedEvent + | HostRevokedEvent -export interface ForkSessionOptions { +export interface ForkSessionOptions

{ /** The thread to copy. */ threadId: string /** The new thread. It must have no messages yet. */ @@ -261,8 +398,12 @@ export interface ForkSessionOptions { * id the thread does not have throws. */ at?: string | null - /** Who opens the new thread, as in `open`. */ - principal?: Principal + /** + * Who opens the new thread, as in `open`. Its `tenantId` is the tenant of + * both threads. With access, it needs the `fork` action on the source, and + * it owns the new thread. + */ + principal?: P /** The log of the source thread, as in `open`. Default: `threadId`. */ logId?: string } @@ -283,24 +424,42 @@ export interface ResumePendingOptions { limit?: number } -/** Runs sessions for one or more harnesses in this process. */ -export interface HarnessHost { +/** + * Runs sessions for one or more harnesses in this process. `P` is the + * principal type: what `access.identify` returns. + */ +export interface HarnessHost

{ /** - * Open a session of `harness`, or return the live one for this harness and - * thread. Session plugins are set up here, so a plugin error rejects the - * promise. + * Open a session of `harness` for `principal`. Each open gives its own + * handle to the live session of this harness, tenant, and thread: a call + * that passes no `principal` acts as the principal of the open. Session + * plugins are set up on the first open, so a plugin error rejects the + * promise. A subagent child thread (its index entry has a + * `parentThreadId`) does not open: open its root thread, and read the + * child with `session.agentTranscript`. + * + * With access, `principal` is required. A thread that exists needs the + * `read` action, else the open throws `AccessDeniedError`. A new thread + * gets the principal as its owner. Then each action of the handle asks + * `authorize`, and the handle acts only as its principal. * * @example * ```ts - * const session = await host.open(assistant, { threadId: 't-1' }) + * const session = await host.open(assistant, { threadId: 't-1', principal }) * ``` */ open: ( harness: THarness, - options: OpenSessionOptions, + options: OpenSessionOptions

, ) => Promise> - /** Is a session of `harness` open on `threadId` on this host? */ - isOpen: (harness: AnyHarness, options: { threadId: string }) => boolean + /** + * Is a session of `harness` open on `threadId` on this host, in the tenant + * of `principal`? + */ + isOpen: ( + harness: AnyHarness, + options: { threadId: string; principal?: P }, + ) => boolean /** * Start a new thread as a copy of another one, and return its open * session of `harness`. It copies the transcript up to and including the @@ -316,7 +475,7 @@ export interface HarnessHost { */ fork: ( harness: THarness, - options: ForkSessionOptions, + options: ForkSessionOptions

, ) => Promise> /** * Continue the work of hosts that stopped. Needs `stores.workClaims`. It @@ -325,8 +484,9 @@ export interface HarnessHost { * 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 + * Known limits: a thread that this host has open already is claimed, but * its session does not recover again. Call `host.recover(threadId)` for it. + * Each thread opens in its own tenant, with no principal. * * @example * ```ts @@ -370,9 +530,10 @@ export interface HarnessHost { close: (options?: { recoverable?: boolean }) => Promise /** * The `durability.reduce` fold of `harness` for a log that has an open - * session in this host. `undefined` when no session has the log open, or - * when the reduce of another harness folded the log (the harness of its - * first session folds a shared log). + * session in this host. `logId` is the `logId` of `open`, in the tenant of + * `principal`. `undefined` when no session has the log open, or when the + * reduce of another harness folded the log (the harness of its first + * session folds a shared log). * * @example * ```ts @@ -382,9 +543,35 @@ export interface HarnessHost { logState: (options: { harness: THarness logId: string + principal?: P }) => HarnessLogStateOf | undefined - /** The session index: list, rename, delete, and fork sessions. */ - sessions: HostSessions + /** The session index: list, rename, delete, fork, and share sessions. */ + sessions: HostSessions

+ /** + * Take the access of the user `principal` to a thread, after your app + * changed its own access tables. `host.events()` yields a `revoked` + * event, so the user's streams close. On a host with access, the open + * handles of that user on the thread throw `AccessDeniedError` from now + * on, and a new `host.open` asks `authorize` again. The `tenantId` of + * `principal` picks the tenant. Only `id` and `tenantId` are read. + * + * @example + * ```ts + * host.revoke({ threadId, principal: { id: 'u3', tenantId: 'acme' } }) + * ``` + */ + revoke: (options: { + threadId: string + principal: { id: string; tenantId?: string } + }) => void + /** + * The user of `request`, from `access.identify`, or `null` when + * `identify` refuses it. Network code other than the HTTP handler (an MCP + * server, your own route) uses it. Throws on a host with no `access`. + */ + identify: (options: { request: Request }) => Promise

+ /** `true` when the host was made with `access`, so every call needs a principal. */ + readonly hasAccess: boolean /** * 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 @@ -416,39 +603,39 @@ function checkHarness( 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. + * Writes to the session index. A write replaces the whole entry, so a change + * reads the entry with its revision, changes a copy, and writes it back with + * `upsertIf`. A change that another host wrote first runs again on the new + * entry. The changes of one key run one at a time on this host, so no change + * loses the fields of another. `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) => { + const serial = (key: SessionIndexKey, task: () => Promise) => { + const id = logIdOf({ tenantId: key.tenantId, logId: key.threadId }) // The tail never rejects, so a failed change does not stop the next one. - const run = (tails.get(threadId) ?? Promise.resolve()).then(task) + const run = (tails.get(id) ?? Promise.resolve()).then(task) const tail = run.then( () => {}, () => {}, ) - tails.set(threadId, tail) + tails.set(id, tail) void tail.then(() => { - if (tails.get(threadId) === tail) tails.delete(threadId) + if (tails.get(id) === tail) tails.delete(id) }) return run } return { - get: async (threadId: string) => store?.get(threadId), + get: async (key: SessionIndexKey) => store?.get(key), /** - * Write `change(entry)` for the thread. `change` gets `undefined` when - * the thread has no entry, and returns `undefined` to write nothing. + * Write `change(entry)` for the key. `change` gets `undefined` when the + * key has no entry, and returns `undefined` to write nothing. */ update: ( - threadId: string, + key: SessionIndexKey, change: ( entry: SessionIndexEntry | undefined, ) => @@ -456,23 +643,16 @@ function sessionIndex( | undefined | Promise, ) => - serial(threadId, async () => { + serial(key, 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 - } + const { threadId } = key for (let attempt = 0; attempt < 5; attempt += 1) { - const current = await store.getVersioned(threadId) + const current = await store.getVersioned(key) const next = await change(current?.entry) if (!next) return undefined const written = await store.upsertIf({ entry: next, - expectedRevision: current?.revision ?? null, + revision: current?.revision ?? null, }) if (!written.ok) continue onChange({ type: 'session', threadId, entry: next }) @@ -482,16 +662,45 @@ function sessionIndex( `The session index entry of ${threadId} changed on another host 5 times in a row.`, ) }), - remove: (threadId: string) => - serial(threadId, async () => { + /** Remove the entry of the key: a hard delete. */ + remove: (key: SessionIndexKey) => + serial(key, async () => { if (!store) return - const entry = await store.get(threadId) - await store.delete(threadId) - if (entry) onChange({ type: 'session-deleted', threadId, entry }) + const entry = await store.get(key) + await store.delete(key) + if (entry) { + onChange({ type: 'session-deleted', threadId: key.threadId, entry }) + } }), } } +/** + * Throw a clear error for a sessions store without compare-and-set. The + * types need both methods, but a store in plain JS can leave them out. + */ +function checkSessionStore(store: SessionIndexStore | undefined) { + if (store === undefined) return + const hasCompareAndSet = + typeof store.getVersioned === 'function' && + typeof store.upsertIf === 'function' + if (!hasCompareAndSet) { + throw new Error( + 'stores.sessions needs getVersioned and upsertIf. The host changes an index entry with compare-and-set, so two hosts do not lose the changes of each other.', + ) + } +} + +/** + * The key of an open session in the host: the harness, the tenant, and the + * thread. One thread id in two tenants is two sessions. + */ +const sessionKey = ( + harness: string, + tenantId: string | undefined, + threadId: string, +) => `${harness}\u0000${logIdOf({ tenantId, logId: threadId })}` + /** @internal How a session writes its index entry. */ export type SessionIndexWriter = ReturnType @@ -527,26 +736,31 @@ function statusAfter(session: HarnessSession, event: StreamChunk) { } } +/** Where a host looks for the open session of a chat turn. */ +interface TurnLookup { + harness: AnyHarness + /** Only the sessions of this tenant. */ + tenantId: string | undefined + operationId: string +} + /** How each host finds the open session of a chat turn. Not public API. */ const turnFinders = new WeakMap< object, - (harness: string, operationId: string) => Promise + (lookup: TurnLookup) => 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. + * @internal The open session of `harness` in `host`, in the tenant + * `tenantId`, 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, +export function sessionOfTurn

( + host: HarnessHost

, + lookup: TurnLookup, ): Promise { - return ( - turnFinders.get(host)?.(harness.name, operationId) ?? - Promise.resolve(undefined) - ) + return turnFinders.get(host)?.(lookup) ?? Promise.resolve(undefined) } /** The log and run stores of each durable host. Not public API. */ @@ -556,10 +770,62 @@ const durableStores = new WeakMap() * @internal The log and run stores of `host` when it is durable * (`stores.log`), else `undefined`. */ -export function durableStoresOf(host: HarnessHost) { +export function durableStoresOf

(host: HarnessHost

) { return durableStores.get(host) } +/** The `access` of each host that has one. Not public API. */ +const accesses = new WeakMap() + +/** + * @internal The `access` of `host`, or `undefined` for a host with no + * checks. The HTTP handler reads `identify` from it. + */ +export function accessOf

(host: HarnessHost

) { + // The host stored its own access under its own key, so it has type P. + return accesses.get(host) as HarnessAccess

| undefined +} + +/** The key of what `host.revoke` took: one user on one thread of a tenant. */ +const revokedKey = ( + tenantId: string | undefined, + threadId: string, + principalId: string, +) => `${logIdOf({ tenantId, logId: threadId })}\u0000${principalId}` + +/** + * The owner and members of a fork. `'self'`: the forker alone. `'all'`: the + * forker as the owner, and the other members with their roles, the old + * owner as a `member`. Without a mode or a forker: those of the source. + */ +function forkedMembers( + source: SessionIndexEntry | undefined, + forker: Principal | undefined, + mode: 'self' | 'all' | undefined, +) { + if (mode === undefined || forker === undefined) { + return { ownerId: source?.ownerId, members: source?.members ?? [] } + } + const owner: SessionMember = { userId: forker.id, role: 'owner' } + if (mode === 'self') return { ownerId: forker.id, members: [owner] } + const oldOwner = source?.ownerId + const others = (source?.members ?? []) + .filter((member) => member.userId !== forker.id) + .map( + (member): SessionMember => + member.role === 'owner' ? { ...member, role: 'member' } : member, + ) + // An owner through `ownerId` alone has no entry in `members`. + const isListed = + oldOwner === undefined || + oldOwner === forker.id || + others.some((member) => member.userId === oldOwner) + const old: Array = isListed + ? [] + : [{ userId: oldOwner, role: 'member' }] + return { ownerId: forker.id, members: [owner, ...old, ...others] } +} + /** Throw a clear error for a durable store set that cannot work. */ function checkDurableStores(persistence: HarnessPersistence) { const { stores } = persistence @@ -582,7 +848,10 @@ function checkDurableStores(persistence: HarnessPersistence) { } /** - * Create a host for harness sessions. + * Create a host for harness sessions. With `access`, the principal type of + * the host is what `access.identify` returns, and each call is checked (see + * `HarnessHostOptions.access`). Throws when `access` has no + * `stores.sessions` to keep the owner and members in. * * @example * ```ts @@ -591,7 +860,9 @@ function checkDurableStores(persistence: HarnessPersistence) { * const turn = await session.send('Write a haiku about the sea.') * ``` */ -export function createHarnessHost(options: HarnessHostOptions = {}) { +export function createHarnessHost

( + options: HarnessHostOptions

= {}, +) { if (!options.persistence && !warned) { warned = true console.warn( @@ -601,6 +872,14 @@ export function createHarnessHost(options: HarnessHostOptions = {}) { const persistence: HarnessPersistence = options.persistence ?? memoryPersistence() checkDurableStores(persistence) + checkSessionStore(persistence.stores.sessions) + const { access } = options + if (access && !persistence.stores.sessions) { + throw new Error( + 'Host access needs stores.sessions: each session keeps its owner and members there.', + ) + } + const policy = access && membersOnlyPolicyOf(access.authorize) 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 @@ -641,13 +920,14 @@ export function createHarnessHost(options: HarnessHostOptions = {}) { const statuses = new Map() const setStatus = ( key: string, - threadId: string, + of: SessionIndexKey, status: HostStatusEvent['status'], ) => { if (statuses.get(key)?.status === status) return const event: HostStatusEvent = { type: 'status', - threadId, + threadId: of.threadId, + ...(of.tenantId !== undefined ? { tenantId: of.tenantId } : {}), status, at: Date.now(), } @@ -656,7 +936,7 @@ export function createHarnessHost(options: HarnessHostOptions = {}) { } const index = sessionIndex(sessionStore, emit) const coalesceMs = options.log?.coalesceMs ?? 100 - /** The shared log of each log id with an open session, by log id. */ + /** The shared log of each log with an open session, by its store log id. */ const logs = new Map>() /** The resolved shared logs, and the reduce that folds each one, for `logState`. */ const resolved = new Map< @@ -664,10 +944,15 @@ export function createHarnessHost(options: HarnessHostOptions = {}) { { log: SharedLog; reduce: ReduceOptions | undefined } >() /** - * The shared log of `logId`. The harness of the session that opens it - * first gives the `project` and `reduce` that fold it. + * The shared log with the store log id `logId`, whose records with no + * `thread` field belong to `rootThread`. The harness of the session that + * opens it first gives the `project` and `reduce` that fold it. */ - const sharedLog = (store: LogStore, logId: string, harness: AnyHarness) => { + const sharedLog = ( + store: LogStore, + { logId, rootThread }: { logId: string; rootThread: string }, + harness: AnyHarness, + ) => { const open = logs.get(logId) if (open) return open const project = harness.durability?.project @@ -675,6 +960,7 @@ export function createHarnessHost(options: HarnessHostOptions = {}) { const loaded = loadLogState({ store, logId, + rootThread, ...(metadata ? { metadata } : {}), ...(project ? { project } : {}), ...(reduce ? { reduce } : {}), @@ -682,6 +968,7 @@ export function createHarnessHost(options: HarnessHostOptions = {}) { const log = new SharedLog({ store, logId, + rootThread, state, coalesceMs, ...(project ? { project } : {}), @@ -711,24 +998,42 @@ export function createHarnessHost(options: HarnessHostOptions = {}) { const closing = new Map>() const hostId = `host-${Date.now().toString(36)}-${Math.random().toString(36).slice(2, 9)}` + /** Throw when the index entry of the thread names a parent thread. */ + const refuseChild = async (key: SessionIndexKey) => { + const parent = (await index.get(key))?.parentThreadId + if (parent === undefined) return + throw new Error( + `${JSON.stringify(key.threadId)} is a child thread of ${JSON.stringify(parent)}. Open the root thread and read the child through it.`, + ) + } + /** - * `host.open`. It gives the session class, so the host can call the - * members that the `HarnessSession` type does not have. + * `host.open`, in the tenant `tenantId`. It gives the session class, so + * the host can call the members that the `HarnessSession` type does not + * have. An internal open can pass a tenant with no principal, so it does + * not claim an entry with no owner. */ const openSession = ( harness: THarness, { threadId, principal, logId: openLogId, promptCache }: OpenSessionOptions, + tenantId: string | undefined, ) => { - const key = `${harness.name}\u0000${threadId}` + const key = sessionKey(harness.name, tenantId, threadId) + const of = { tenantId, 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 + // The one place where the log id of the app becomes the id in the + // store: the tenant is part of it. + const storeLogId = logIdOf({ tenantId, logId }) const created = new SessionImpl({ harness, threadId, logId, + storeLogId, + ...(tenantId !== undefined ? { tenantId } : {}), persistence, inbox, credentials, @@ -743,11 +1048,12 @@ export function createHarnessHost(options: HarnessHostOptions = {}) { ? { project: harness.durability.project } : {}), open: async (sessionThread, onFailure) => { - let log = await sharedLog(logStore, logId, harness) + const ids = { logId: storeLogId, rootThread: logId } + let log = await sharedLog(logStore, ids, harness) // The log closed while this open waited, and the host // dropped it. while (!log.isOpen) { - log = await sharedLog(logStore, logId, harness) + log = await sharedLog(logStore, ids, harness) } return log.view(sessionThread, onFailure) }, @@ -757,9 +1063,10 @@ export function createHarnessHost(options: HarnessHostOptions = {}) { ...(options.ownership ? { lease: options.ownership } : {}), ...(principal ? { principal } : {}), ...(promptCache ? { promptCache } : {}), + ...(access ? { access: { envKeys: access.envKeys === true } } : {}), onEvent: (event) => { const status = statusAfter(created, event) - if (status) setStatus(key, threadId, status) + if (status) setStatus(key, of, status) }, onClose: () => { // A sweep may have put a new session under this key already. It @@ -772,12 +1079,13 @@ export function createHarnessHost(options: HarnessHostOptions = {}) { // the thread first. const closed = closing.get(key) ?? Promise.resolve() session = closed + .then(() => refuseChild(of)) .then(() => created.open()) .then( () => { // A recovered turn can run already. Else the new session is // idle. - if (!statuses.has(key)) setStatus(key, threadId, 'idle') + if (!statuses.has(key)) setStatus(key, of, 'idle') return created }, (error: unknown) => { @@ -791,61 +1099,321 @@ export function createHarnessHost(options: HarnessHostOptions = {}) { return session as Promise> } - const host: HarnessHost = { - open: openSession, - isOpen(harness, { threadId }) { - return sessions.has(`${harness.name}\u0000${threadId}`) - }, - async fork(harness, { threadId, newThreadId, at, principal, logId }) { - if (newThreadId === threadId) { - throw new Error('A fork needs a new thread: newThreadId is threadId.') - } - // A source that only this fork opens closes again at the end. - const wasOpen = host.isOpen(harness, { threadId }) - const source = await openSession(harness, { - threadId, - ...(logId ? { logId } : {}), + /** + * How often `host.revoke` took each user's access to each thread. A handle + * keeps the count of its open, and refuses each call once it changed. + */ + // ponytail: one count per revoked user and thread, for the life of the + // host. Drop the old ones if a host revokes millions. + const revocations = new Map() + + /** Count a revoke, and tell the readers of `host.events()`. */ + const revokeUser = ({ + tenantId, + threadId, + principalId, + }: { + tenantId: string | undefined + threadId: string + principalId: string + }) => { + const key = revokedKey(tenantId, threadId, principalId) + revocations.set(key, (revocations.get(key) ?? 0) + 1) + emit({ + type: 'revoked', + ...(tenantId !== undefined ? { tenantId } : {}), + threadId, + principalId, + at: Date.now(), + }) + } + + /** With access, the principal of a public call: it is required. */ + const need = (principal: P | undefined, call: string) => { + if (principal === undefined) { + throw new Error(`${call} needs a principal: the host has access.`) + } + return principal + } + + /** The entry of the root thread of `entry`. A child is checked through it. */ + const rootOf = async (entry: SessionIndexEntry) => { + let root = entry + const seen = new Set([root.threadId]) + while ( + root.parentThreadId !== undefined && + !seen.has(root.parentThreadId) + ) { + seen.add(root.parentThreadId) + const parent = await index.get({ + tenantId: root.tenantId, + threadId: root.parentThreadId, }) - try { - 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)}.`, - ) + if (!parent) break + root = parent + } + return root + } + + /** + * Throw `AccessDeniedError` unless `authorize` lets `principal` do + * `action` on the thread of `entry`. Without access, it allows all. The + * entries are read in the tenant of the principal, so a session of + * another tenant is never found. + */ + const allow = async ({ + principal, + action, + threadId, + entry, + run, + }: { + principal: P + action: AccessAction + threadId: string + entry: SessionIndexEntry | undefined + run?: AccessRun + }) => { + if (!access) return + const { tenantId } = principal + // A thread with no entry has no owner and no members. + const session: AccessSession = entry + ? await rootOf(entry) + : { + ...(tenantId !== undefined ? { tenantId } : {}), + threadId, + members: [], } - const isUsed = logStore - ? (await logStore.read({ logId: 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 isAllowed = await access.authorize({ + principal, + action, + session, + ...(run ? { run } : {}), + }) + if (!isAllowed) throw new AccessDeniedError({ action, threadId }) + } + + /** + * The entry of a thread for a public call, after the check of `action`. + * With access, the principal is required, and a thread with no entry + * needs no check: the call finds nothing. + */ + const checkEntry = async ( + call: string, + action: AccessAction, + { threadId, principal }: { threadId: string; principal?: P | undefined }, + ) => { + const user = access ? need(principal, call) : undefined + const entry = await index.get({ tenantId: principal?.tenantId, threadId }) + if (user && entry) await allow({ principal: user, action, threadId, entry }) + return entry + } + + /** With access, the source of a fork needs an entry and the `fork` action. */ + const checkFork = async ( + call: string, + target: { threadId: string; principal?: P | undefined }, + ) => { + const source = await checkEntry(call, 'fork', target) + // No entry is not found: the same answer as a refusal. + if (access && !source) { + throw new AccessDeniedError({ action: 'fork', threadId: target.threadId }) + } + return source + } + + /** `share` and `unshare` change the members list of `membersOnly()`. */ + const checkShare = async ( + call: string, + target: { threadId: string; principal?: P | undefined }, + ) => { + if (!policy) { + throw new Error( + `${call} needs membersOnly() as the authorize of the host access. With your own access tables, change them and call host.revoke.`, + ) + } + await checkEntry(call, 'share', target) + return need(target.principal, call) + } + + /** The checks of a handle of `principal` on `threadId`. */ + const guardOf = (principal: P, threadId: string): HandleGuard => { + const key = revokedKey(principal.tenantId, threadId, principal.id) + const count = revocations.get(key) ?? 0 + return { + isRevoked: () => (revocations.get(key) ?? 0) !== count, + authorize: async (action, run) => + allow({ + principal, + action, + threadId, + entry: await index.get({ tenantId: principal.tenantId, threadId }), + ...(run ? { run } : {}), + }), + } + } + + /** The handle of `session` for `principal`. With access, it checks. */ + const handleOf = ( + session: SessionImpl, + principal: P | undefined, + ) => + sessionHandle({ + session, + principal, + ...(access && principal + ? { guard: guardOf(principal, session.threadId) } + : {}), + }) + + /** The open sessions of a thread of a tenant on this host, of any harness. */ + const openSessionsOf = async ( + tenantId: string | undefined, + threadId: string, + ) => { + const suffix = `\u0000${logIdOf({ tenantId, logId: threadId })}` + const live = await Promise.allSettled( + [...sessions].flatMap(([key, session]) => + key.endsWith(suffix) ? [session] : [], + ), + ) + return live.flatMap((entry) => + entry.status === 'fulfilled' ? [entry.value] : [], + ) + } + + /** `host.fork` after its checks: the session class of the new thread. */ + const forkInto = async ( + harness: THarness, + { threadId, newThreadId, at, principal, logId }: ForkSessionOptions

, + ) => { + if (newThreadId === threadId) { + throw new Error('A fork needs a new thread: newThreadId is threadId.') + } + const tenantId = principal?.tenantId + // A source that only this fork opens closes again at the end. + const wasOpen = sessions.has(sessionKey(harness.name, tenantId, threadId)) + // In the tenant of the fork, with no principal: the fork does not + // become the owner of a source with no owner. + const source = await openSession( + harness, + { threadId, ...(logId ? { logId } : {}) }, + tenantId, + ) + try { + 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)}.`, + ) + } + // The new thread opens with no `logId`, so its log is its own. Its + // transcript key in `stores.messages` is the same id. + const newLogId = logIdOf({ tenantId, logId: newThreadId }) + // With access, a thread with an entry belongs to someone already. + const isTaken = + access !== undefined && + (await index.get({ tenantId, threadId: newThreadId })) !== undefined + const isUsed = + isTaken || + (logStore + ? (await logStore.read({ logId: newLogId, limit: 1 })).length > 0 + : ((await persistence.stores.messages?.loadThread(newLogId)) + ?.length ?? 0) > 0) + if (isUsed) { + throw new Error( + `Thread ${JSON.stringify(newThreadId)} is not empty. Fork into a new thread.`, + ) + } + const fork = await openSession( + harness, + { threadId: newThreadId, ...(principal ? { principal } : {}) }, + tenantId, + ) + await fork.adoptFork(source, transcript.slice(0, end)) + return fork + } finally { + // A failed close must not hide the fork's own error. + if (!wasOpen) + await source + .close() + .catch((error: unknown) => + console.error( + '[@tanstack/ai-harness] The source session did not close.', + error, + ), ) + } + } + + const host: HarnessHost

= { + async open(harness, openOptions) { + if (!access) { + const { principal } = openOptions + const session = await openSession( + harness, + openOptions, + principal?.tenantId, + ) + return handleOf(session, principal) + } + const principal = need(openOptions.principal, 'host.open') + const { threadId } = openOptions + const key = { tenantId: principal.tenantId, threadId } + const before = await index.get(key) + // The open refuses a child thread, and says why. + const isRoot = before !== undefined && before.parentThreadId === undefined + if (isRoot) { + await allow({ principal, action: 'read', threadId, entry: before }) + } + const session = await openSession( + harness, + openOptions, + principal.tenantId, + ) + if (!before) { + // Two users can create one thread at once. Only the user whose open + // wrote the entry needs no check. + const created = await index.get(key) + if (created?.ownerId !== principal.id) { + await allow({ principal, action: 'read', threadId, entry: created }) } - const fork = await openSession(harness, { - threadId: newThreadId, - ...(principal ? { principal } : {}), - }) - await fork.adoptFork(source, transcript.slice(0, end)) - return fork - } finally { - // A failed close must not hide the fork's own error. - if (!wasOpen) - await source - .close() - .catch((error: unknown) => - console.error( - '[@tanstack/ai-harness] The source session did not close.', - error, - ), - ) } + return handleOf(session, principal) + }, + isOpen(harness, { threadId, principal }) { + if (access) need(principal, 'host.isOpen') + return sessions.has( + sessionKey(harness.name, principal?.tenantId, threadId), + ) + }, + async fork(harness, forkOptions) { + await checkFork('host.fork', forkOptions) + return handleOf( + await forkInto(harness, forkOptions), + forkOptions.principal, + ) + }, + revoke: ({ threadId, principal }) => + revokeUser({ + tenantId: principal.tenantId, + threadId, + principalId: principal.id, + }), + hasAccess: access !== undefined, + async identify({ request }) { + if (!access) { + throw new Error( + 'host.identify needs host access ({ identify, authorize }).', + ) + } + return await access.identify({ request }) }, async resumePending({ harnesses, close = 'whenIdle', limit = 100 }) { const claims = persistence.stores.workClaims @@ -868,30 +1436,34 @@ export function createHarnessHost(options: HarnessHostOptions = {}) { expired = expired.slice(0, limit) const opened: Array<{ threadId: string; harness: string }> = [] const failures: Array = [] - for (const { threadId, harness: name } of expired) { + for (const { threadId: claimed, 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, + threadId: claimed, harness: name, ownerId: hostId, until: Date.now() + ttlMs, }) if (!isMine) continue - const key = `${name}\u0000${threadId}` + // A claim is on the store key of the thread: its tenant and its id. + const { tenantId, logId: threadId } = parseLogId(claimed) + const key = sessionKey(name, tenantId, threadId) // Only a session that this sweep opens closes when idle. A session // that the app has open stays open. const isOwn = close === 'whenIdle' && !sessions.has(key) if (isOwn) sweptOnly.add(key) let session: SessionImpl try { - session = await openSession(harness, { threadId }) + session = await openSession(harness, { threadId }, tenantId) } catch (error) { // One thread that cannot open must not stop the sweep. The claim // goes back, so the next sweep tries again. sweptOnly.delete(key) - await claims.release({ threadId, ownerId: hostId }).catch(() => {}) + await claims + .release({ threadId: claimed, ownerId: hostId }) + .catch(() => {}) failures.push(error) continue } @@ -962,8 +1534,11 @@ export function createHarnessHost(options: HarnessHostOptions = {}) { .map((entry) => entry.value.close(closeOptions)), ) }, - logState: ({ harness, logId }) => { - const folded = resolved.get(logId) + logState: ({ harness, logId, principal }) => { + if (access) need(principal, 'host.logState') + const folded = resolved.get( + logIdOf({ tenantId: principal?.tenantId, logId }), + ) // The reduce of another harness folded this log: its state has // another type. if (folded?.reduce !== harness.durability?.reduce) return undefined @@ -974,40 +1549,107 @@ export function createHarnessHost(options: HarnessHostOptions = {}) { | undefined }, sessions: { - list: async (listOptions) => - (await sessionStore?.list({ + async list({ principal, ...filters } = {}) { + if (!sessionStore) return { entries: [] } + // Last, so a filter never names another tenant. + const tenantId = principal?.tenantId + if (!access) { + return sessionStore.list({ + parentThreadId: null, + ...filters, + tenantId, + }) + } + const user = need(principal, 'host.sessions.list') + const parent = filters.parentThreadId + // The children of a thread are read through their root: one check, + // and no member filter, as a child has no members. + if (typeof parent === 'string') { + await checkEntry('host.sessions.list', 'list', { + threadId: parent, + principal: user, + }) + return sessionStore.list({ ...filters, tenantId }) + } + const page = await sessionStore.list({ parentThreadId: null, - ...listOptions, - })) ?? { - entries: [], - }, - get: (threadId) => index.get(threadId), - rename: ({ threadId, title }) => - index.update( - threadId, + ...filters, + tenantId, + // The store keeps only the sessions of the member, so a page of + // `membersOnly()` stays full. + ...(policy ? { memberId: user.id } : {}), + }) + const allowed = await Promise.all( + page.entries.map(async (session) => + access.authorize({ principal: user, action: 'list', session }), + ), + ) + return { + ...page, + entries: page.entries.filter((_entry, at) => allowed[at]), + } + }, + get: (target) => checkEntry('host.sessions.get', 'read', target), + async rename({ threadId, title, principal }) { + const call = 'host.sessions.rename' + await checkEntry(call, 'configure', { threadId, principal }) + return index.update( + { tenantId: principal?.tenantId, threadId }, (entry) => entry && { ...entry, title, updatedAt: Date.now() }, - ), - update: async ({ harness, threadId, update }) => + ) + }, + async update({ harness, threadId, update, principal }) { + const call = 'host.sessions.update' + await checkEntry(call, 'configure', { threadId, principal }) // 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, + return (await index.update( + { tenantId: principal?.tenantId, threadId }, + (entry) => { + checkHarness(entry, harness) + return changedEntry( + entry, + update as SessionEntryUpdate>, + harness.metadataSchema, + ) + }, + )) as SessionIndexEntry> | undefined + }, + async delete({ threadId, principal }) { + const tenantId = principal?.tenantId + const entry = await checkEntry('host.sessions.delete', 'delete', { + threadId, + principal, + }) + // The transcript goes with the entry, so the next opener of this + // thread id cannot read it. ponytail: the log or the messages only; + // PR 4 adds a purge of blobs, metadata, and children. + if (entry?.harness !== undefined) { + const open = sessions.get( + sessionKey(entry.harness, tenantId, threadId), ) - })) as SessionIndexEntry> | undefined, - delete: (threadId) => index.remove(threadId), - async fork({ harness, threadId, at }) { + if (open) await (await open).close() + } + const storeId = logIdOf({ tenantId, logId: threadId }) + const { log, messages } = persistence.stores + if (log) await log.deleteLog({ logId: storeId }) + else await messages?.saveThread(storeId, []) + await index.remove({ tenantId, threadId }) + }, + async fork({ harness, threadId, at, members, principal }) { + const tenantId = principal?.tenantId // 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) + const source = await checkFork('host.sessions.fork', { + threadId, + principal, + }) 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 }) + // With no principal, as `host.fork` opens the source. + const opened = await openSession(harness, { threadId }, tenantId) const transcript = await opened.transcript() const position = transcript.findIndex(({ id }) => id === at.before) if (position === -1) { @@ -1023,32 +1665,92 @@ export function createHarnessHost(options: HarnessHostOptions = {}) { last = previous } } - const forked = await host.fork(harness, { + const forked = await forkInto(harness, { threadId, newThreadId: `thread-${crypto.randomUUID()}`, at: last, - ...(source?.principal ? { principal: source.principal } : {}), + ...(principal ? { principal } : {}), }) + const mode = + members ?? (access ? (policy?.fork.members ?? 'self') : undefined) + const { ownerId, members: people } = forkedMembers( + source, + principal, + mode, + ) const now = Date.now() // A sibling of the old thread, not a child: no parentThreadId. const entry: SessionIndexEntry = { + ...(tenantId !== undefined ? { tenantId } : {}), threadId: forked.threadId, + ...(ownerId !== undefined ? { ownerId } : {}), + members: people, createdAt: now, updatedAt: now, harness: harness.name, ...(source?.title ? { title: `${source.title} (fork)` } : {}), - ...(source?.principal ? { principal: source.principal } : {}), } - // The open of the fork wrote an entry already: keep what it holds. return ( - (await index.update(entry.threadId, (current) => ({ - ...entry, - ...current, - ...(entry.title ? { title: entry.title } : {}), - updatedAt: now, - }))) ?? entry + (await index.update( + { tenantId, threadId: entry.threadId }, + (current) => { + // The open of the fork wrote an entry already: keep what it + // holds, but not its owner. The owner and the members come + // from `members`. + const { ownerId: _opener, ...opened } = current ?? entry + return { + ...entry, + ...opened, + ...(entry.title ? { title: entry.title } : {}), + members: entry.members, + updatedAt: now, + } + }, + )) ?? entry ) }, + async share({ threadId, userId, role, principal }) { + const call = 'host.sessions.share' + const user = await checkShare(call, { threadId, principal }) + return index.update( + { tenantId: user.tenantId, threadId }, + (entry) => + entry && { + ...entry, + members: [ + ...entry.members.filter((member) => member.userId !== userId), + { userId, role }, + ], + updatedAt: Date.now(), + }, + ) + }, + async unshare({ threadId, userId, principal }) { + const call = 'host.sessions.unshare' + const user = await checkShare(call, { threadId, principal }) + const { tenantId } = user + const changed = await index.update({ tenantId, threadId }, (entry) => { + if (!entry) return undefined + if (entry.ownerId === userId) { + throw new Error( + `${JSON.stringify(userId)} owns thread ${JSON.stringify(threadId)}: the owner cannot be removed.`, + ) + } + return { + ...entry, + members: entry.members.filter((member) => member.userId !== userId), + updatedAt: Date.now(), + } + }) + // After the write, so no new work of the user gets through. The log + // keeps who cancelled it. + const open = await openSessionsOf(tenantId, threadId) + for (const session of open) { + await session.cancelWorkOf({ userId, principal: user }) + } + revokeUser({ tenantId, threadId, principalId: userId }) + return changed + }, }, async *events({ signal } = {}) { // The current statuses first. A change while the reader handles an @@ -1084,9 +1786,13 @@ export function createHarnessHost(options: HarnessHostOptions = {}) { const { runs } = persistence.stores durableStores.set(host, { log: logStore, ...(runs ? { runs } : {}) }) } - turnFinders.set(host, async (name, operationId) => { + if (access) accesses.set(host, access) + turnFinders.set(host, async ({ harness, tenantId, operationId }) => { + // The key of a session with an empty thread id is the start of the keys + // of the harness and tenant: the encoded tenant has no `:`. + const prefix = sessionKey(harness.name, tenantId, '') for (const [key, opening] of sessions) { - if (!key.startsWith(`${name}\u0000`)) continue + if (!key.startsWith(prefix)) continue const session = await opening.catch(() => undefined) if (session?.operation(operationId)?.kind === 'chat') return session } diff --git a/packages/ai-harness/src/http.ts b/packages/ai-harness/src/http.ts index ad76c53eae..5c87488d17 100644 --- a/packages/ai-harness/src/http.ts +++ b/packages/ai-harness/src/http.ts @@ -9,8 +9,10 @@ import { toServerSentEventsStream, } from '@tanstack/ai' import { parseRangeHeader, resolveBlobRange } from '@tanstack/ai-persistence' -import { durableStoresOf, sessionOfTurn } from './host' +import { AccessDeniedError } from './access' +import { accessOf, durableStoresOf, sessionOfTurn } from './host' import { MediaError } from './media' +import { tenantStores } from './tenant-stores' import { HARNESS_PROTOCOL_VERSION, applyInput, @@ -18,6 +20,7 @@ import { describeForClient, parseControlFrame, parseHarnessInput, + turnReceipt, } from './protocol' import { base64url, isRecord } from './utils' import { SessionMetadataError } from './session-entry' @@ -28,7 +31,7 @@ import type { UIMessage, WebSocketLike, } from '@tanstack/ai' -import type { SessionIndexEntry } from '@tanstack/ai-persistence' +import type { SessionIndexEntry, SessionMember } from '@tanstack/ai-persistence' import type { AnyHarness } from './define' import type { HarnessHost, HostEvent } from './host' import type { HostFrame } from './protocol' @@ -36,30 +39,28 @@ import type { HarnessSession } from './session/index' import type { Principal, SessionEvent, UserInput } from './types' /** - * Decide who sends a request. Return `null` to refuse it with 401. Every - * endpoint calls it, except a signed media URL: the handler signs one only - * for a principal that `authorize` and `canAccess` let in. + * Who a handler or a socket serves when its host has no `access`. `true`: + * each request runs with no principal, in the default tenant (one user). + * `{ principal }`: each request runs as that principal, for example the one + * local user of a CLI. */ -export type Authorize = ( - request: Request, -) => Principal | null | Promise - -/** May this principal use this thread? */ -export type CanAccess = ( - principal: Principal, - threadId: string, -) => boolean | Promise +export type OpenAccess

= + | true + | { principal: P } -export interface HarnessHandlerOptions { - host: HarnessHost +export interface HarnessHandlerOptions

{ + /** + * The host. With `access`, `access.identify` reads the user of each + * request, and the host checks each action. + */ + host: HarnessHost

harness: AnyHarness - authorize: Authorize /** - * May this principal use this thread? Required: use it to keep users out - * of each other's threads. Pass `() => true` to let every principal use - * every thread. + * For a host without `access` only: who the requests run as. Without + * host access and without this option, the handler throws at the start, + * so a server never serves every thread to everyone by mistake. */ - canAccess: CanAccess + openAccess?: OpenAccess

/** * How the handler signs media URLs. Required on a durable host * (`stores.log`), where another instance can serve the next request. @@ -102,20 +103,29 @@ const json = (body: unknown, status = 200) => headers: { 'Content-Type': 'application/json' }, }) +const forbidden = () => json({ error: 'forbidden' }, 403) + /** * An SSE response: one event per frame, with `data` as JSON and `id` when - * the frame has one. When the client leaves, `signal` aborts. + * the frame has one. `watch` gets the reader: aborting it ends the stream. + * When the client leaves, the reader aborts. */ function sseResponse( frames: ( signal: AbortSignal, ) => AsyncIterable<{ id?: string; data: unknown }>, + watch?: (reader: AbortController) => void, ) { // Not `request.signal`: some servers abort it once the body is read. const reader = new AbortController() + let isCancelled = false const body = new ReadableStream({ - cancel: () => reader.abort(), + cancel: () => { + isCancelled = true + reader.abort() + }, async start(controller) { + watch?.(reader) try { for await (const { id, data } of frames(reader.signal)) { const idLine = id === undefined ? '' : `id: ${id}\n` @@ -124,7 +134,9 @@ function sseResponse( ) } } finally { - if (!reader.signal.aborted) controller.close() + if (!isCancelled) controller.close() + // The watch ends with the stream. + reader.abort() } }, }) @@ -134,19 +146,27 @@ function sseResponse( /** * Session events as an AG-UI SSE response. With `withIds`, each event has its * cursor as `id`, so a client that loses the connection resumes with - * `Last-Event-ID`. When the client leaves, the read stops. + * `Last-Event-ID`. `watch` gets the reader, as in `sseResponse`. When the + * client leaves, the read stops. */ function cursorResponse( read: (signal: AbortSignal) => AsyncIterable, withIds: boolean, + watch?: (reader: AbortController) => void, ) { // Not `request.signal`: some servers abort it once the body is read. const reader = new AbortController() + watch?.(reader) const cursors: Array = [] async function* chunks() { - for await (const entry of read(reader.signal)) { - cursors.push(entry.cursor) - yield entry.event + try { + for await (const entry of read(reader.signal)) { + cursors.push(entry.cursor) + yield entry.event + } + } finally { + // The watch ends with the stream. + reader.abort() } } return new Response( @@ -159,14 +179,20 @@ function cursorResponse( ) } -/** A handler error as JSON: a `MediaError` keeps its status, the rest is 400. */ -const failed = (error: unknown) => - error instanceof MediaError - ? json({ error: error.message }, error.status) - : json( - { error: error instanceof Error ? error.message : String(error) }, - 400, - ) +/** + * A handler error as JSON: a refused action is 403, a `MediaError` keeps its + * status, the rest is 400. + */ +const failed = (error: unknown) => { + if (error instanceof AccessDeniedError) return forbidden() + if (error instanceof MediaError) { + return json({ error: error.message }, error.status) + } + return json( + { error: error instanceof Error ? error.message : String(error) }, + 400, + ) +} /** * The content of the last user message of an AG-UI request, as TanStack @@ -185,10 +211,14 @@ function lastUserInput(messages: Array) { return input } +const isRole = (value: unknown): value is SessionMember['role'] => + value === 'owner' || value === 'member' || value === 'viewer' + /** A `POST sessions` body, or `undefined` when it is not a valid one. */ function parseSessionOp(body: unknown) { if (!isRecord(body)) return undefined const { op, threadId, title, before, through, pinned, metadata } = body + const { userId, role } = body if (typeof threadId !== 'string') return undefined switch (op) { case 'rename': @@ -221,24 +251,19 @@ function parseSessionOp(body: unknown) { if (typeof through === 'string' && before === undefined) return { op: 'fork' as const, threadId, at: { through } } return undefined + case 'share': + return typeof userId === 'string' && isRole(role) + ? { op: 'share' as const, threadId, userId, role } + : undefined + case 'unshare': + return typeof userId === 'string' + ? { op: 'unshare' as const, threadId, userId } + : undefined default: return undefined } } -/** - * The entry belongs to `principal`: the same id, and the same tenant when - * the principal has one. The same rule as the `principal` filter of the - * index `list`, so a user changes only the sessions the user can list. - */ -function isOwnedBy(entry: SessionIndexEntry, principal: Principal) { - const owner = entry.principal - if (!owner || owner.id !== principal.id) return false - return ( - principal.tenantId === undefined || owner.tenantId === principal.tenantId - ) -} - /** Base64url to bytes, or `undefined` for a value that is not base64url. */ function fromBase64url(value: string) { try { @@ -285,6 +310,71 @@ async function mediaResponse( }) } +/** + * Who sends each request of a handler or a socket: the host's + * `access.identify`, else the `openAccess` principal. Throws for a host + * without access and without `openAccess`, and for a host with access and + * `openAccess`. Resolves to `null` for a request that `identify` refuses. + */ +function requestUsers

( + call: string, + host: HarnessHost

, + openAccess: OpenAccess

| undefined, +) { + const access = accessOf(host) + if (access && openAccess !== undefined) { + throw new Error( + `${call}: the host has access, so access.identify reads the user of each request. Remove openAccess.`, + ) + } + if (!access && openAccess === undefined) { + throw new Error( + `${call} needs host access ({ identify, authorize }) or openAccess: true`, + ) + } + const fixed = openAccess === true ? undefined : openAccess?.principal + return async (request: Request) => { + if (!access) return { principal: fixed } + const principal = await access.identify({ request }) + return principal === null ? null : { principal } + } +} + +/** `event` is the revoke of `principal` on the thread `threadId`. */ +const isRevokeOf = ( + event: HostEvent, + principal: Principal | undefined, + threadId: string, +) => + event.type === 'revoked' && + principal !== undefined && + event.principalId === principal.id && + event.tenantId === principal.tenantId && + event.threadId === threadId + +/** + * Call `onRevoke` when `host.revoke` (or `unshare`) takes the access of + * `principal` to the thread. The watch ends when `signal` aborts. + */ +function watchRevoke

( + host: HarnessHost

, + { + principal, + threadId, + signal, + }: { principal: P | undefined; threadId: string; signal: AbortSignal }, + onRevoke: () => void, +) { + if (principal === undefined) return + void (async () => { + for await (const event of host.events({ signal })) { + if (!isRevokeOf(event, principal, threadId)) continue + onRevoke() + return + } + })() +} + /** * A fetch handler for a harness. Mount it on a route that ends with any of * these paths: @@ -319,37 +409,53 @@ async function mediaResponse( * 413 when it is over `media.maxBytes`, 415 for a type the harness does not * take. * - `GET .../media-url?threadId=&id=`: `{ path, expiresAt }`, a signed path - * to the file, relative to the handler. It works for 1 hour. + * to the file, relative to the handler. It works for 1 hour, and only for + * the user that it was signed for: each request checks that user again. * - `GET .../media?threadId=&id=[&exp=&sig=]`: the bytes, with `Range` - * support. A signed request needs no `authorize`, and a bad or expired + * support. A signed request needs no login, and a bad or expired * signature is 403. * - `GET .../sessions?limit=&cursor=&parentThreadId=`: one page of the - * session index, newest first. Only the sessions of this principal that - * `canAccess` lets in. Default: top-level sessions only. - * - `POST .../sessions`: change a session of this principal. The body is + * session index, newest first: the sessions that `host.sessions.list` + * gives this principal. Default: top-level sessions only. + * - `POST .../sessions`: change a session. The body is * `{ op: 'rename', threadId, title }`, `{ op: 'update', threadId, title?, * pinned?, metadata? }` (each field replaces the stored one, and the * harness `metadataSchema` checks `metadata`. Without a schema, a client * cannot write `metadata`: 403 `not_exposed`), `{ op: 'delete', threadId }` - * (it removes the index entry only), or `{ op: 'fork', threadId, before }` - * or `{ op: 'fork', threadId, through }` with a message id. Rename, - * update, and fork answer with the entry, delete with 204. 400 - * (`invalid_metadata`, with `issues`) for metadata the schema refuses, 403 - * when `canAccess` refuses the thread, 404 for a thread this principal - * does not own, and 409 (`other_harness`) for an update or a fork of a - * thread that another harness runs. - * - `GET .../host-events`: the `host.events()` of this principal as SSE: - * status changes and session index changes. First the current status of - * each open session. Only the threads that `canAccess` lets in and that - * this principal owns in the session index, so it needs `stores.sessions`. + * (it removes the index entry only), `{ op: 'fork', threadId, before }` + * or `{ op: 'fork', threadId, through }` with a message id, and with + * `membersOnly()`, `{ op: 'share', threadId, userId, role }` and + * `{ op: 'unshare', threadId, userId }`. Rename, update, fork, share, and + * unshare answer with the entry, delete with 204. 400 + * (`invalid_metadata`, with `issues`) for metadata the schema refuses, 404 + * for a thread with no entry, and 409 (`other_harness`) for an update or a + * fork of a thread that another harness runs. + * - `GET .../host-events`: the `host.events()` that this principal may + * read, as SSE: status changes, session index changes, and the revokes of + * this principal. First the current status of each open session. + * + * With host `access`, `access.identify` reads the user of each request + * (`null` is 401), and the host checks each action: a refused action is 403. + * A read of a thread with no entry is 404, and it does not start a session. + * Only `GET run?threadId=` answers such a thread with an empty hydration, as + * a `ChatClient` reads it before the first message. + * The streams of a user close when `host.revoke` or `unshare` takes their + * access. Each chat input runs as the principal of its request. * - * Each chat input runs as the principal that `authorize` returned for its - * request, not as the one that opened the session first. + * Throws at the start for a host without `access` and without + * `openAccess`, and for a host with both. + * + * @example + * ```ts + * const handler = createHarnessHandler({ host, harness: assistant }) + * ``` */ -export function createHarnessHandler( - options: HarnessHandlerOptions, -): (request: Request) => Promise { - const { host, harness, authorize, canAccess, media } = options +export function createHarnessHandler

( + options: HarnessHandlerOptions

, +) { + const { host, harness, media } = options + const usersOf = requestUsers('createHarnessHandler', host, options.openAccess) + const hasAccess = accessOf(host) !== undefined const durable = durableStoresOf(host) if (media !== undefined && media.secret.length === 0) { throw new Error( @@ -362,11 +468,35 @@ export function createHarnessHandler( ) } - const openFor = async (principal: Principal, threadId: string) => { - if (!(await canAccess(principal, threadId))) return null - return host.open(harness, { threadId, principal }) + const openFor = (principal: P | undefined, threadId: string) => + host.open(harness, { threadId, ...(principal ? { principal } : {}) }) + + /** The run store of a durable host, in the tenant of `principal`. */ + const runsOf = (principal: P | undefined) => + durable?.runs && + tenantStores({ runs: durable.runs }, principal?.tenantId).runs + + /** + * Open a thread for a read. With access, a thread with no entry is not + * found, so a read never starts a session: `undefined`. + */ + const openToRead = async (principal: P | undefined, threadId: string) => { + const isMissing = + hasAccess && + (await host.sessions.get({ + threadId, + ...(principal ? { principal } : {}), + })) === undefined + return isMissing ? undefined : openFor(principal, threadId) } + /** The stream watch that ends a stream of `principal` on its revoke. */ + const closeOnRevoke = + (principal: P | undefined, threadId: string) => (reader: AbortController) => + watchRevoke(host, { principal, threadId, signal: reader.signal }, () => + reader.abort(), + ) + // ponytail: without `media.secret` (only allowed without a log store), a // random key per handler, so signed URLs stop working after a restart. const secret = @@ -379,36 +509,69 @@ export function createHarnessHandler( 'sign', 'verify', ])) - const signed = (threadId: string, id: string, exp: string) => - encoder.encode(`${threadId}\n${id}\n${exp}`) - const sign = async (threadId: string, id: string, exp: string) => { + /** What a signed media URL holds: who may read which file, until when. */ + interface SignedMedia { + tenantId: string + threadId: string + id: string + userId: string + exp: string + } + const signed = (media: SignedMedia) => + encoder.encode( + [ + harness.name, + media.tenantId, + media.threadId, + media.id, + media.userId, + media.exp, + ].join('\n'), + ) + const sign = async (media: SignedMedia) => { const signature = await crypto.subtle.sign( 'HMAC', await mediaKey(), - signed(threadId, id, exp), + signed(media), ) return base64url(new Uint8Array(signature)) } - /** Serve a signed media URL. The signature was made after `canAccess` passed. */ + /** + * Serve a signed media URL. The URL names its user: the handle opens as + * that user, so a user who lost access (unshare, revoke) is refused. + */ const serveSigned = async (request: Request, url: URL) => { const param = (name: string) => url.searchParams.get(name) ?? '' - const threadId = param('threadId') - const id = param('id') - const exp = param('exp') + const payload: SignedMedia = { + tenantId: param('tenantId'), + threadId: param('threadId'), + id: param('id'), + userId: param('userId'), + exp: param('exp'), + } const signature = fromBase64url(param('sig')) const isValid = signature !== undefined && - Number(exp) > Date.now() && + Number(payload.exp) > Date.now() && // `verify` compares in constant time. (await crypto.subtle.verify( 'HMAC', await mediaKey(), signature, - signed(threadId, id, exp), + signed(payload), )) - if (!isValid) return json({ error: 'forbidden' }, 403) - return mediaResponse(request, await host.open(harness, { threadId }), id) + if (!isValid) return forbidden() + const { tenantId, userId } = payload + // A signed URL keeps the id and the tenant of its user, as the log does, + // and `authorize` gets those. The other fields of the principal are not + // in the URL. + const user = userId + ? ({ id: userId, ...(tenantId ? { tenantId } : {}) } as P) + : undefined + const session = await openToRead(user, payload.threadId) + if (!session) return json({ error: 'not found' }, 404) + return mediaResponse(request, session, payload.id) } /** @@ -419,43 +582,50 @@ export function createHarnessHandler( */ const joinTurn = async ( request: Request, - principal: Principal, + principal: P | undefined, runId: string, ) => { // `offset=-1` is the start. A reconnect sends the last cursor it saw. const offset = new URL(request.url).searchParams.get('offset') const resumeFrom = request.headers.get('Last-Event-ID') ?? (offset === '-1' ? null : offset) - const session = await sessionOfTurn(host, harness, runId) - const turn = session?.operation(runId) - if (session && turn) { - if (!(await canAccess(principal, session.threadId))) { - return json({ error: 'forbidden' }, 403) - } + const running = await sessionOfTurn(host, { + harness, + tenantId: principal?.tenantId, + operationId: runId, + }) + if (running?.operation(runId)) { + const { threadId } = running + const session = await openToRead(principal, threadId) + const turn = session?.operation(runId) + if (!session || !turn) return json({ error: 'unknown run' }, 404) const startedCursor = session .snapshot() .activeOperations.find((item) => item.id === runId)?.startedCursor const from = resumeFrom ?? startedCursor ?? '0' - return cursorResponse((signal) => turn.events({ from, signal }), true) + return cursorResponse( + (signal) => turn.events({ from, signal }), + true, + closeOnRevoke(principal, threadId), + ) } // Another instance runs the turn. A durable host finds its thread in the // run store, and its events in the shared log. Not a child run or an // agent run: only a chat turn of a thread. - const record = await durable?.runs?.get(runId) + const record = await runsOf(principal)?.get(runId) const isChatTurn = record != null && record.kind !== 'agent' && record.parentRunId === undefined if (!record || !isChatTurn) return json({ error: 'unknown run' }, 404) const { threadId } = record - if (!(await canAccess(principal, threadId))) { - return json({ error: 'forbidden' }, 403) - } - const entry = await host.sessions.get(threadId) + const user = principal ? { principal } : {} + const entry = await host.sessions.get({ threadId, ...user }) if (entry?.harness !== undefined && entry.harness !== harness.name) { return json({ error: 'unknown run' }, 404) } - const shared = await host.open(harness, { threadId, principal }) + const shared = await openToRead(principal, threadId) + if (!shared) return json({ error: 'unknown run' }, 404) // ponytail: from the start of the log without a cursor. Store the // started position on the run record if long threads make this slow. return cursorResponse( @@ -465,324 +635,373 @@ export function createHarnessHandler( runId, ), true, + closeOnRevoke(principal, threadId), ) } - return async (request) => { + /** + * The `host.events()` that `principal` may read. With access, a thread + * needs the `read` action. The decision of each thread is kept for the + * stream, until a revoke or an index change of that thread. + */ + const watchedEvents = (principal: P | undefined) => { + const user = principal ? { principal } : {} + const decisions = new Map>() + const mayRead = (threadId: string) => { + const known = decisions.get(threadId) + if (known) return known + const decision = host.sessions.get({ threadId, ...user }).then( + (entry) => entry !== undefined, + () => false, + ) + decisions.set(threadId, decision) + return decision + } + return async (event: HostEvent) => { + const tenantId = + event.type === 'session' || event.type === 'session-deleted' + ? event.entry.tenantId + : event.tenantId + if (tenantId !== principal?.tenantId) return false + if (event.type === 'revoked') { + decisions.delete(event.threadId) + // A user hears of their own revokes only. + return event.principalId === principal?.id + } + if (!hasAccess) return true + if (event.type === 'session') decisions.delete(event.threadId) + if (event.type === 'session-deleted') { + decisions.delete(event.threadId) + // The entry is gone: decide on the removed one. + const access = accessOf(host) + return principal !== undefined && access !== undefined + ? access.authorize({ + principal, + action: 'read', + session: event.entry, + }) + : false + } + return mayRead(event.threadId) + } + } + + const handle = async (request: Request, principal: P | undefined) => { const url = new URL(request.url) const route = url.pathname.split('/').at(-1) - const isSignedMedia = - request.method === 'GET' && - route === 'media' && - url.searchParams.has('sig') - // Before `authorize`: a signed URL is its own permission. - if (isSignedMedia) return serveSigned(request, url).catch(failed) - const principal = await authorize(request) - if (!principal) return json({ error: 'unauthorized' }, 401) + const user = principal ? { principal } : {} - try { - if (request.method === 'GET' && route === 'capabilities') { - return json(capabilitiesOf(harness)) - } + if (request.method === 'GET' && route === 'capabilities') { + return json(capabilitiesOf(harness)) + } - if (request.method === 'POST' && route === 'run') { - const params = await chatParamsFromRequestBody(await request.json()) - const message = lastUserInput(params.messages) - const session = await openFor(principal, params.threadId) - if (!session) return json({ error: 'forbidden' }, 403) - // The turn runs as the request's run id, so an AG-UI client matches - // its events and interrupts. The run id is the input id too, so a - // retried request runs once. It runs as this request's principal. - const ids = { inputId: params.runId, runId: params.runId, principal } - const lastEventId = request.headers.get('Last-Event-ID') - if (params.resume && params.resume.length > 0) { - const answers = { resume: params.resume, ...ids } - // Only the session class takes the run id. `host.open` gives one. - const receipt = await (session instanceof SessionImpl - ? session.resolveForRun(answers) - : session.resolve(answers)) - if (receipt.status === 'rejected') { - return json({ error: receipt.reason }, 409) - } - if (!receipt.operationId) return json({ error: 'rejected' }, 409) - const { operationId } = receipt - return cursorResponse( - (signal) => - followOperation( - session.events({ - ...(durable && lastEventId ? { from: lastEventId } : {}), - signal, - }), - operationId, - ), - durable !== undefined, - ) - } - if (message === undefined) - return json({ error: 'no user message' }, 400) - // AG-UI `forwardedProps` are the input's context. - const { forwardedProps } = params - const operation = session.send(message, { - ...ids, - ...(Object.keys(forwardedProps).length > 0 - ? { context: forwardedProps } - : {}), - }) - const receipt = await operation.receipt + if (request.method === 'POST' && route === 'run') { + const params = await chatParamsFromRequestBody(await request.json()) + const message = lastUserInput(params.messages) + const session = await openFor(principal, params.threadId) + // The turn runs as the request's run id, so an AG-UI client matches + // its events and interrupts. The run id is the input id too, so a + // retried request runs once. It runs as this request's principal. + const ids = { inputId: params.runId, runId: params.runId, ...user } + const lastEventId = request.headers.get('Last-Event-ID') + const watch = closeOnRevoke(principal, params.threadId) + if (params.resume && params.resume.length > 0) { + const answers = { resume: params.resume, ...ids } + // Only the session class takes the run id. `host.open` gives one. + const receipt = await (session instanceof SessionImpl + ? session.resolveForRun(answers) + : session.resolve(answers)) if (receipt.status === 'rejected') { - return json({ error: receipt.reason ?? 'rejected' }, 409) + return json({ error: receipt.reason }, 409) } - // A durable host sends each log cursor as the event id. A client cut - // off sends the same request with `Last-Event-ID`, maybe to another - // instance: the run id finds the same turn, and it streams from there. - // Without a log a cursor means nothing on another instance, so no id. - const from = (durable && lastEventId) || '0' + if (!receipt.operationId) return json({ error: 'rejected' }, 409) + const { operationId } = receipt return cursorResponse( - (signal) => operation.events({ from, signal }), + (signal) => + followOperation( + session.events({ + ...(durable && lastEventId ? { from: lastEventId } : {}), + signal, + }), + operationId, + ), durable !== undefined, + watch, ) } + if (message === undefined) return json({ error: 'no user message' }, 400) + // AG-UI `forwardedProps` are the input's context. + const { forwardedProps } = params + const operation = session.send(message, { + ...ids, + ...(Object.keys(forwardedProps).length > 0 + ? { context: forwardedProps } + : {}), + }) + // A turn that the host refused throws, and the answer is 403. + const receipt = await turnReceipt(operation) + if (receipt.status === 'rejected') { + return json({ error: receipt.reason ?? 'rejected' }, 409) + } + // A durable host sends each log cursor as the event id. A client cut + // off sends the same request with `Last-Event-ID`, maybe to another + // instance: the run id finds the same turn, and it streams from there. + // Without a log a cursor means nothing on another instance, so no id. + const from = (durable && lastEventId) || '0' + return cursorResponse( + (signal) => operation.events({ from, signal }), + durable !== undefined, + watch, + ) + } - if (request.method === 'GET' && route === 'run') { - // The `joinRun` of a ChatClient after a reload names the run only. - const runId = resolveResumeRunId(request) - if (runId) return await joinTurn(request, principal, runId) - const threadId = url.searchParams.get('threadId') - if (!threadId) return json({ error: 'threadId is required' }, 400) - const session = await openFor(principal, threadId) - if (!session) return json({ error: 'forbidden' }, 403) - return json(await hydration(session, durable?.runs)) + if (request.method === 'GET' && route === 'run') { + // The `joinRun` of a ChatClient after a reload names the run only. + const runId = resolveResumeRunId(request) + if (runId) return await joinTurn(request, principal, runId) + const threadId = url.searchParams.get('threadId') + if (!threadId) return json({ error: 'threadId is required' }, 400) + const session = await openToRead(principal, threadId) + // A ChatClient hydrates a new thread before its first message, and a + // non-OK answer is an error there. A new thread has nothing to load. + if (!session) { + return json({ messages: [], activeRun: null, interrupts: null }) } + return json(await hydration(session, runsOf(principal))) + } - if (request.method === 'GET' && route === 'events') { - const threadId = url.searchParams.get('threadId') - if (!threadId) return json({ error: 'threadId is required' }, 400) - const session = await openFor(principal, threadId) - if (!session) return json({ error: 'forbidden' }, 403) - const from = - request.headers.get('Last-Event-ID') ?? - url.searchParams.get('from') ?? - undefined - return sseResponse(async function* (signal) { + if (request.method === 'GET' && route === 'events') { + const threadId = url.searchParams.get('threadId') + if (!threadId) return json({ error: 'threadId is required' }, 400) + const session = await openToRead(principal, threadId) + if (!session) return json({ error: 'not found' }, 404) + const from = + request.headers.get('Last-Event-ID') ?? + url.searchParams.get('from') ?? + undefined + return sseResponse( + async function* (signal) { const events = session.events({ ...(from ? { from } : {}), signal }) for await (const entry of events) { const frame: HostFrame = { type: 'harness.event', ...entry } yield { id: entry.cursor, data: frame } } - }) - } + }, + closeOnRevoke(principal, threadId), + ) + } - if (request.method === 'GET' && route === 'host-events') { - // The same rule as `POST sessions`: `canAccess`, and the owner in - // the index. A deleted entry comes with the event. - const mayWatch = async (event: HostEvent) => { - if (!(await canAccess(principal, event.threadId))) return false - const entry = - event.type === 'status' - ? await host.sessions.get(event.threadId) - : event.entry - return entry !== undefined && isOwnedBy(entry, principal) + if (request.method === 'GET' && route === 'host-events') { + const mayWatch = watchedEvents(principal) + return sseResponse(async function* (signal) { + for await (const event of host.events({ signal })) { + if (await mayWatch(event)) yield { data: event } } - return sseResponse(async function* (signal) { - for await (const event of host.events({ signal })) { - if (await mayWatch(event)) yield { data: event } - } - }) - } + }) + } - if (request.method === 'POST' && route === 'control') { - const body: unknown = await request.json() - if ( - typeof body !== 'object' || - body === null || - !('threadId' in body) - ) { - return json({ error: 'threadId is required' }, 400) - } - const threadId = body.threadId - if (typeof threadId !== 'string') - return json({ error: 'threadId is required' }, 400) - const input = parseHarnessInput( - 'input' in body ? body.input : undefined, - ) - const session = await openFor(principal, threadId) - if (!session) return json({ error: 'forbidden' }, 403) - return json(await applyInput({ harness, session, input, principal })) + if (request.method === 'POST' && route === 'control') { + const body: unknown = await request.json() + if (typeof body !== 'object' || body === null || !('threadId' in body)) { + return json({ error: 'threadId is required' }, 400) } + const threadId = body.threadId + if (typeof threadId !== 'string') + return json({ error: 'threadId is required' }, 400) + const input = parseHarnessInput('input' in body ? body.input : undefined) + const session = await openFor(principal, threadId) + return json(await applyInput({ harness, session, input, ...user })) + } - if ( - request.method === 'GET' && - (route === 'transcript' || route === 'describe') - ) { - const threadId = url.searchParams.get('threadId') - if (!threadId) return json({ error: 'threadId is required' }, 400) - const session = await openFor(principal, threadId) - if (!session) return json({ error: 'forbidden' }, 403) - return json( - route === 'transcript' - ? await session.transcript() - : describeForClient(harness, session), - ) - } + if ( + request.method === 'GET' && + (route === 'transcript' || route === 'describe') + ) { + const threadId = url.searchParams.get('threadId') + if (!threadId) return json({ error: 'threadId is required' }, 400) + const session = await openToRead(principal, threadId) + if (!session) return json({ error: 'not found' }, 404) + return json( + route === 'transcript' + ? await session.transcript() + : describeForClient(harness, session), + ) + } - if (request.method === 'GET' && route === 'agent-transcript') { - const threadId = url.searchParams.get('threadId') - const subagentRunId = url.searchParams.get('subagentRunId') - if (!threadId) return json({ error: 'threadId is required' }, 400) - if (!subagentRunId) - return json({ error: 'subagentRunId is required' }, 400) - const session = await openFor(principal, threadId) - if (!session) return json({ error: 'forbidden' }, 403) - return json(await session.agentTranscript(subagentRunId)) - } + if (request.method === 'GET' && route === 'agent-transcript') { + const threadId = url.searchParams.get('threadId') + const subagentRunId = url.searchParams.get('subagentRunId') + if (!threadId) return json({ error: 'threadId is required' }, 400) + if (!subagentRunId) + return json({ error: 'subagentRunId is required' }, 400) + const session = await openToRead(principal, threadId) + if (!session) return json({ error: 'not found' }, 404) + return json(await session.agentTranscript(subagentRunId)) + } - if (request.method === 'GET' && route === 'snapshot') { - const threadId = url.searchParams.get('threadId') - if (!threadId) return json({ error: 'threadId is required' }, 400) - const session = await openFor(principal, threadId) - if (!session) return json({ error: 'forbidden' }, 403) - return json(session.snapshot()) - } + if (request.method === 'GET' && route === 'snapshot') { + const threadId = url.searchParams.get('threadId') + if (!threadId) return json({ error: 'threadId is required' }, 400) + const session = await openToRead(principal, threadId) + if (!session) return json({ error: 'not found' }, 404) + return json(session.snapshot()) + } - if (request.method === 'GET' && route === 'sessions') { - const limit = url.searchParams.get('limit') - const cursor = url.searchParams.get('cursor') - const parentThreadId = url.searchParams.get('parentThreadId') - const isBadLimit = - limit !== null && - !(Number.isInteger(Number(limit)) && Number(limit) > 0) - if (isBadLimit) - return json({ error: 'limit must be a positive integer' }, 400) - const search = url.searchParams.get('search') - // `metadata.=` asks for that exact metadata value. - const metadata = Object.fromEntries( - [...url.searchParams] - .filter(([key]) => key.startsWith('metadata.')) - .map(([key, value]) => [key.slice('metadata.'.length), value]), - ) - // The store keeps only this principal's entries of this harness. - // `canAccess` can still refuse some of them. - const page = await host.sessions.list({ - principal, + if (request.method === 'GET' && route === 'sessions') { + const limit = url.searchParams.get('limit') + const cursor = url.searchParams.get('cursor') + const parentThreadId = url.searchParams.get('parentThreadId') + const isBadLimit = + limit !== null && + !(Number.isInteger(Number(limit)) && Number(limit) > 0) + if (isBadLimit) + return json({ error: 'limit must be a positive integer' }, 400) + const search = url.searchParams.get('search') + // `metadata.=` asks for that exact metadata value. + const metadata = Object.fromEntries( + [...url.searchParams] + .filter(([name]) => name.startsWith('metadata.')) + .map(([name, value]) => [name.slice('metadata.'.length), value]), + ) + // The host lists the entries of this harness in the tenant of this + // principal, and only the ones that it may list. + return json( + await host.sessions.list({ + ...user, harness: harness.name, metadata, ...(limit ? { limit: Number(limit) } : {}), ...(cursor ? { cursor } : {}), ...(parentThreadId ? { parentThreadId } : {}), ...(search ? { search } : {}), - }) - const allowed = await Promise.all( - page.entries.map(async (entry) => - canAccess(principal, entry.threadId), - ), - ) - return json({ - ...page, - entries: page.entries.filter((_entry, index) => allowed[index]), - }) - } + }), + ) + } - if (request.method === 'POST' && route === 'sessions') { - const op = parseSessionOp(await request.json()) - if (!op) return json({ error: 'not a valid sessions op' }, 400) - if (!(await canAccess(principal, op.threadId))) - return json({ error: 'forbidden' }, 403) - // The session of another user answers as a missing one, so its - // thread id tells the caller nothing. - const entry = await host.sessions.get(op.threadId) - if (!entry || !isOwnedBy(entry, principal)) - return json({ error: 'not found' }, 404) - switch (op.op) { - case 'rename': { - const renamed = await host.sessions.rename({ - threadId: op.threadId, - title: op.title, - }) - return renamed ? json(renamed) : json({ error: 'not found' }, 404) - } - case 'update': { - // The host can run more harnesses than this handler serves. - if (entry.harness !== undefined && entry.harness !== harness.name) - return json({ error: 'other_harness' }, 409) - // The schema is the contract for what a client may write. Without - // one, metadata stays a server field. - if (op.change.metadata !== undefined && !harness.metadataSchema) - return json({ error: 'not_exposed' }, 403) - try { - const updated = await host.sessions.update({ + if (request.method === 'POST' && route === 'sessions') { + const op = parseSessionOp(await request.json()) + if (!op) return json({ error: 'not a valid sessions op' }, 400) + const target = { threadId: op.threadId, ...user } + const entry = await host.sessions.get(target) + if (!entry) return json({ error: 'not found' }, 404) + /** A changed entry, or 404 when the entry went away meanwhile. */ + const answer = (changed: SessionIndexEntry | undefined) => + changed ? json(changed) : json({ error: 'not found' }, 404) + // The host can run more harnesses than this handler serves. + const isOtherHarness = + entry.harness !== undefined && entry.harness !== harness.name + switch (op.op) { + case 'rename': + return answer( + await host.sessions.rename({ ...target, title: op.title }), + ) + case 'update': { + if (isOtherHarness) return json({ error: 'other_harness' }, 409) + // The schema is the contract for what a client may write. Without + // one, metadata stays a server field. + if (op.change.metadata !== undefined && !harness.metadataSchema) + return json({ error: 'not_exposed' }, 403) + try { + return answer( + await host.sessions.update({ + ...target, harness, - threadId: op.threadId, update: op.change, - }) - return updated ? json(updated) : json({ error: 'not found' }, 404) - } catch (error) { - if (error instanceof SessionMetadataError) - return json( - { error: 'invalid_metadata', issues: error.issues }, - 400, - ) - throw error - } - } - case 'delete': - await host.sessions.delete(op.threadId) - return new Response(null, { status: 204 }) - case 'fork': - // The host can run more harnesses than this handler serves. - if (entry.harness !== undefined && entry.harness !== harness.name) - return json({ error: 'other_harness' }, 409) - return json( - await host.sessions.fork({ - harness, - threadId: op.threadId, - at: op.at, }), ) + } catch (error) { + if (error instanceof SessionMetadataError) + return json( + { error: 'invalid_metadata', issues: error.issues }, + 400, + ) + throw error + } } + case 'delete': + await host.sessions.delete(target) + return new Response(null, { status: 204 }) + case 'fork': + if (isOtherHarness) return json({ error: 'other_harness' }, 409) + return json( + await host.sessions.fork({ ...target, harness, at: op.at }), + ) + case 'share': + return answer( + await host.sessions.share({ + ...target, + userId: op.userId, + role: op.role, + }), + ) + case 'unshare': + return answer( + await host.sessions.unshare({ ...target, userId: op.userId }), + ) } + } - if (request.method === 'POST' && route === 'media') { - const threadId = url.searchParams.get('threadId') - const name = url.searchParams.get('name') - if (!threadId || !name) - return json({ error: 'threadId and name are required' }, 400) - const session = await openFor(principal, threadId) - if (!session) return json({ error: 'forbidden' }, 403) - const record = await session.putMedia({ - body: request.body ?? new Uint8Array(), - mimeType: request.headers.get('Content-Type') ?? '', - name, - }) - return json(record) - } + if (request.method === 'POST' && route === 'media') { + const threadId = url.searchParams.get('threadId') + const name = url.searchParams.get('name') + if (!threadId || !name) + return json({ error: 'threadId and name are required' }, 400) + const session = await openFor(principal, threadId) + const record = await session.putMedia({ + body: request.body ?? new Uint8Array(), + mimeType: request.headers.get('Content-Type') ?? '', + name, + }) + return json(record) + } - if ( - request.method === 'GET' && - (route === 'media' || route === 'media-url') - ) { - const threadId = url.searchParams.get('threadId') - const id = url.searchParams.get('id') - if (!threadId || !id) - return json({ error: 'threadId and id are required' }, 400) - const session = await openFor(principal, threadId) - if (!session) return json({ error: 'forbidden' }, 403) - if (route === 'media') return await mediaResponse(request, session, id) - if (!(await session.getMedia(id))) - return json({ error: 'not found' }, 404) - const expiresAt = Date.now() + SIGNED_URL_TTL_MS - const exp = String(expiresAt) - const query = new URLSearchParams({ - threadId, - id, - exp, - sig: await sign(threadId, id, exp), - }) - return json({ path: `media?${query}`, expiresAt }) + if ( + request.method === 'GET' && + (route === 'media' || route === 'media-url') + ) { + const threadId = url.searchParams.get('threadId') + const id = url.searchParams.get('id') + if (!threadId || !id) + return json({ error: 'threadId and id are required' }, 400) + const session = await openToRead(principal, threadId) + if (!session) return json({ error: 'not found' }, 404) + if (route === 'media') return await mediaResponse(request, session, id) + if (!(await session.getMedia(id))) + return json({ error: 'not found' }, 404) + const expiresAt = Date.now() + SIGNED_URL_TTL_MS + const payload: SignedMedia = { + tenantId: principal?.tenantId ?? '', + threadId, + id, + userId: principal?.id ?? '', + exp: String(expiresAt), } + // An empty field stays out of the URL. The served URL reads it as ''. + const fields = Object.entries(payload).filter(([, value]) => value !== '') + const query = new URLSearchParams([ + ...fields, + ['sig', await sign(payload)], + ]) + return json({ path: `media?${query}`, expiresAt }) + } + return json({ error: 'not found' }, 404) + } + + return async (request: Request) => { + const url = new URL(request.url) + const isSignedMedia = + request.method === 'GET' && + url.pathname.split('/').at(-1) === 'media' && + url.searchParams.has('sig') + try { + // Before `identify`: a signed URL is its own login. + if (isSignedMedia) return await serveSigned(request, url) + const user = await usersOf(request) + if (!user) return json({ error: 'unauthorized' }, 401) + return await handle(request, user.principal) } catch (error) { return failed(error) } - return json({ error: 'not found' }, 404) } } @@ -833,28 +1052,53 @@ async function* followOperation( } } -export interface HarnessSocketOptions { - host: HarnessHost +export interface HarnessSocketOptions

{ + host: HarnessHost

harness: AnyHarness socket: WebSocketLike - /** The principal your upgrade handler authorized. */ - principal: Principal - /** May this principal use this thread? Pass `() => true` to allow all. */ - canAccess: CanAccess + /** + * The upgrade request. With host `access`, `access.identify` reads the + * user from it. + */ + request: Request + /** For a host without `access` only. See `HarnessHandlerOptions`. */ + openAccess?: OpenAccess

} +/** The text of an error frame: a refused action is `forbidden`. */ +const errorMessage = (error: unknown) => + error instanceof AccessDeniedError + ? 'forbidden' + : error instanceof Error + ? error.message + : String(error) + /** * Serve the session tier over one WebSocket. The first frame must be - * `harness.subscribe`. Authorize the upgrade request before you call this. + * `harness.subscribe`. The user comes from the upgrade `request`, with the + * same rules as `createHarnessHandler`: a user that `identify` refuses gets + * an `unauthorized` error frame and the socket closes. The host checks each + * input. When `host.revoke` or `unshare` takes the access of the user, the + * socket sends a `forbidden` error frame and closes. + * + * Throws at the start for a host without `access` and without `openAccess`, + * and for a host with both. + * + * @example + * ```ts + * handleHarnessSocket({ host, harness: assistant, socket, request }) + * ``` */ -export function handleHarnessSocket(options: HarnessSocketOptions): void { - const { host, harness, socket, principal, canAccess } = options +export function handleHarnessSocket

( + options: HarnessSocketOptions

, +) { + const { host, harness, socket, request } = options + const usersOf = requestUsers('handleHarnessSocket', host, options.openAccess) const reader = new AbortController() + const who = usersOf(request).catch(() => null) // Set by the first subscribe, before its awaits, so a second subscribe is // refused and an input sent right after the subscribe waits for the open. - let opened: - | Promise> | undefined> - | undefined + let opened: Promise | undefined const send = (frame: HostFrame) => { try { socket.send(JSON.stringify(frame)) @@ -862,26 +1106,37 @@ export function handleHarnessSocket(options: HarnessSocketOptions): void { reader.abort() } } + const refuse = (message: 'unauthorized' | 'forbidden') => { + send({ type: 'harness.error', message }) + socket.close(message === 'unauthorized' ? 4401 : 4403, message) + reader.abort() + } socket.addEventListener('close', () => reader.abort()) socket.addEventListener('error', () => reader.abort()) socket.addEventListener('message', (message) => { void (async () => { try { + const user = await who + if (!user) return refuse('unauthorized') + const { principal } = user const frame = parseControlFrame(String(message.data)) if (frame.type === 'harness.subscribe') { if (opened) throw new Error('Already subscribed.') const { threadId } = frame - opened = (async () => - (await canAccess(principal, threadId)) - ? host.open(harness, { threadId, principal }) - : undefined)() + opened = host + .open(harness, { threadId, ...(principal ? { principal } : {}) }) + .catch((error: unknown) => { + if (!(error instanceof AccessDeniedError)) throw error + return undefined + }) const session = await opened - if (!session) { - send({ type: 'harness.error', message: 'forbidden' }) - socket.close(4403, 'forbidden') - return - } + if (!session) return refuse('forbidden') + watchRevoke( + host, + { principal, threadId, signal: reader.signal }, + () => refuse('forbidden'), + ) send({ type: 'harness.hello', v: HARNESS_PROTOCOL_VERSION, @@ -896,10 +1151,7 @@ export function handleHarnessSocket(options: HarnessSocketOptions): void { send({ type: 'harness.event', ...entry }) })() .catch((error: unknown) => - send({ - type: 'harness.error', - message: error instanceof Error ? error.message : String(error), - }), + send({ type: 'harness.error', message: errorMessage(error) }), ) // The feed ended while the socket is open: close it, so the // client subscribes again. @@ -918,7 +1170,7 @@ export function handleHarnessSocket(options: HarnessSocketOptions): void { harness, session, input: frame.input, - principal, + ...(principal ? { principal } : {}), }) send({ type: 'harness.receipt', @@ -926,10 +1178,7 @@ export function handleHarnessSocket(options: HarnessSocketOptions): void { ...receipt, }) } catch (error) { - send({ - type: 'harness.error', - message: error instanceof Error ? error.message : String(error), - }) + send({ type: 'harness.error', message: errorMessage(error) }) } })() }) diff --git a/packages/ai-harness/src/index.ts b/packages/ai-harness/src/index.ts index b97477297f..8f0ad2795f 100644 --- a/packages/ai-harness/src/index.ts +++ b/packages/ai-harness/src/index.ts @@ -27,7 +27,7 @@ export type { ToolStep, } from './durable-tool' -export { LogVersionError, logMessageStore } from './log' +export { LogVersionError, logIdOf, logMessageStore } from './log' export type { ProjectOptions, ProjectRecord, ReduceOptions } from './log' export type { LeaseOptions } from './resume' export type { @@ -100,6 +100,7 @@ export type { HarnessHostOptions, HarnessPersistence, HostEvent, + HostRevokedEvent, HostSessionDeletedEvent, HostSessionEvent, HostSessions, @@ -107,11 +108,23 @@ export type { OpenSessionOptions, ResumePendingOptions, } from './host' +export { AccessDeniedError, membersOnly, membersOnlyPolicyOf } from './access' +export type { + AccessAction, + AccessRun, + AccessSession, + Authorize, + AuthorizeInput, + HarnessAccess, + MembersOnlyOptions, + MembersOnlyPolicy, +} from './access' export type { SessionIndexEntry, SessionIndexListOptions, SessionIndexPage, SessionIndexStore, + SessionMember, } from '@tanstack/ai-persistence' export type { @@ -152,6 +165,7 @@ export type { Cursor, ForkPoint, HarnessInput, + HarnessRunContext, InputSettlement, LiveEvent, MediaKind, @@ -163,6 +177,7 @@ export type { PendingQuestion, Principal, Receipt, + RunContext, SendOptions, SessionEvent, StreamChunkOf, diff --git a/packages/ai-harness/src/log.ts b/packages/ai-harness/src/log.ts index 99884387f6..830aa909dd 100644 --- a/packages/ai-harness/src/log.ts +++ b/packages/ai-harness/src/log.ts @@ -525,12 +525,52 @@ export function sessionOf(state: SharedLogState, thread: string) { return session } -/** The session of a record: its `thread`, or the session whose thread id is the log id. */ -export const threadOf = (record: LogRecord, logId: string) => - typeof record.thread === 'string' ? record.thread : logId +/** + * The id of a log in the `LogStore`: the tenant and the log id, each + * URL-encoded, with a `:` between them. Without `tenantId`, the log is in the + * default tenant. So one log id in two tenants is two logs. `logId` is the + * `logId` of `host.open`, and its default is the thread id. + * + * The other stores that the harness keys by thread (messages, runs, + * interrupts, metadata, inbox, work claims, leases, and media) use the same + * key, with the thread id as `logId`. + * + * @example + * ```ts + * const entries = await log.read({ logId: logIdOf({ tenantId: 'org-1', logId: threadId }) }) + * ``` + */ +export function logIdOf({ + tenantId, + logId, +}: { + tenantId?: string + logId: string +}) { + return `${encodeURIComponent(tenantId ?? '')}:${encodeURIComponent(logId)}` +} + +/** + * The tenant and the log id of a store id from {@link logIdOf}. The first + * `:` splits them. An empty tenant is the default tenant: no `tenantId`. + */ +export function parseLogId(id: string) { + const at = id.indexOf(':') + const tenantId = decodeURIComponent(id.slice(0, Math.max(at, 0))) + const logId = decodeURIComponent(id.slice(at + 1)) + return tenantId === '' ? { logId } : { tenantId, logId } +} + +/** + * The session of a record: its `thread`, or the root thread of the log. The + * root thread is the thread whose id is the `logId` of `host.open`. + */ +export const threadOf = (record: LogRecord, rootThread: string) => + typeof record.thread === 'string' ? record.thread : rootThread interface FoldOptions { - logId: string + /** The thread of the records that have no `thread` field. */ + rootThread: string project?: ProjectRecord reduce?: ReduceOptions } @@ -556,7 +596,7 @@ export function foldLogEntry( ) } foldEntry( - sessionOf(state, threadOf(record, options.logId)), + sessionOf(state, threadOf(record, options.rootThread)), entry, options.project, ) @@ -700,17 +740,20 @@ const versionsOf = ( }) /** - * Fold the log `logId`. With `metadata`, start from the newest fold - * checkpoint when it is valid for this log and these versions. + * Fold the log `logId`, the id in the store (see {@link logIdOf}). The + * records with no `thread` field belong to `rootThread`. With `metadata`, + * start from the newest fold checkpoint when it is valid for this log and + * these versions. */ export async function loadLogState(options: { store: LogStore logId: string + rootThread: string metadata?: MetadataStore project?: ProjectOptions reduce?: ReduceOptions }) { - const { store, logId, metadata, project, reduce } = options + const { store, logId, rootThread, metadata, project, reduce } = options const saved = metadata ? await metadata.get(CHECKPOINT_NAMESPACE, logId).catch(() => null) : null @@ -727,7 +770,7 @@ export async function loadLogState(options: { if (!atCheckpoint) state = emptySharedLogState(reduce) } const fold: FoldOptions = { - logId, + rootThread, ...(project ? { project: project.record } : {}), ...(reduce ? { reduce } : {}), } @@ -774,6 +817,7 @@ function deltaKey(thread: string, operationId: string, event: StreamChunk) { /** A session's view of its log. It is also the session's {@link EventFeed}. */ export interface LogWriter extends EventFeed { readonly threadId: string + /** The id of the log in the store. See {@link logIdOf}. */ readonly logId: string /** * The fold of this session. In a shared log, `seq` is this session's last @@ -810,7 +854,10 @@ export interface LogWriter extends EventFeed { * rejects. */ export class SharedLog { + /** The id of the log in the store. See {@link logIdOf}. */ readonly logId: string + /** The thread of the records that have no `thread` field. */ + readonly rootThread: string readonly state: SharedLogState private readonly store: LogStore private readonly fold: FoldOptions @@ -852,7 +899,10 @@ export class SharedLog { constructor(options: { store: LogStore + /** The id of the log in the store. See {@link logIdOf}. */ logId: string + /** The thread of the records that have no `thread` field. */ + rootThread: string state: SharedLogState coalesceMs: number project?: ProjectOptions @@ -863,12 +913,13 @@ export class SharedLog { }) { this.store = options.store this.logId = options.logId + this.rootThread = options.rootThread this.state = options.state this.metadata = options.metadata this.coalesceMs = options.coalesceMs this.onIdle = options.onIdle ?? (() => {}) this.fold = { - logId: options.logId, + rootThread: options.rootThread, ...(options.project ? { project: options.project.record } : {}), ...(options.reduce ? { reduce: options.reduce } : {}), } @@ -972,9 +1023,9 @@ export class SharedLog { }) } - /** Stamp a record with its session, unless it names one or the session is the log's own. */ + /** Stamp a record with its session, unless it names one or the session is the root thread. */ private stamp(record: LogRecord, thread: string): LogRecord { - if (typeof record.thread === 'string' || thread === this.logId) { + if (typeof record.thread === 'string' || thread === this.rootThread) { return record } return { ...record, thread } @@ -1221,7 +1272,10 @@ export class SharedLog { foldLogEntry(this.state, entry, this.fold) const event = sessionEventOf(entry) if (!event) return - this.tail.push({ thread: threadOf(entry.record, this.logId), entry: event }) + this.tail.push({ + thread: threadOf(entry.record, this.rootThread), + entry: event, + }) if (this.tail.length > MAX_TAIL) { const dropped = this.tail.shift() if (dropped) this.tailFrom = Number(dropped.entry.cursor) @@ -1271,7 +1325,7 @@ export class SharedLog { if (entry.seq > upTo || signal?.aborted) break after = entry.seq const event = sessionEventOf(entry) - const isOwn = threadOf(entry.record, this.logId) === thread + const isOwn = threadOf(entry.record, this.rootThread) === thread if (event && isOwn && accepts(event)) yield event } if (page.length === 0) after = upTo @@ -1346,20 +1400,27 @@ function sessionEventOf(entry: LogEntry) { return event } +/** What a durable session needs for the message stores of its log. */ +interface SessionLogView { + writer: LogWriter + store: LogStore + project?: ProjectOptions + /** The tenant of the session. The logs of its subagent children are in it. */ + tenantId?: string +} + /** * The `MessageStore` a durable session gives `withPersistence`. Its own * thread goes through the session's writer. Another thread (a subagent - * transcript) reads and appends that thread's log directly. + * transcript) reads and appends that thread's log directly, in the tenant of + * the session. */ -export function sessionMessageStore(options: { - writer: LogWriter - store: LogStore - project?: ProjectOptions -}) { +export function sessionMessageStore(options: SessionLogView) { const { writer } = options const other = logMessageStore({ store: options.store, ...(options.project ? { project: options.project } : {}), + ...(options.tenantId !== undefined ? { tenantId: options.tenantId } : {}), }) const messages: MessageStore = { loadThread: async (threadId: string) => @@ -1384,11 +1445,7 @@ export function sessionMessageStore(options: { * engine's list and gives the model the fold, so the two are the same at * every model call. */ -export function engineMessageStore(options: { - writer: LogWriter - store: LogStore - project?: ProjectOptions -}) { +export function engineMessageStore(options: SessionLogView) { const { writer } = options const shared = sessionMessageStore(options) const fold = () => [...writer.state.messages] @@ -1494,19 +1551,31 @@ export function engineMessageStore(options: { export function logMessageStore(options: { store: LogStore project?: ProjectOptions - /** Read the threads of this log. Default: each thread is its own log. */ + /** + * Read the threads of this log: the `logId` of `host.open`. Default: each + * thread is its own log. + */ logId?: string + /** + * The tenant of the threads: the `principal.tenantId` of `host.open`. + * Default: the default tenant. + */ + tenantId?: string }) { - const { store, project } = options + const { store, project, tenantId } = options // ponytail: each read and save folds the whole log of the thread, with no // checkpoint: O(n) per save, so O(n²) over a long child transcript. Pass // `metadata` to loadLogState and write checkpoints here when a child with // many tool phases gets slow. - const logOf = (threadId: string) => options.logId ?? threadId + /** The root thread of the log of `threadId`. */ + const rootOf = (threadId: string) => options.logId ?? threadId + const logOf = (threadId: string) => + logIdOf({ tenantId, logId: rootOf(threadId) }) const load = (threadId: string) => loadLogState({ store, logId: logOf(threadId), + rootThread: rootOf(threadId), ...(project ? { project } : {}), }) const messages: MessageStore = { @@ -1521,7 +1590,7 @@ export function logMessageStore(options: { sessionOf(state, threadId).messages, list, ).map((record) => - logOf(threadId) === threadId + rootOf(threadId) === threadId ? record : { ...record, thread: threadId }, ) diff --git a/packages/ai-harness/src/media.ts b/packages/ai-harness/src/media.ts index a41090ce59..b3f26f2633 100644 --- a/packages/ai-harness/src/media.ts +++ b/packages/ai-harness/src/media.ts @@ -6,6 +6,7 @@ import { retrieveBlob, withGenerationPersistence, } from '@tanstack/ai-persistence' +import { logIdOf } from './log' import { kindOf, mediaIdOf } from './media-ref' import type { AnyTranscriptionAdapter, @@ -71,17 +72,25 @@ const ALL_KINDS: ReadonlyArray = [ 'document', ] -// Uploads are artifact records with the run id `upload:`. Their -// bytes use the key `artifacts/upload:/`: the thread id -// comes from the client, and a raw `../` would leave the folder of a -// filesystem blob store. Generated media keeps what +// The artifact records of a thread have the store key of the thread as their +// `threadId`: `logIdOf` of the tenant and the thread id. So one thread id in +// two tenants has two sets of media. Uploads have the run id +// `upload:`, and their bytes use the key +// `artifacts/upload:/`. The store key is URL-encoded: the +// thread id comes from the client, and a raw `../` would leave the folder of +// a filesystem blob store. Generated media keeps what // `withGenerationPersistence` gives it: a generation run id // `[:...]:-`. -function uploadRunId(threadId: string) { - return `upload:${threadId}` +function uploadRunId(storeThreadId: string) { + return `upload:${storeThreadId}` } -function toMediaRecord(record: ArtifactRecord, kind: MediaKind) { +/** The media record of `record`, an artifact of the thread `threadId`. */ +function toMediaRecord( + record: ArtifactRecord, + kind: MediaKind, + threadId: string, +) { const isUpload = record.runId === uploadRunId(record.threadId) const segments = isUpload ? [] : record.runId.split(':') const [runId] = segments @@ -90,7 +99,7 @@ function toMediaRecord(record: ArtifactRecord, kind: MediaKind) { ) const base = { id: record.artifactId, - threadId: record.threadId, + threadId, kind, mimeType: record.mimeType, name: record.name, @@ -145,8 +154,8 @@ export function toBase64(bytes: Uint8Array) { /** * 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. + * stores. It sees only the media of `threadId` in the tenant `tenantId`: an + * id of another thread or another tenant 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` @@ -160,13 +169,17 @@ export function toBase64(bytes: Uint8Array) { export function createMediaStore({ persistence, threadId, + tenantId, options = {}, }: { persistence: MediaPersistence threadId: string + /** The tenant of the thread. Without it, the default tenant. */ + tenantId?: string options?: MediaOptions }) { const { artifacts, blobs } = persistence.stores + const storeThreadId = logIdOf({ tenantId, logId: threadId }) const maxBytes = options.maxBytes ?? DEFAULT_MAX_BYTES const kinds = options.kinds ?? ALL_KINDS const tooBig = () => @@ -177,13 +190,13 @@ export function createMediaStore({ async function artifact(id: string) { const record = await artifacts.get(id) - return record?.threadId === threadId ? record : null + return record?.threadId === storeThreadId ? 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 + return record && kind ? toMediaRecord(record, kind, threadId) : null } async function load(id: string, range?: BlobRange) { @@ -216,8 +229,8 @@ export function createMediaStore({ if (size !== undefined && size > maxBytes) throw tooBig() const id = crypto.randomUUID() - const runId = uploadRunId(threadId) - const blobKey = `artifacts/upload:${encodeURIComponent(threadId)}/${id}` + const runId = uploadRunId(storeThreadId) + 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 = @@ -242,7 +255,7 @@ export function createMediaStore({ const record: ArtifactRecord = { artifactId: id, runId, - threadId, + threadId: storeThreadId, blobKey, name: info.name, mimeType: info.mimeType, @@ -250,7 +263,7 @@ export function createMediaStore({ createdAt: Date.now(), } await artifacts.save(record) - return toMediaRecord(record, kind) + return toMediaRecord(record, kind, threadId) }, /** The record of `id`, or `null` when it is missing or another thread's. */ get, @@ -294,6 +307,7 @@ export function createMediaStore({ export function mediaCapture({ persistence, threadId, + tenantId, options = {}, publish, onError, @@ -304,13 +318,17 @@ export function mediaCapture({ generationRuns: GenerationRunStore }> threadId: string + /** The tenant of the thread. Without it, the default tenant. */ + tenantId?: string options?: MediaOptions publish: (record: MediaRecord) => void | Promise onError?: (error: unknown) => void }) { - const store = createMediaStore({ persistence, threadId, options }) + const store = createMediaStore({ persistence, threadId, tenantId, options }) + // The runs and the artifacts are under the store key of the thread, as in + // `createMediaStore`. const persist = withGenerationPersistence(persistence, { - threadId, + threadId: logIdOf({ tenantId, logId: threadId }), maxArtifactBytes: options.maxBytes ?? DEFAULT_MAX_BYTES, }) const keep: GenerationMiddleware = { diff --git a/packages/ai-harness/src/plugins.ts b/packages/ai-harness/src/plugins.ts index 485ae44de9..e296d1bd9f 100644 --- a/packages/ai-harness/src/plugins.ts +++ b/packages/ai-harness/src/plugins.ts @@ -22,7 +22,7 @@ import type { AgentResultOf, AnyAgent, } from './agents' -import type { AgentRun, Operation, TurnInfo } from './types' +import type { AgentRun, Operation, Principal, TurnInfo } from './types' import type { CredentialsAccess } from './auth' import type { AnyCommand, PluginSessionApi } from './commands' import type { ConfigOption } from './config' @@ -39,19 +39,23 @@ export type PluginLifetime = 'session' | 'turn' /** * 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. + * function `text` is called for each turn with that turn, so it can show + * current state, for example the state of the turn's user (`turn.run`). */ export interface PluginPrompt { id: string - text: string | (() => string) + text: string | ((turn: TurnInfo) => 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> + /** + * System prompt text for every chat turn. A function runs for each turn, + * with that turn. + */ + prompts?: ReadonlyArray string) | PluginPrompt> /** Chat middleware, the same type as `chat({ middleware })`. */ middleware?: ReadonlyArray /** @@ -88,10 +92,13 @@ export interface PluginContributions { ) => 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. + * signed in to after the session opened. Called before each chat turn, + * with that turn: `turn.run` has the turn's user and that user's + * credentials. A name that another tool already uses is skipped. */ - discoverTools?: () => ReadonlyArray | Promise> + discoverTools?: ( + turn: TurnInfo, + ) => 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 @@ -112,7 +119,13 @@ export interface PluginState { update: (change: (current: T) => T) => Promise } -/** Run agents from a plugin: in the foreground, in the background, or as a group. */ +/** + * Run agents from a plugin: in the foreground, in the background, or as a + * group. The runs are for the user of a turn plugin's turn. In a session + * plugin, they are for the principal that opened the session: to start an + * agent for the user of a tool call, pass `principal: ctx.context.run.principal` + * to `start`. + */ export interface PluginAgentActions { run: { ( @@ -126,13 +139,13 @@ export interface PluginAgentActions { agent: TAgent, input?: AgentInputOf, /** See `AgentStartOptions`. */ - options?: { wake?: boolean; resume?: boolean }, + options?: { wake?: boolean; resume?: boolean; principal?: Principal }, ): AgentRun> ( name: string, input?: unknown, /** See `AgentStartOptions`. */ - options?: { wake?: boolean; resume?: boolean }, + options?: { wake?: boolean; resume?: boolean; principal?: Principal }, ): AgentRun } /** @@ -269,28 +282,33 @@ export interface PluginSetupContext> { /** 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. + * The credentials of the user of this plugin's turn, for a plugin with + * `lifetime: 'turn'`. + * + * A session plugin has no user: each call throws. Read the credentials of + * the run that calls instead: `ctx.context.run.credentials` in a tool or a + * middleware hook, and `run.credentials` (or `credentials`) of the command + * context in a command. */ 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`. + * The model provider keys of the same user as `credentials`: the key saved + * with `/connect `, else the provider's env var. Build a + * `keyedAdapter(...)` with `await keys.adapter(adapter)` just before the + * call. A missing key throws `AuthRequiredError`. + * + * A session plugin has no user: each call throws. Read + * `ctx.context.run.keys` in a tool or a middleware hook, and `run.keys` of + * the command context in a command. */ 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. + * The turn this plugin is set up for: its input, context, sender, and run. + * Set only for a plugin with `lifetime: 'turn'`. A session plugin has no + * turn at setup. */ turn?: TurnInfo } @@ -421,14 +439,14 @@ interface Owned { export interface MountedPlugins { tools: Array /** Strings, or functions to call for each turn. */ - prompts: Array string)> + prompts: Array string)> commands: Map config: Map /** Contributions to extension points, by point name. */ extensions: Map> adapters: Array> discoverers: Array<{ - discover: () => ReadonlyArray | Promise> + discover: NonNullable owner: string }> preparers: Array<{ @@ -497,6 +515,29 @@ const unavailable = (what: string) => () => { throw new Error(`${what} is only available inside a harness session.`) } +/** Each call of a service of a session plugin, which has no user, fails. */ +const noUser = (what: 'credentials' | 'keys') => async () => { + throw new Error( + `ctx.${what} of a session plugin has no user. In a tool or a middleware hook, read ctx.context.run.${what}. In a command, use the ${what} of its run context.`, + ) +} + +/** @internal `ctx.credentials` of a session plugin. Each call fails. */ +export const SESSION_PLUGIN_CREDENTIALS: CredentialsAccess = { + get: noUser('credentials'), + require: noUser('credentials'), + set: noUser('credentials'), + delete: noUser('credentials'), + list: noUser('credentials'), +} + +/** @internal `ctx.keys` of a session plugin. Each call fails. */ +export const SESSION_PLUGIN_KEYS: ProviderKeys = { + get: noUser('keys'), + require: noUser('keys'), + adapter: noUser('keys'), +} + const NO_SESSION: PluginSessionApi = { threadId: '', principal: undefined, @@ -658,7 +699,7 @@ function checkCapabilityOrder( } function promptOf( - prompt: string | (() => string) | PluginPrompt, + prompt: string | ((turn: TurnInfo) => string) | PluginPrompt, owner: string, index: number, ): PluginPrompt { diff --git a/packages/ai-harness/src/protocol.ts b/packages/ai-harness/src/protocol.ts index f606d1394d..15ab562cc1 100644 --- a/packages/ai-harness/src/protocol.ts +++ b/packages/ai-harness/src/protocol.ts @@ -1,11 +1,18 @@ import { convertSchemaToJsonSchema } from '@tanstack/ai' +import { AccessDeniedError } from './access' import { isRecord } from './utils' import { createInputId, errorText } from './session/helpers' import type { StreamChunk } from '@tanstack/ai' import type { AnyAgent } from './agents' import type { AnyHarness } from './define' import type { HarnessSession, SessionSnapshot } from './session/index' -import type { Cursor, HarnessInput, Principal, Receipt } from './types' +import type { + Cursor, + HarnessInput, + Operation, + Principal, + Receipt, +} from './types' /** The session-tier protocol version. Sent in `subscribe` and `hello`. */ export const HARNESS_PROTOCOL_VERSION = 1 @@ -208,9 +215,10 @@ export function parseControlFrame(data: string): ControlFrame { * only the agents in `expose.agents`. It can run only the commands in * `expose.commands` and the tools in `expose.tools`, and change only the * settings in `expose.settings` and the config keys in `expose.config`. - * Resolves to the receipt. `principal` is who - * sent the input, from your `authorize`, never from the input itself. A chat - * input runs with its credentials. + * Resolves to the receipt. `principal` is who sent the input, from the + * login of the request, never from the input itself. A chat input runs with + * its credentials. Rejects with `AccessDeniedError` when the host refuses + * the input. * * @example * ```ts @@ -228,9 +236,11 @@ export async function applyInput({ input: HarnessInput principal?: Principal }): Promise { + // The log keeps who sent each input, for every input kind. + const user = principal ? { principal } : {} const id = { ...(input.inputId === undefined ? {} : { inputId: input.inputId }), - ...(principal ? { principal } : {}), + ...user, } const sent = { ...id, @@ -244,52 +254,41 @@ export async function applyInput({ reason: 'not_exposed', } switch (input.op) { - case 'prompt': { - const operation = session.send(input.message, { - ...(input.busy ? { busy: input.busy } : {}), - ...(input.ephemeral ? { ephemeral: input.ephemeral } : {}), - ...sent, - }) - // Nobody may await a turn a client started. Its receipt is the answer. - operation.then( - () => {}, - () => {}, + case 'prompt': + return turnReceipt( + session.send(input.message, { + ...(input.busy ? { busy: input.busy } : {}), + ...(input.ephemeral ? { ephemeral: input.ephemeral } : {}), + ...sent, + }), ) - return operation.receipt - } - case 'continue': { - const operation = session.continue({ - ...(input.ephemeral ? { ephemeral: input.ephemeral } : {}), - ...sent, - }) - operation.then( - () => {}, - () => {}, + case 'continue': + return turnReceipt( + session.continue({ + ...(input.ephemeral ? { ephemeral: input.ephemeral } : {}), + ...sent, + }), ) - return operation.receipt - } case 'steer': - case 'followUp': { - const operation = session.send(input.message, { ...sent, mode: input.op }) - operation.then( - () => {}, - () => {}, + case 'followUp': + return turnReceipt( + session.send(input.message, { ...sent, mode: input.op }), ) - return operation.receipt - } case 'resolve': return session.resolve({ resume: input.resume, ...id }) case 'cancel': - return session.cancel({ operationId: input.operationId }) - // It runs no new code, so it needs no `expose` entry, like `cancel`. + return session.cancel({ operationId: input.operationId, ...user }) + // It runs no new code, so it needs no `expose` entry, like `cancel`. It + // stores no input, so there is no sender to keep. case 'background': return session.background(input.toolCallId) case 'cancelInput': - return session.cancelInput(input.inputId) + return session.cancelInput(input.inputId, user) case 'setDelivery': return session.setDelivery({ inputId: input.inputId, delivery: input.delivery, + ...user, }) case 'command': { // A command can change the session, for example `/mode bypass`. @@ -299,8 +298,10 @@ export async function applyInput({ const operation = session.command({ name: input.name, input: input.input, - ...(principal ? { principal } : {}), + ...user, }) + // ponytail: a command has no receipt, so it is accepted at once. A + // command that the host refuses fails with a `RUN_ERROR` event. operation.then( () => {}, () => {}, @@ -315,28 +316,25 @@ export async function applyInput({ return session.answer({ questionId: input.questionId, value: input.value, + ...user, }) case 'config': if (!(harness.expose?.config ?? []).includes(input.key)) { return notExposed } - return session.setConfig({ key: input.key, value: input.value }) + return session.setConfig({ key: input.key, value: input.value, ...user }) case 'tool': { if (!(harness.expose?.tools ?? []).includes(input.name)) { return notExposed } - const operation = session.tool({ - name: input.name, - args: input.args, - ...(input.meta ? { meta: input.meta } : {}), - ...id, - }) - // Nobody may await a turn a client started. Its receipt is the answer. - operation.then( - () => {}, - () => {}, + return turnReceipt( + session.tool({ + name: input.name, + args: input.args, + ...(input.meta ? { meta: input.meta } : {}), + ...id, + }), ) - return operation.receipt } case 'append': { // A client adds only the exposed types, and only to this thread. @@ -380,8 +378,12 @@ export async function applyInput({ if (!handle) { return { inputId: '', status: 'rejected', reason: 'unknown_agent' } } + // The run is the client's: its credentials and keys are theirs. + // ponytail: accepted at once, as a command. A run that the host + // refuses fails with a `RUN_ERROR` event. const operation = handle.start(input.input, { wake: input.detached === true, + ...user, }) return { inputId: operation.id, @@ -407,20 +409,35 @@ export async function applyInput({ reason: 'not_exposed', } } - const operation = session.send(input.message, { - to: input.operationId, - ...id, - ...(input.mode ? { mode: input.mode } : {}), - }) - operation.then( - () => {}, - () => {}, + return turnReceipt( + session.send(input.message, { + to: input.operationId, + ...id, + ...(input.mode ? { mode: input.mode } : {}), + }), ) - return operation.receipt } } } +/** + * The receipt of a turn that a client started. Nobody may await the turn, + * so its failure is handled here. A turn that the host refused throws its + * `AccessDeniedError`, as the other inputs do, so a transport answers 403. + */ +export async function turnReceipt(operation: Operation) { + const failure = operation.then( + () => undefined, + (error: unknown) => error, + ) + const receipt = await operation.receipt + if (receipt.status !== 'rejected') return receipt + // A refused input fails its operation at once. + const error = await failure + if (error instanceof AccessDeniedError) throw error + return receipt +} + /** * What `describe` sends a client: only the commands in `expose.commands` and * the config keys in `expose.config`, so a UI shows only what `applyInput` diff --git a/packages/ai-harness/src/server.ts b/packages/ai-harness/src/server.ts index c0235d3d35..2c5f647615 100644 --- a/packages/ai-harness/src/server.ts +++ b/packages/ai-harness/src/server.ts @@ -2,10 +2,9 @@ // and the protocol helpers that a custom transport uses. export { createHarnessHandler, handleHarnessSocket } from './http' export type { - Authorize, - CanAccess, HarnessHandlerOptions, HarnessSocketOptions, + OpenAccess, } from './http' export { diff --git a/packages/ai-harness/src/session/agents.ts b/packages/ai-harness/src/session/agents.ts index 223ab867a7..e3d751090f 100644 --- a/packages/ai-harness/src/session/agents.ts +++ b/packages/ai-harness/src/session/agents.ts @@ -29,6 +29,7 @@ import { isSameSender, resetCut, referenceNote, + runMiddleware, AGENT_STOPPED, SHUTDOWN_REASON, } from './helpers' @@ -55,6 +56,7 @@ import type { Operation, Principal, Receipt, + RunContext, UserInput, } from '../types' import type { @@ -82,11 +84,7 @@ export abstract class SessionAgents< */ 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 ?? {}), - } + return this.agentHandle(name) } protected override sendToRun( @@ -310,18 +308,21 @@ export abstract class SessionAgents< } /** - * 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. + * What every agent run of `operation` gets: the run context `run`, 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 provider keys of `run` as `ctx.keys`, and + * their tools read `run` as `ctx.context.run`. Their model calls count in + * the usage of `principal`, the user of `run`. With `agents` (the + * subagents of a turn), the agent tools can start them in the background, + * for `principal`. */ protected binding( operation: OperationImpl | OperationImpl, captured: Array, principal: Principal | undefined, + run: RunContext, runPlugins?: MountedPlugins, agents?: ReadonlyArray, ) { @@ -337,7 +338,7 @@ export abstract class SessionAgents< const operation = this.runAgent( agent, call.input, - startOptions, + { ...startOptions, ...(principal ? { principal } : {}) }, undefined, { subagentRunId, @@ -356,6 +357,7 @@ export abstract class SessionAgents< ...mediaCapture({ persistence: this.media, threadId: this.threadId, + tenantId: this.tenantId, options, publish: (record) => { captured.push(record) @@ -370,6 +372,8 @@ export abstract class SessionAgents< }), ], chatMiddleware: [ + // First, so the plugin middleware after it reads the run too. + runMiddleware(run), ...(this.sessionPlugins?.agentMiddleware ?? []), ...(runPlugins?.agentMiddleware ?? []), // A routed root agent gets the thread too, so a reset cuts it as well. @@ -384,7 +388,7 @@ export abstract class SessionAgents< transcribe: options?.transcribe, }), ], - keys: this.keys, + keys: run.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 }), @@ -408,10 +412,11 @@ export abstract class SessionAgents< } 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. + // The agent runs for the user who starts it, also after the turn that + // starts it ends. A child of agent code runs for its parent run's user. + // Without either, it runs for the principal that opened the session. const principal = - options.principal ?? (parent ? parent.current.principal : this.sender()) + options.principal ?? (parent ? parent.current.principal : this.principal) const sender = principal ? { principal } : {} const fromParent = parent ? { @@ -621,11 +626,16 @@ export abstract class SessionAgents< return store.loadThread(agentThread) } - /** Run a group of agents. Every child settles before the group returns. */ + /** + * Run a group of agents, for `principal` (else the principal that opened + * the session). Every child settles before the group returns. + */ protected override async agentGroup( options: { onFailure?: 'cancel-siblings' | 'collect' }, body: (group: AgentGroup) => Promise, + principal?: Principal, ): Promise { + const runOptions = { wake: false, ...(principal ? { principal } : {}) } const started: Array> = [] const cancelOthers = (failed: OperationImpl) => { if (options.onFailure === 'collect') return @@ -641,7 +651,7 @@ export abstract class SessionAgents< } const group: AgentGroup = { run: (target: string | AnyAgent, input?: unknown) => { - const operation = this.runAgent(target, input, { wake: false }) + const operation = this.runAgent(target, input, runOptions) started.push(operation) return Promise.resolve(operation).catch((error: unknown) => { cancelOthers(operation) @@ -649,7 +659,7 @@ export abstract class SessionAgents< }) }, runSettled: (target: string | AnyAgent, input?: unknown) => { - const operation = this.runAgent(target, input, { wake: false }) + const operation = this.runAgent(target, input, runOptions) started.push(operation) return Promise.resolve(operation).then( (value) => ({ ok: true as const, value }), @@ -872,7 +882,15 @@ export abstract class SessionAgents< message && !saved.some((item) => item.id === message.inputId) ? [...saved, asUserMessage(message)] : saved - const binding = this.binding(operation, [], principal) + // The run of this agent: its user's credentials and keys, for its own + // chat calls and its tools, also while another user's turn runs. + const run = this.runContext({ + principal, + runId: operation.id, + ...(chain.parentRunId ? { parentRunId: chain.parentRunId } : {}), + signal: operation.abortController.signal, + }) + const binding = this.binding(operation, [], principal, run) const durable = options.resume ? this.resumable(inputId) : undefined const harnessContext: HarnessAgentContext = { step: durable?.step ?? NO_SESSION_CONTEXT.step, @@ -907,13 +925,16 @@ export abstract class SessionAgents< chatMiddleware: [ // Every run keeps its transcript in the chain's thread. ...(this.chatPersistence - ? [withPersistence(this.chatPersistence)] + ? [ + withPersistence(this.chatPersistence, { + tenantId: this.tenantId, + }), + ] : []), ...(durable?.middleware ?? []), this.agentSteering(chain, inputId, operation), ...binding.chatMiddleware, ], - keys: this.keysOf(principal), budget: chain.budget, // No parent chat() call to take it from, so the session's retention. promptCache: this.promptCache.retention, @@ -1035,38 +1056,45 @@ export abstract class SessionAgents< } /** - * 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. + * The index entry of an agent run: thread `subagent:`, in + * the tenant of this thread, with this thread as its parent and the + * harness of this thread. Its owner is the owner of this thread, else the + * user of the run. It has no members: access to a child goes through its + * parent. A write 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 { tenantId } = this const threadId = `subagent:${child.subagentRunId}` const now = Date.now() - await this.index - .update(threadId, (entry) => ({ + try { + const parent = await this.index.get(this.indexKey) + const ownerId = parent?.ownerId ?? principal?.id + await this.index.update({ tenantId, threadId }, (entry) => ({ ...entry, + ...(tenantId !== undefined ? { tenantId } : {}), threadId, parentThreadId: this.threadId, ...(child.parentToolCallId !== undefined && { parentToolCallId: child.parentToolCallId, }), harness: this.harness.name, + ...(ownerId !== undefined ? { ownerId } : {}), + members: [], 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)}`, - ), + } catch (error) { + this.warn( + operation, + 'harness:sessions', + `The session index was not updated. ${String(error)}`, ) + } } /** diff --git a/packages/ai-harness/src/session/base.ts b/packages/ai-harness/src/session/base.ts index f028bdb178..c8e408f99e 100644 --- a/packages/ai-harness/src/session/base.ts +++ b/packages/ai-harness/src/session/base.ts @@ -16,6 +16,9 @@ import { onUsageOf, } from '../feed' import { createMediaStore } from '../media' +import { logIdOf } from '../log' +import { tenantStores } from '../tenant-stores' +import { SESSION_PLUGIN_CREDENTIALS, SESSION_PLUGIN_KEYS } from '../plugins' import { LEASE } from '../resume' import { HARNESS_EVENTS } from '../types' import { emptyUsage } from '../usage' @@ -25,6 +28,7 @@ import { customEvent, cacheObject, scopeOf, + storedPrincipal, notOpenMessages, tapFeed, } from './helpers' @@ -44,6 +48,8 @@ import type { InboxEntry, InboxStore, MessageStore, + RunStore, + SessionIndexKey, SessionIndexStore, } from '@tanstack/ai-persistence' import type { AuthRequiredError, CredentialsAccess } from '../auth' @@ -56,7 +62,7 @@ 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 { SessionIndexWriter } from '../host' import type { AgentGroup, HarnessPlugin, @@ -74,15 +80,18 @@ import type { Operation, Principal, Receipt, + RunContext, SessionEvent, ThreadSettings, SendOptions, + TurnInfo, } from '../types' import type { AgentRunOptions, AgentStartOptions, AgentChild, AgentHandles, + DynamicAgentHandle, SessionInspection, SessionDependencies, QueuedTurn, @@ -122,9 +131,34 @@ export abstract class SessionBase { readonly registry: AgentRegistryView protected readonly harness: THarness - protected readonly persistence: HarnessPersistence + /** The tenant of this session: the tenant of the principal that opened it. */ + protected readonly tenantId: string | undefined + /** The id of the log of this session in the store, with the tenant. */ + protected readonly storeLogId: string + /** + * The key of this thread in the stores that the session keys by thread + * id: the tenant and the thread id, from `logIdOf`. The metadata rows, the + * inbox, the work claims, and the leases use it. `persistence` gives the + * messages, runs, and interrupts in the tenant already. + */ + protected readonly storeThreadId: string + /** The key of the session index entry of this thread. */ + protected readonly indexKey: SessionIndexKey + /** + * The stores of the host. Its messages, runs, and interrupts are in the + * tenant of this session (see `tenantStores`): they take the thread id of + * the app. + */ + protected readonly persistence: AIPersistence + /** + * The run store of the host, in all tenants. A run id is one key for all + * tenants, so the check of a client run id reads this one. + */ + protected readonly hostRuns: RunStore | undefined protected readonly inbox: InboxStore protected readonly principal: Principal | undefined + /** The owner of the index entry, once the session read it. */ + protected owner: Principal | undefined /** The session log writer of a durable host. It is also the feed. */ protected writer: LogWriter | undefined protected feed: ListenedFeed @@ -272,8 +306,6 @@ export abstract class SessionBase { protected readonly stateWrites = 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 /** Each secret value a credentials access of this session handed out. */ protected readonly secrets = new Set() /** @@ -282,29 +314,39 @@ export abstract class SessionBase { */ 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. + * The credentials of one principal: the user of a run. A save there can + * answer that user's sign-ins. With `canWait`, a `require` with `wait` can + * stop the run for the sign-in. */ protected readonly credentialsOf: ( principal: Principal | undefined, + canWait?: (id: string) => boolean, ) => CredentialsAccess - /** The provider keys of one principal, for an agent that outlives a turn. */ + /** + * The model provider keys of one principal: saved with + * `/connect `, else the env var. + */ protected readonly keysOf: (principal: Principal | undefined) => ProviderKeys + /** What session plugins get. A turn plugin gets `pluginServices(turn)`. */ 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 + /** The host checks access. See `SessionDependencies.access`. */ + protected readonly hasAccess: boolean constructor(deps: SessionDependencies) { this.harness = deps.harness as THarness this.threadId = deps.threadId this.logId = deps.logId + this.tenantId = deps.tenantId + this.storeLogId = deps.storeLogId + this.storeThreadId = logIdOf({ + tenantId: deps.tenantId, + logId: deps.threadId, + }) + this.indexKey = { tenantId: deps.tenantId, threadId: deps.threadId } // The session value wins over the harness value, field by field. const own = cacheObject(deps.promptCache) const shared = cacheObject(this.harness.promptCache) @@ -312,12 +354,17 @@ export abstract class SessionBase { retention: own?.retention ?? shared?.retention ?? 'short', key: own?.key ?? shared?.key ?? deps.threadId, } - this.persistence = deps.persistence + const { stores } = deps.persistence + this.persistence = { + stores: { ...stores, ...tenantStores(stores, deps.tenantId) }, + } + this.hostRuns = stores.runs this.inbox = deps.inbox this.media = deps.media this.mediaStore = createMediaStore({ persistence: deps.media, threadId: deps.threadId, + ...(deps.tenantId !== undefined ? { tenantId: deps.tenantId } : {}), options: this.harness.media, }) this.principal = deps.principal @@ -329,7 +376,7 @@ export abstract class SessionBase { this.lease = deps.lease this.log = deps.log // A durable host gets its log view in `open()`. - this.messages = deps.persistence.stores.messages ?? notOpenMessages + this.messages = this.persistence.stores.messages ?? notOpenMessages const onMissing = (error: AuthRequiredError) => this.feed.publish( 'session', @@ -338,71 +385,26 @@ export abstract class SessionBase { ...(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(), - secrets: this.secrets, - }, - ) - this.keys = providerKeysFor(this.credentialAccess, onMissing, this.secrets) - this.credentialsOf = (principal) => + this.credentialsOf = (principal, canWait) => credentialsFor( deps.credentials, scopeOf(deps.threadId, principal), onMissing, - { onSet: () => this.resumeSignIns(), secrets: this.secrets }, + { + ...(canWait ? { canWait } : {}), + onSet: () => this.resumeSignIns(), + secrets: this.secrets, + }, ) + this.hasAccess = deps.access !== undefined + // Without access, one user runs the host, so the env vars are theirs. + const envFallback = deps.access?.envKeys ?? true this.keysOf = (principal) => - providerKeysFor(this.credentialsOf(principal), onMissing, this.secrets) - 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, {}), - ), - } + providerKeysFor(this.credentialsOf(principal), onMissing, { + secrets: this.secrets, + envFallback, + }) + this.services = this.pluginServices() this.registry = this.agentRegistry for (const agent of this.harness.agents ?? []) { this.agentRegistry.add(agent, 'the harness') @@ -411,18 +413,26 @@ export abstract class SessionBase { 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 ?? {}), - } - }, + get: (_target, name) => + typeof name === 'string' ? this.agentHandle(name) : undefined, }) } + /** + * @internal `run` and `start` of the agent `name`, for `principal` when + * given. `session.agents` and `session.agent` use it, and so does the + * handle of `host.open`, which passes its own principal. + */ + agentHandle(name: string, principal?: Principal): DynamicAgentHandle { + const user = principal ? { principal } : {} + return { + run: (input?: unknown, options?: AgentRunOptions) => + this.runAgent(name, input, { ...options, ...user, wake: false }), + start: (input?: unknown, options?: AgentStartOptions) => + this.runAgent(name, input, { ...options, ...user }), + } + } + /** * An operation of this session by id: a running one, an agent run, or one * of the last 500 other operations that settled. @@ -577,12 +587,100 @@ export abstract class SessionBase { /** * 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. + * the session. Turns run one at a time, so one value is enough. It never + * picks credentials or keys: those come from the run context of each run. */ protected sender(): Principal | undefined { return this.turnPrincipal ?? this.principal } + /** + * The run context of one run: the user who started it, and the + * credentials and provider keys of that user. `canWait` says if a sign-in + * wait of a tool can stop this run now. + */ + protected runContext(options: { + principal: Principal | undefined + runId: string + parentRunId?: string + signal: AbortSignal + canWait?: (id: string) => boolean + }): RunContext { + const { principal } = storedPrincipal(options.principal) + return { + ...(this.tenantId !== undefined ? { tenantId: this.tenantId } : {}), + principal, + logId: this.storeLogId, + threadId: this.threadId, + runId: options.runId, + ...(options.parentRunId ? { parentRunId: options.parentRunId } : {}), + signal: options.signal, + credentials: this.credentialsOf(options.principal, options.canWait), + keys: this.keysOf(options.principal), + } + } + + /** + * What plugins get from the session. A session plugin has no user: its + * `credentials` and `keys` throw, its `session.principal` is the principal + * that opened the session, and its agents run for that principal. A turn + * plugin gets its `turn`: they are all of the run of that turn. + */ + protected pluginServices(turn?: TurnInfo): PluginServices { + const principal = turn?.principal + const operationId = turn?.operationId + return { + 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: turn?.run.credentials ?? SESSION_PLUGIN_CREDENTIALS, + keys: turn?.run.keys ?? SESSION_PLUGIN_KEYS, + session: this.pluginApi(principal, undefined, operationId), + sessionWith: (metadataSchema) => + this.pluginApi(principal, metadataSchema, operationId), + agents: this.pluginAgents(principal), + commandsChanged: () => + this.feed.publish( + 'session', + customEvent(HARNESS_EVENTS.commandsChanged, {}), + ), + } + } + + /** + * `ctx.agents` of a plugin. The runs are for `principal`, else for the + * principal that opened the session. A `principal` option of `start` wins. + */ + protected pluginAgents(principal?: Principal): PluginAgentActions { + const user = principal ? { principal } : {} + return { + run: ((target: string | AnyAgent, input?: unknown) => + this.runAgent(target, input, { + ...user, + wake: false, + })) as PluginAgentActions['run'], + start: (( + target: string | AnyAgent, + input?: unknown, + options?: AgentStartOptions, + ) => + this.runAgent(target, input, { + ...user, + ...options, + })) as PluginAgentActions['start'], + group: (options, body) => this.agentGroup(options, body, principal), + } + } + protected emitPluginEvent( plugin: string, name: string, @@ -607,7 +705,7 @@ export abstract class SessionBase { protected pluginState(plugin: string, initial: T): PluginState { const metadata = this.persistence.stores.metadata const namespace = `plugin:${plugin}` - const key = this.threadId + const key = this.storeThreadId const read = async (): Promise<{ value: T; revision: string | null }> => { if (metadata?.getVersioned) { const stored = await metadata.getVersioned(namespace, key) @@ -696,7 +794,7 @@ export abstract class SessionBase { // returns false. The timer tries again while this thread has work, // so the work gets a claim once that one expires. const isHeld = await claims.claim({ - threadId: this.threadId, + threadId: this.storeThreadId, harness: this.harness.name, ownerId: this.hostId, until: Date.now() + ttlMs, @@ -738,7 +836,7 @@ export abstract class SessionBase { protected async isOwnedElsewhere(): Promise { const claims = this.persistence.stores.workClaims if (!this.writer || !claims) return false - const held = await claims.get(this.threadId) + const held = await claims.get(this.storeThreadId) return ( held !== null && held.ownerId !== this.hostId && held.until > Date.now() ) @@ -758,7 +856,9 @@ export abstract class SessionBase { } const check = async () => { if (this.closing) return stop() - const pending = await inbox.listPending(this.threadId).catch(() => null) + const pending = await inbox + .listPending(this.storeThreadId) + .catch(() => null) if (pending?.length === 0) return stop() void this.drainInbox() } @@ -785,10 +885,10 @@ export abstract class SessionBase { const claims = this.persistence.stores.workClaims if (!this.writer || !inbox || !claims || this.closing) return try { - const pending = await inbox.listPending(this.threadId) + const pending = await inbox.listPending(this.storeThreadId) if (pending.length === 0) return const isHeld = await claims.claim({ - threadId: this.threadId, + threadId: this.storeThreadId, harness: this.harness.name, ownerId: this.hostId, until: Date.now() + (this.lease?.ttlMs ?? LEASE.ttlMs), @@ -808,11 +908,11 @@ export abstract class SessionBase { // The host that stored an answer did not know if it is a secret. // This host does, and logs it hidden. The entry goes. if ((entry.input as HarnessInput).op === 'answer') { - await inbox.delete({ threadId: this.threadId, inputId }) + await inbox.delete({ threadId: this.storeThreadId, inputId }) continue } await inbox.markApplied({ - threadId: this.threadId, + threadId: this.storeThreadId, inputId, operationId: 'forwarded', }) @@ -831,9 +931,10 @@ export abstract class SessionBase { const input = entry.input as HarnessInput const { inputId } = entry const principal = entry.principal + const user = principal ? { principal } : {} const sent = { inputId, - ...(principal ? { principal } : {}), + ...user, ...('context' in input && input.context !== undefined ? { context: input.context } : {}), @@ -888,13 +989,18 @@ export abstract class SessionBase { }), ) case 'cancel': - return session.cancel({ operationId: input.operationId, inputId }) + return session.cancel({ + operationId: input.operationId, + inputId, + ...user, + }) case 'cancelInput': - return session.cancelInput(input.inputId) + return session.cancelInput(input.inputId, user) case 'setDelivery': return session.setDelivery({ inputId: input.inputId, delivery: input.delivery, + ...user, }) case 'command': return settle( @@ -902,7 +1008,7 @@ export abstract class SessionBase { name: input.name, input: input.input, inputId, - ...(principal ? { principal } : {}), + ...user, }), ) case 'answer': @@ -910,12 +1016,14 @@ export abstract class SessionBase { questionId: input.questionId, value: input.value, inputId, + ...user, }) case 'config': return session.setConfig({ key: input.key, value: input.value, inputId, + ...user, }) case 'tool': return stored( @@ -924,7 +1032,7 @@ export abstract class SessionBase { name: input.name, args: input.args as never, inputId, - ...(principal ? { principal } : {}), + ...user, ...(input.meta ? { meta: input.meta } : {}), }), ) @@ -945,7 +1053,7 @@ export abstract class SessionBase { try { await this.claimWrite await this.persistence.stores.workClaims?.release({ - threadId: this.threadId, + threadId: this.storeThreadId, ownerId: this.hostId, }) } catch (error) { @@ -1172,6 +1280,7 @@ export abstract class SessionBase { protected abstract agentGroup( options: { onFailure?: 'cancel-siblings' | 'collect' }, body: (group: AgentGroup) => Promise, + principal?: Principal, ): Promise protected abstract pluginApi( principal?: Principal, diff --git a/packages/ai-harness/src/session/handle.ts b/packages/ai-harness/src/session/handle.ts new file mode 100644 index 0000000000..b1ce2e2a79 --- /dev/null +++ b/packages/ai-harness/src/session/handle.ts @@ -0,0 +1,193 @@ +import { AccessDeniedError } from '../access' +import { isSameSender } from './helpers' +import type { AccessAction, AccessRun } from '../access' +import type { AnyHarness } from '../define' +import type { + AgentSendOptions, + HarnessInput, + Principal, + SendOptions, + UserInput, +} from '../types' +import type { SessionImpl } from './harness-session' +import type { AgentRunOptions, AgentStartOptions } from './types' + +/** What the handle of a host with access checks. */ +export interface HandleGuard { + /** Throws `AccessDeniedError` unless `authorize` allows `action` now. */ + authorize: (action: AccessAction, run?: AccessRun) => Promise + /** `host.revoke` took the access of the principal of the handle. */ + isRevoked: () => boolean +} + +/** The action that `authorize` gets for an input. */ +function actionOf(input: HarnessInput): AccessAction { + switch (input.op) { + case 'resolve': + case 'answer': + return 'answer' + case 'cancel': + case 'cancelInput': + return 'cancel' + case 'configure': + case 'config': + return 'configure' + case 'prompt': + case 'steer': + case 'followUp': + case 'continue': + case 'tool': + case 'reset': + case 'command': + case 'agent': + case 'agentMessage': + case 'setDelivery': + case 'background': + case 'append': + case 'revert': + case 'unrevert': + return 'send' + } +} + +type Session = SessionImpl + +/** + * The session that `host.open` returns, bound to `principal`: a call that + * passes no `principal` acts as it. It has every member of the session, and + * `this` stays the session. + * + * With `guard` (the host has access), each action asks `authorize` first. An + * input is checked before the session stores it. A call that passes another + * principal throws. After `host.revoke`, every call throws + * `AccessDeniedError`. Reads ask `authorize` once, when the host opens the + * handle. + */ +export function sessionHandle({ + session, + principal, + guard, +}: { + session: Session + principal: Principal | undefined + guard?: HandleGuard +}) { + const { threadId } = session + const live = (action: AccessAction) => { + if (guard?.isRevoked()) throw new AccessDeniedError({ action, threadId }) + } + /** Who acts in a call that passes `passed`. */ + const actor = (passed: Principal | undefined) => { + if (passed === undefined) return principal + const isOther = guard !== undefined && !isSameSender(passed, principal) + if (isOther) { + throw new Error( + `This handle acts as ${JSON.stringify(principal?.id)}. Open another handle for ${JSON.stringify(passed.id)}.`, + ) + } + return passed + } + /** `options`, with the principal that acts. */ + const as = (options: T) => { + const acting = actor(options.principal) + return acting ? { ...options, principal: acting } : options + } + /** A call that stores an input: the session checks the input first. */ + const storing = (action: AccessAction, call: () => T) => { + live(action) + if (!guard) return call() + return session.checked( + ({ input, run }) => guard.authorize(actionOf(input), run), + call, + ) + } + /** A call that stores no input: it runs after `authorize` allows `action`. */ + const acting = async (action: AccessAction, call: () => Promise) => { + live(action) + await guard?.authorize(action) + return call() + } + const agentOf = (name: string) => ({ + run: (input?: unknown, options?: AgentRunOptions) => + storing('send', () => + session.agentHandle(name, principal).run(input, options), + ), + start: (input?: unknown, options?: AgentStartOptions) => + storing('send', () => + session + .agentHandle(name, actor(options?.principal)) + .start(input, options), + ), + }) + const agents = new Proxy(session.agents, { + get: (_target, name) => + typeof name === 'string' ? agentOf(name) : undefined, + }) + const methods = new Map( + Object.entries({ + send: (message: UserInput, options?: SendOptions | AgentSendOptions) => + storing('send', () => + options?.to === undefined + ? session.send(message, as(options ?? {})) + : session.send(message, as(options)), + ), + tool: (options: Parameters['tool']>[0]) => + storing('send', () => session.tool(as(options))), + continue: (options?: Parameters['continue']>[0]) => + storing('send', () => session.continue(as(options ?? {}))), + command: (options: Parameters['command']>[0]) => + storing('send', () => session.command(as(options))), + reset: (options?: Parameters['reset']>[0]) => + storing('send', () => session.reset(as(options ?? {}))), + setDelivery: (options: Parameters['setDelivery']>[0]) => + storing('send', () => session.setDelivery(as(options))), + resolve: (options: Parameters['resolve']>[0]) => + storing('answer', () => session.resolve(as(options))), + resolveForRun: ( + options: Parameters['resolveForRun']>[0], + ) => storing('answer', () => session.resolveForRun(as(options))), + answer: (options: Parameters['answer']>[0]) => + storing('answer', () => session.answer(as(options))), + cancel: (options?: Parameters['cancel']>[0]) => + storing('cancel', () => session.cancel(as(options ?? {}))), + cancelInput: ( + target: string, + options?: Parameters['cancelInput']>[1], + ) => + storing('cancel', () => session.cancelInput(target, as(options ?? {}))), + configure: (options: Parameters['configure']>[0]) => + storing('configure', () => session.configure(as(options))), + setConfig: (options: Parameters['setConfig']>[0]) => + storing('configure', () => session.setConfig(as(options))), + agents, + agent: (name: string) => + session.agent(name) === undefined ? undefined : agentOf(name), + // These store no input, so the handle asks `authorize` itself. + revert: (messageId: string) => + acting('send', () => session.revert(messageId)), + unrevert: () => acting('send', () => session.unrevert()), + background: (toolCallId?: string) => + acting('send', () => session.background(toolCallId)), + append: (records: Parameters['append']>[0]) => + acting('send', () => session.append(records)), + putMedia: (upload: Parameters['putMedia']>[0]) => + acting('send', () => session.putMedia(upload)), + updateEntry: (update: Parameters['updateEntry']>[0]) => + acting('configure', () => session.updateEntry(update)), + reload: () => acting('configure', () => session.reload()), + }), + ) + // ponytail: a Proxy, so the handle has every member of the session and + // stays an instance of it. Only the members above change. + return new Proxy(session, { + get(target, key) { + if (methods.has(key)) return methods.get(key) + const value: unknown = Reflect.get(target, key) + if (typeof value !== 'function') return value + return (...args: Array) => { + live('read') + return Reflect.apply(value, target, args) + } + }, + }) +} diff --git a/packages/ai-harness/src/session/harness-session.ts b/packages/ai-harness/src/session/harness-session.ts index 550a96b2e0..49fcb68d81 100644 --- a/packages/ai-harness/src/session/harness-session.ts +++ b/packages/ai-harness/src/session/harness-session.ts @@ -40,7 +40,13 @@ export interface HarnessSession< > extends Omit< SessionImpl, // The public members that only the host and the handler call. - 'open' | 'adoptFork' | 'onIdle' | 'resolveForRun' + | 'open' + | 'adoptFork' + | 'onIdle' + | 'resolveForRun' + | 'checked' + | 'cancelWorkOf' + | 'agentHandle' > {} /** @@ -67,7 +73,7 @@ export class SessionImpl< : revertOf( await this.persistence.stores.metadata?.get( 'harness:revert', - this.threadId, + this.storeThreadId, ), ) await this.loadInterrupted() @@ -225,10 +231,12 @@ export class SessionImpl< this.stopOnLogFailure(error), ) this.feed = tapFeed(this.writer, this.notifyHost, this.eventListeners) + // The logs of the subagent children are in the tenant of this session. const view = { writer: this.writer, store, ...(project ? { project } : {}), + ...(this.tenantId !== undefined ? { tenantId: this.tenantId } : {}), } this.messages = sessionMessageStore(view) this.engine = engineMessageStore(view) @@ -343,7 +351,7 @@ export class SessionImpl< /** 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, + this.storeThreadId, ) for (const entry of pending) { const input = entry.input as HarnessInput @@ -679,7 +687,7 @@ export class SessionImpl< const isAlive = leases && input.operationId ? await leases.isAlive({ - threadId: this.threadId, + threadId: this.storeThreadId, inputId: input.inputId, operationId: input.operationId, attempt: input.attempt, diff --git a/packages/ai-harness/src/session/helpers.ts b/packages/ai-harness/src/session/helpers.ts index 6689d0a5b6..8ba1dd2cf9 100644 --- a/packages/ai-harness/src/session/helpers.ts +++ b/packages/ai-harness/src/session/helpers.ts @@ -25,6 +25,7 @@ import type { MediaKind, OperationStatus, Principal, + RunContext, ThreadSettings, } from '../types' import type { @@ -243,6 +244,31 @@ export function turnContext(harness: unknown, input: unknown) { return harness !== undefined ? harness : input } +/** + * The chat() context of a run: `context` with the server's `run`. A `run` + * key that a client sent is replaced. + */ +export function withRun(context: unknown, run: RunContext) { + if (context === undefined || context === null) return { run } + // ponytail: a context that is not a plain object (a class instance, an + // array, a string) stays as it is, with no `run`: a copy would lose its + // prototype. Tools that read `ctx.context.run` then fail. Wrap such a + // value in a plain object if a harness needs both. + return isPlainObject(context) ? { ...context, run } : context +} + +/** + * Middleware that gives each chat() call of an agent run the run context, + * as `ctx.context.run`. A tool reads the context when it runs, so it gets it + * too. A nested child gets this middleware from the binding as well. + */ +export const runMiddleware = (run: RunContext): AnyChatMiddleware => ({ + name: 'harness:run-context', + setup: (ctx) => { + ctx.context = withRun(ctx.context, run) + }, +}) + /** Where `stores.metadata` keeps the stored settings of each thread. */ export const SETTINGS = 'harness:settings' diff --git a/packages/ai-harness/src/session/inputs.ts b/packages/ai-harness/src/session/inputs.ts index fb94fa1532..35606a3088 100644 --- a/packages/ai-harness/src/session/inputs.ts +++ b/packages/ai-harness/src/session/inputs.ts @@ -13,6 +13,7 @@ import { checkHostRecords, } from './helpers' import { KEPT_SETTLED, SessionBase } from './base' +import type { AccessRun } from '../access' import type { OperationImpl } from '../operation' import type { LogRecord } from '@tanstack/ai-persistence' import type { HarnessRecord, InputState } from '../log' @@ -40,6 +41,15 @@ export type Admission = | 'forwarded' | 'owned_elsewhere' +/** + * The access check of one input, before `accept` stores it: the input, and + * the run that an answer or a cancel targets. It throws to refuse the input. + */ +export type InputCheck = (check: { + input: HarnessInput + run: AccessRun | undefined +}) => Promise + /** * A layer of {@link HarnessSession}: admitting inputs, receipts, and * settlements. @@ -50,6 +60,97 @@ export abstract class SessionInputs< /** Chat inputs that this host passed to the owner of the thread. */ protected readonly forwarded = new Set() + /** The check that the next `accept` runs. See `checked`. */ + private pendingCheck: InputCheck | undefined + + /** + * @internal Run `call`, and run `check` on the input that it stores, + * before the input is stored. The handle of `host.open` checks access + * through it. Each method that stores an input calls `accept` before its + * first `await`, so the check reaches the input of `call` and no other. A + * call that returns before `accept` stores nothing. + */ + checked(check: InputCheck, call: () => T): T { + this.pendingCheck = check + try { + return call() + } finally { + this.pendingCheck = undefined + } + } + + /** + * The run that `input` answers or cancels, and who started it, when the + * session knows. Other inputs target no run. + */ + protected runOf(input: HarnessInput): AccessRun | undefined { + switch (input.op) { + case 'resolve': { + const stopped = this.interrupted + if (!stopped) return undefined + return { runId: stopped.runId, ...idOf(stopped.principal) } + } + case 'answer': { + const operationId = this.questions.get(input.questionId)?.operationId + return operationId === undefined + ? undefined + : this.startedBy(operationId) + } + case 'cancel': { + const operationId = input.operationId ?? this.activeTurn?.id + return operationId === undefined + ? undefined + : this.startedBy(operationId) + } + case 'cancelInput': + return { runId: input.inputId, ...idOf(this.senderOf(input.inputId)) } + case 'prompt': + case 'steer': + case 'followUp': + case 'continue': + case 'tool': + case 'reset': + case 'command': + case 'agent': + case 'agentMessage': + case 'setDelivery': + case 'background': + case 'append': + case 'configure': + case 'config': + case 'revert': + case 'unrevert': + return undefined + } + } + + /** The operation `operationId`, with the id of who started it when known. */ + protected startedBy(operationId: string): AccessRun { + const chain = this.chainOf.get(operationId) + const agentRun = [...(chain?.runs ?? []), ...(chain?.followUps ?? [])].find( + (run) => run.operation.id === operationId, + ) + const inputId = this.operationInputs.get(operationId) + const principal = + this.activeTurn?.id === operationId + ? this.turnPrincipal + : agentRun + ? agentRun.principal + : inputId !== undefined + ? this.senderOf(inputId) + : undefined + return { runId: operationId, ...idOf(principal) } + } + + /** Who sent the input `inputId`: from the log, or from what waits. */ + protected senderOf(inputId: string) { + return ( + this.writer?.state.inputs.get(inputId)?.principal ?? + this.steerQueue.find((steer) => steer.inputId === inputId)?.principal ?? + this.queue.find((turn) => turn.inputId === inputId)?.principal + ) + } + /** * 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` @@ -234,7 +335,7 @@ export abstract class SessionInputs< if (!this.writer) { this.inbox .markApplied({ - threadId: this.threadId, + threadId: this.storeThreadId, inputId: steer.inputId, operationId, }) @@ -327,7 +428,7 @@ export abstract class SessionInputs< protected async applyControl(inputId: string, change: HarnessRecord) { if (!this.writer) { await this.inbox.markApplied({ - threadId: this.threadId, + threadId: this.storeThreadId, inputId, operationId: 'session', }) @@ -351,7 +452,9 @@ export abstract class SessionInputs< /** * 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. + * tick is a duplicate too. `principal` is who sent it. The access check of + * a `checked` call runs first: a refused input is not stored, and the + * error rejects this call. */ protected async accept( inputId: string, @@ -360,6 +463,9 @@ export abstract class SessionInputs< /** The AG-UI run id of a turn, so a forwarded turn runs as it. */ runId?: string, ): Promise { + const check = this.pendingCheck + this.pendingCheck = undefined + if (check) await check({ input, run: this.runOf(input) }) const sender = storedPrincipal(principal) const payload = inputKey(input, sender.principal) const local = this.admitted.get(inputId) @@ -396,7 +502,7 @@ export abstract class SessionInputs< } else { const stored = await this.inbox.append({ inputId, - threadId: this.threadId, + threadId: this.storeThreadId, input, createdAt: at, ...sender, @@ -445,7 +551,7 @@ export abstract class SessionInputs< } await inbox.append({ inputId, - threadId: this.threadId, + threadId: this.storeThreadId, // The owner runs the turn as `runId`, so this host streams it from the log. input: runId === undefined ? input : { ...input, runId }, createdAt: Date.now(), @@ -459,11 +565,12 @@ export abstract class SessionInputs< } /** - * 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. + * Does `stores.runs` have a run with this id, on any thread of any tenant? + * 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 + const runs = this.hostRuns if (!runs) return false return runs.get(runId).then( (run) => run !== null, @@ -478,7 +585,7 @@ export abstract class SessionInputs< ): Promise { if (!this.writer) { await this.inbox.markApplied({ - threadId: this.threadId, + threadId: this.storeThreadId, inputId, operationId, }) @@ -517,7 +624,7 @@ export abstract class SessionInputs< .catch(() => {}) } else { this.inbox - .markRejected({ threadId: this.threadId, inputId, reason }) + .markRejected({ threadId: this.storeThreadId, inputId, reason }) .catch((error: unknown) => this.warnSession('harness:inbox', error)) } const receipt: Receipt = { inputId, status: 'rejected', reason } @@ -745,6 +852,10 @@ export abstract class SessionInputs< } } +/** The `principalId` field of an {@link AccessRun}, when the run has a user. */ +const idOf = (principal: { id: string } | undefined) => + principal ? { principalId: principal.id } : {} + /** The reason of a rejected input record. A rejected record always has one. */ function rejectedReason(logged: InputState): string { if (logged.reason === undefined) { diff --git a/packages/ai-harness/src/session/requests.ts b/packages/ai-harness/src/session/requests.ts index db6dabff4c..de9d6aeb34 100644 --- a/packages/ai-harness/src/session/requests.ts +++ b/packages/ai-harness/src/session/requests.ts @@ -559,7 +559,9 @@ export abstract class SessionRequests< * Stop the operation `operationId`, or the running turn without it. A turn * that waits in the queue never runs. Refused with `'not_running'` when * nothing runs or the operation ended. `inputId` means the same as in - * `send`: a second cancel with the same id is a duplicate. + * `send`: a second cancel with the same id is a duplicate. `principal` is + * who cancels. Default: the principal that opened the session. The log + * keeps it. * * @example * ```ts @@ -567,11 +569,15 @@ export abstract class SessionRequests< * ``` */ async cancel( - options?: CancelOptions & { inputId?: string }, + options?: CancelOptions & { inputId?: string; principal?: Principal }, ): Promise { const operationId = options?.operationId const inputId = options?.inputId ?? createInputId() - const admission = await this.accept(inputId, { op: 'cancel', operationId }) + const admission = await this.accept( + inputId, + { op: 'cancel', operationId }, + options?.principal, + ) if (admission !== 'new') return this.duplicateReceipt(inputId, admission) const target = operationId ? this.operations.get(operationId) @@ -656,11 +662,34 @@ export abstract class SessionRequests< return { inputId, status: 'accepted', operationId: target.id } } + /** + * @internal Cancel the queued and running work that the user `userId` + * started: `host.sessions.unshare` stops it. `principal` is who cancels, + * and the log keeps it. + */ + async cancelWorkOf({ + userId, + principal, + }: { + userId: string + principal: Principal + }): Promise { + const theirs = [...this.operations.values()].filter( + (operation) => + !operation.isSettled() && + this.startedBy(operation.id).principalId === userId, + ) + for (const operation of theirs) { + await this.cancel({ operationId: operation.id, principal }) + } + } + /** * 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 })`. + * `cancel({ operationId })`. `principal` is who cancels. Default: the + * principal that opened the session. The log keeps it. * * @example * ```ts @@ -668,12 +697,16 @@ export abstract class SessionRequests< * await session.cancelInput('tests') * ``` */ - async cancelInput(target: string): Promise { + async cancelInput( + target: string, + options?: { principal?: Principal }, + ): Promise { const inputId = createInputId() - const admission = await this.accept(inputId, { - op: 'cancelInput', - inputId: target, - }) + const admission = await this.accept( + inputId, + { op: 'cancelInput', inputId: target }, + options?.principal, + ) if (isElsewhere(admission)) return this.duplicateReceipt(inputId, admission) const taken = this.takeWaiting(target) ?? this.takeQueuedContinue(target) if (!taken) return this.reject(inputId, 'not_waiting') @@ -685,7 +718,7 @@ export abstract class SessionRequests< // pending entries. if (!this.writer) await this.inbox.markRejected({ - threadId: this.threadId, + threadId: this.storeThreadId, inputId: target, reason: 'cancelled', }) @@ -708,6 +741,8 @@ export abstract class SessionRequests< * 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. + * `principal` is who moves it. Default: the principal that opened the + * session. The log keeps it. * * @example * ```ts @@ -718,16 +753,18 @@ export abstract class SessionRequests< async setDelivery({ inputId: target, delivery, + principal, }: { inputId: string delivery: WaitingInput['delivery'] + principal?: Principal }): Promise { const inputId = createInputId() - const admission = await this.accept(inputId, { - op: 'setDelivery', - inputId: target, - delivery, - }) + const admission = await this.accept( + inputId, + { op: 'setDelivery', inputId: target, delivery }, + principal, + ) if (isElsewhere(admission)) return this.duplicateReceipt(inputId, admission) const current = this.inputs().find((input) => input.inputId === target) if (current?.delivery !== delivery) { @@ -779,20 +816,28 @@ export abstract class SessionRequests< } /** - * The session API of plugins. With `principal` (a command), it acts for - * that user. Without it, it follows the sender of the running turn. With - * `operationId` (a command), its questions belong to that operation. + * The session API of plugins. With `principal` (a command, a turn plugin), + * it acts for that user. Without it (a session plugin), its `principal` is + * the principal that opened the session. With `operationId` (a command, a + * turn plugin), its questions belong to that operation. */ protected override pluginApi( principal?: Principal, metadataSchema?: SchemaInput, operationId?: string, ): PluginSessionApi { + // ponytail: D9. Without `principal`, `send`, `note`, and `ask` follow the + // running turn: its sender, and its operation for a question. PR 7 + // rebuilds them on the run that calls. const sender = () => principal ?? this.sender() + // A session plugin acts for the owner of the entry, not for whoever + // opened the session first on this host. The owner is known only after + // the session read its entry, so it is read at each call. + const owner = () => principal ?? this.owner ?? this.principal return { threadId: this.threadId, get principal() { - return sender() + return owner() }, snapshot: () => this.snapshot(), // A plugin turn always waits its turn. It cannot see a busy rejection. @@ -815,7 +860,7 @@ export abstract class SessionRequests< transcript: () => this.transcript(), replaceTranscript: (messages) => this.messages.saveThread(this.threadId, messages), - entry: () => this.index.get(this.threadId), + entry: () => this.index.get(this.indexKey), updateEntry: (update) => { // Only these fields: a plugin cannot change the owner, the parent, // or the pin. @@ -829,7 +874,7 @@ export abstract class SessionRequests< metadata: change.metadata, usage: change.usage, } - return this.index.update(this.threadId, (entry) => + return this.index.update(this.indexKey, (entry) => changedEntry( entry, typeof update === 'function' @@ -848,7 +893,7 @@ export abstract class SessionRequests< 'session', customEvent(HARNESS_EVENTS.authRequired, { ...info }), ), - setConfig: (change) => this.setConfig(change), + setConfig: (change) => this.setConfig({ ...change, principal: owner() }), settings: () => this.settings(), } } diff --git a/packages/ai-harness/src/session/state.ts b/packages/ai-harness/src/session/state.ts index a1b16c6d8b..1c068173e3 100644 --- a/packages/ai-harness/src/session/state.ts +++ b/packages/ai-harness/src/session/state.ts @@ -8,6 +8,7 @@ import { loadChildTranscript } from '@tanstack/ai-persistence' import { checkConfigValue } from '../config' import { isRecord } from '../utils' import { changedEntry } from '../session-entry' +import { threadOf } from '../log' import { mediaIdOf, mediaOfMessage, mediaPart } from '../media-ref' import { BackgroundJobs, @@ -89,19 +90,46 @@ export abstract class SessionState< /** * 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. + * time, in the tenant of this session, else a new `updatedAt`. The owner + * is who opened the thread first, so an open by another user does not + * change it. The owner is the first member, with the role `owner`. */ 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), - })) + const { principal, tenantId, threadId } = this + const opened = await this.index.update(this.indexKey, (entry) => { + const current: SessionIndexEntry = entry ?? { + ...(tenantId !== undefined ? { tenantId } : {}), + threadId, + members: [], + createdAt: now, + updatedAt: now, + } + // A new entry belongs to its opener. Without access, an entry with no + // owner goes to the next opener with a principal too. With access, it + // stays without an owner: `authorize` decides who may open it. + const isClaim = + principal !== undefined && + (entry === undefined || + (current.ownerId === undefined && !this.hasAccess)) + return { + ...current, + harness: current.harness ?? this.harness.name, + updatedAt: now, + ...(isClaim + ? { + ownerId: principal.id, + members: [{ userId: principal.id, role: 'owner' as const }], + } + : {}), + } + }) + if (opened?.ownerId !== undefined) { + this.owner = { + id: opened.ownerId, + ...(tenantId !== undefined ? { tenantId } : {}), + } + } } snapshot(): SessionSnapshot { @@ -248,7 +276,7 @@ export abstract class SessionState< const { metadata } = this.persistence.stores this.threadSettings = source.settings() if (Object.keys(this.threadSettings).length > 0) { - await metadata?.set(SETTINGS, this.threadId, this.threadSettings) + await metadata?.set(SETTINGS, this.storeThreadId, this.threadSettings) } for (const [key, value] of source.configValues) { if (this.sessionPlugins?.config.has(key)) @@ -257,7 +285,7 @@ export abstract class SessionState< if (this.configValues.size > 0) { await metadata?.set( 'harness:config', - this.threadId, + this.storeThreadId, Object.fromEntries(this.configValues), ) } @@ -286,18 +314,28 @@ export abstract class SessionState< return result } - /** Change a session setting. It applies at the next turn. */ + /** + * Change a session setting. It applies at the next turn. `principal` is + * who changes it. Default: the principal that opened the session. The log + * keeps it. + */ async setConfig({ key, value, inputId = createInputId(), + principal, }: { key: string value: unknown /** @internal The id of a forwarded input, so a replay is a duplicate. */ inputId?: string + principal?: Principal }): Promise { - const admission = await this.accept(inputId, { op: 'config', key, value }) + const admission = await this.accept( + inputId, + { op: 'config', key, value }, + principal, + ) if (admission !== 'new') return this.duplicateReceipt(inputId, admission) const entry = this.sessionPlugins?.config.get(key) if (!entry) return this.reject(inputId, 'unknown_config') @@ -311,7 +349,7 @@ export abstract class SessionState< this.configValues.set(key, checked) await this.persistence.stores.metadata?.set( 'harness:config', - this.threadId, + this.storeThreadId, Object.fromEntries(this.configValues), ) await this.applied(inputId, 'session') @@ -364,7 +402,11 @@ export abstract class SessionState< this.reject(inputId, reason) return this.keep({ inputId, status: 'rejected', reason }) } - await this.persistence.stores.metadata?.set(SETTINGS, this.threadId, next) + await this.persistence.stores.metadata?.set( + SETTINGS, + this.storeThreadId, + next, + ) const movedFrom = this.threadSettings.cwd this.threadSettings = next // Tell the model about a new folder, so it stops using the old paths. @@ -420,7 +462,7 @@ export abstract class SessionState< SessionIndexEntry> | undefined > { // The schema checked the metadata on each write. - return (await this.index.get(this.threadId)) as + return (await this.index.get(this.indexKey)) as | SessionIndexEntry> | undefined } @@ -435,7 +477,7 @@ export abstract class SessionState< async updateEntry( update: SessionEntryUpdate>, ): Promise> | undefined> { - return (await this.index.update(this.threadId, (entry) => + return (await this.index.update(this.indexKey, (entry) => changedEntry( entry, update as SessionEntryUpdate>, @@ -573,7 +615,7 @@ export abstract class SessionState< } else { await this.persistence.stores.metadata?.set( 'harness:revert', - this.threadId, + this.storeThreadId, revert ?? null, ) } @@ -592,13 +634,13 @@ export abstract class SessionState< let after = 0 for (;;) { const entries = await this.log.store.read({ - logId: this.logId, + logId: this.storeLogId, after, limit: 256, }) for (const { seq, record } of entries) { after = seq - const isOwn = (record.thread ?? this.logId) === this.threadId + const isOwn = threadOf(record, this.logId) === this.threadId if (isOwn && record.type === type) found.push(record) } // A store can return fewer than `limit` before the end: stop only at @@ -621,16 +663,21 @@ export abstract class SessionState< } satisfies SessionDescription } - /** Answer a question from `ctx.session.ask`. */ + /** + * Answer a question from `ctx.session.ask`. `principal` is who answers. + * Default: the principal that opened the session. The log keeps it. + */ async answer({ questionId, value, inputId = createInputId(), + principal, }: { questionId: string value: unknown /** @internal The id of a forwarded input, so a replay is a duplicate. */ inputId?: string + principal?: Principal }): Promise { // A secret (a key or a password) stays out of the inbox and the log. A // restart does not replay an answer, so they never need it. A host that @@ -638,11 +685,15 @@ export abstract class SessionState< // entry of the answer after it applies it. const isSecret = this.questions.get(questionId)?.secret === true // A host that does not have the question passes the value on to its owner. - const admission = await this.accept(inputId, { - op: 'answer', - questionId, - value: isSecret ? '[secret]' : value, - }) + const admission = await this.accept( + inputId, + { + op: 'answer', + questionId, + value: isSecret ? '[secret]' : value, + }, + principal, + ) if (admission !== 'new') return this.duplicateReceipt(inputId, admission) const question = this.questions.get(questionId) if (!question) return this.reject(inputId, 'unknown_question') @@ -683,7 +734,7 @@ export abstract class SessionState< protected async loadConfig(): Promise { const stored = await this.persistence.stores.metadata?.get( 'harness:config', - this.threadId, + this.storeThreadId, ) if (typeof stored !== 'object' || stored === null) return for (const [key, value] of Object.entries(stored)) { @@ -701,7 +752,7 @@ export abstract class SessionState< protected async loadSettings(): Promise { const stored = await this.persistence.stores.metadata?.get( SETTINGS, - this.threadId, + this.storeThreadId, ) if (!isRecord(stored)) return const names = this.settingNames() @@ -723,7 +774,7 @@ export abstract class SessionState< protected async loadUsage(): Promise { const stored = await this.persistence.stores.metadata?.get( USAGE, - this.threadId, + this.storeThreadId, ) if (isSessionUsage(stored)) this.usageTotals = stored } @@ -732,7 +783,7 @@ export abstract class SessionState< protected async loadInterrupted(): Promise { const stored = await this.persistence.stores.metadata?.get( INTERRUPTED, - this.threadId, + this.storeThreadId, ) if (isInterruptedTurn(stored)) this.interrupted = stored await this.refreshInterrupted() @@ -822,14 +873,14 @@ export abstract class SessionState< let after = 0 for (;;) { const entries = await this.log.store.read({ - logId: this.logId, + logId: this.storeLogId, after, limit: 256, }) for (const { seq, record } of entries) { after = seq if ( - (record.thread ?? this.logId) !== this.threadId || + threadOf(record, this.logId) !== this.threadId || record.type !== 'harness.event' || !isRecord(record.event) ) @@ -920,20 +971,20 @@ export abstract class SessionState< if (next) await this.persistence.stores.metadata?.set( INTERRUPTED, - this.threadId, + this.storeThreadId, next, ) else if (previous !== undefined) { try { await this.persistence.stores.metadata?.delete( INTERRUPTED, - this.threadId, + this.storeThreadId, ) } catch (error) { if (!consumed || observed !== undefined) throw error await this.persistence.stores.metadata?.set( INTERRUPTED, - this.threadId, + this.storeThreadId, null, ) } @@ -998,7 +1049,7 @@ export abstract class SessionState< stopped: InterruptedTurn, ): Promise { await this.persistence.stores.metadata - ?.set(INTERRUPTED, this.threadId, stopped) + ?.set(INTERRUPTED, this.storeThreadId, stopped) .catch((error: unknown) => this.warn(operation, 'harness:interrupts', error), ) diff --git a/packages/ai-harness/src/session/turn-parts.ts b/packages/ai-harness/src/session/turn-parts.ts index 94aae4f3e3..5114a9a95c 100644 --- a/packages/ai-harness/src/session/turn-parts.ts +++ b/packages/ai-harness/src/session/turn-parts.ts @@ -35,6 +35,7 @@ import type { ChatTurnResult, MediaRecord, Principal, + TurnInfo, UserInput, } from '../types' import type { QueuedTurn, AgentRunImpl, AgentChain } from './types' @@ -70,19 +71,21 @@ export abstract class SessionTurnParts< } /** - * Ask each plugin for the tools it found since the last turn. A plugin - * that fails (for example an MCP server that is down) is skipped, and - * clients get a `harness.plugin.warning` event. + * Ask each plugin for the tools it found since the last turn, for `turn`. + * A plugin that fails (for example an MCP server that is down) is skipped, + * and clients get a `harness.plugin.warning` event. */ protected async discoverTools( discoverers: MountedPlugins['discoverers'], taken: Set, operation: OperationImpl, + turn: TurnInfo, ): Promise> { const found: Array = [] for (const { discover, owner } of discoverers) { try { - for (const tool of await discover()) { + const tools = await discover(turn) + for (const tool of tools) { if (taken.has(tool.name)) continue taken.add(tool.name) found.push(tool) @@ -138,9 +141,13 @@ export abstract class SessionTurnParts< /** * Middleware that gives each tool call of a turn `threadId`, `runId`, and * `detach` in its context. `detach` lets `background()` move the call to - * the background. A tool that calls `detach` supports it. + * the background. A tool that calls `detach` supports it. The wake turn of + * a detached job runs as `principal`, the sender of the turn. */ - protected detachTools(runId: string): AnyChatMiddleware { + protected detachTools( + runId: string, + principal: Principal | undefined, + ): AnyChatMiddleware { // Each model call gets the tools of the last one. Wrap each tool once. const wrapped = new WeakSet() const wrap = (tool: AnyTool): AnyTool => { @@ -153,8 +160,6 @@ export abstract class SessionTurnParts< if (!context || id === undefined) return execute(args, context) context.threadId = this.threadId context.runId = runId - // The wake turn runs as the sender of the turn that detached the job. - const principal = this.sender() context.detach = (work) => new Promise((moved) => { this.detachable.set(id, () => { @@ -261,7 +266,7 @@ export abstract class SessionTurnParts< if (metadata) { const totals = structuredClone(this.usageTotals) this.usageSaved = this.usageSaved.then(() => - metadata.set(USAGE, this.threadId, totals), + metadata.set(USAGE, this.storeThreadId, totals), ) await this.usageSaved } @@ -593,7 +598,7 @@ export abstract class SessionTurnParts< const ttlMs = this.lease?.ttlMs ?? LEASE.ttlMs const renewMs = this.lease?.renewMs ?? LEASE.renewMs const lease = () => ({ - threadId: this.threadId, + threadId: this.storeThreadId, inputId, operationId: operation.id, attempt, diff --git a/packages/ai-harness/src/session/turn.ts b/packages/ai-harness/src/session/turn.ts index 8831d5ab46..93490427fa 100644 --- a/packages/ai-harness/src/session/turn.ts +++ b/packages/ai-harness/src/session/turn.ts @@ -23,6 +23,7 @@ import { toolCallOf, storedPrincipal, turnContext, + withRun, withoutPlugins, settingTools, resetMarker, @@ -153,14 +154,19 @@ export abstract class SessionTurn< operation.setStatus('running') await this.applied(inputId, operation.id) this.publishStarted(operation) - const credentials = this.credentialsOf(principal) + const run = this.runContext({ + principal, + runId: operation.id, + signal: operation.abortController.signal, + }) try { const result: unknown = await entry.command.run(checked, { signal: operation.abortController.signal, // The command acts for its sender, also while another user's turn runs. session: this.pluginApi(principal, undefined, operation.id), ...(principal ? { principal } : {}), - credentials, + credentials: run.credentials, + run, }) operation.publish( customEvent('harness.command.result', { @@ -364,7 +370,19 @@ export abstract class SessionTurn< // What plugins and the router see of this turn. A turn that recovery // runs again gets the message its input sent. const sentMessage = turn.message ?? turn.sentMessage - const sender = this.sender() + const sender = turn.principal ?? this.principal + const runParent = turn.answers?.runId ?? turn.parentRunId + // The run of this turn: its tools, its middleware, its subagents, and + // its turn plugins get the credentials and keys of its sender. A sign-in + // wait of one of its tools stops this turn only. + const run = this.runContext({ + principal: sender, + runId: operation.id, + ...(runParent ? { parentRunId: runParent } : {}), + signal: operation.abortController.signal, + canWait: (id) => + this.activeTurn === operation && !this.declinedSignIns.has(id), + }) const info: TurnInfo = { operationId: operation.id, ...(turn.inputId ? { inputId: turn.inputId } : {}), @@ -372,6 +390,7 @@ export abstract class SessionTurn< ...(turn.context !== undefined ? { context: turn.context } : {}), ...(sender ? { principal: sender } : {}), ...(turn.overrides ? { overrides: turn.overrides } : {}), + run, } try { await this.flushNotes() @@ -404,7 +423,8 @@ export abstract class SessionTurn< inheritedExtensions: this.sessionPlugins.extensions, } : {}), - services: this.services, + // Bound to this turn's run, not read when a service is called. + services: this.pluginServices(info), }) } // A removed session plugin stays mounted, but gives this turn nothing. @@ -458,9 +478,9 @@ export abstract class SessionTurn< ? new ToolTurnAdapter(turn.tool, `tool-${turn.inputId ?? operation.id}`) : (overrides?.adapter ?? named ?? picked ?? this.harness.adapter) if (!model) throw new Error(NO_MODEL) - const adapter: AnyTextAdapter = await this.keys.adapter(model) - const resolvePrompt = (prompt: string | (() => string)) => - typeof prompt === 'function' ? prompt() : prompt + const adapter: AnyTextAdapter = await run.keys.adapter(model) + const resolvePrompt = (prompt: string | ((turn: TurnInfo) => string)) => + typeof prompt === 'function' ? prompt(info) : prompt const turnTools = overrides?.tools ?? [] const turnToolNames = new Set(turnTools.map((tool) => tool.name)) const staticTools = [ @@ -473,6 +493,7 @@ export abstract class SessionTurn< [...(session?.discoverers ?? []), ...(runPlugins?.discoverers ?? [])], new Set(staticTools.map((tool) => tool.name)), operation, + info, ) // Before the chat() options below: prompts may describe these tools. // The `tools` setting picks from them first, so a preparer (code mode) @@ -502,7 +523,11 @@ export abstract class SessionTurn< keptTools.length > 0 ? keepTurnTools(keptTools) : undefined const reasoning = overrides?.reasoning ?? stored.reasoning ?? this.harness.reasoning - const context = turnContext(this.harness.context, turn.context) + // The server's run wins over a `run` key of the client's context. + const context = withRun( + turnContext(this.harness.context, turn.context), + run, + ) // The turn value wins over the session value, field by field. const turnCache = cacheObject(overrides?.promptCache) const promptCache = { @@ -514,7 +539,7 @@ export abstract class SessionTurn< ? await this.userMessage(turn.message, turn.wake) : undefined let resume = turn.resume - let parentRunId = turn.answers?.runId ?? turn.parentRunId + let parentRunId = runParent const { chatPersistence, checkpoint } = this if (!chatPersistence || !checkpoint) { throw new Error('The session is not open yet.') @@ -654,7 +679,8 @@ export abstract class SessionTurn< middleware: [ ...bridges, detachOnShutdown, - withPersistence(chatPersistence), + // The index entries of subagent children go in this tenant. + withPersistence(chatPersistence, { tenantId: this.tenantId }), checkpoint, // The same text as the recovery repair. Early, so a later // middleware can still replace the result of a cut call. @@ -673,7 +699,7 @@ export abstract class SessionTurn< ...(runPlugins?.middleware ?? []), ...(keepTools ? [keepTools] : []), ...(durableTools ? [durableTools] : []), - this.detachTools(operation.id), + this.detachTools(operation.id, sender), ...(rootBag ? [] : [this.syncFold()]), // Again: a middleware that returns `messages` (as `syncFold` // does) resets what the model gets to the whole transcript. @@ -705,6 +731,7 @@ export abstract class SessionTurn< operation, captured, sender, + run, runPlugins, bag.agents, ), @@ -722,7 +749,7 @@ export abstract class SessionTurn< ? { interrupts: this.harness.interrupts } : {}), // The context type narrows per harness; chat() wants it exact. - ...(context !== undefined ? { context: context as never } : {}), + context: context as never, ...(reasoning !== undefined ? { reasoning } : {}), promptCache, threadId: this.threadId, @@ -760,7 +787,7 @@ export abstract class SessionTurn< runId: operation.id, interrupts, ...(messages ? { routed: { messages, root: isRootPart } } : {}), - ...storedPrincipal(this.sender()), + ...storedPrincipal(sender), ...(turn.context !== undefined ? { context: turn.context } : {}), @@ -999,7 +1026,7 @@ export abstract class SessionTurn< const { total } = this.usage() await this.index .update( - this.threadId, + this.indexKey, (entry) => entry && { ...entry, diff --git a/packages/ai-harness/src/session/types.ts b/packages/ai-harness/src/session/types.ts index 2919a7475b..6234341112 100644 --- a/packages/ai-harness/src/session/types.ts +++ b/packages/ai-harness/src/session/types.ts @@ -72,8 +72,8 @@ export interface AgentStartOptions extends AgentRunOptions { resume?: boolean /** * Who starts the agent. The run, its credentials, and its wake turn are - * theirs. Default: the sender of the running turn, else the principal that - * opened the session. + * theirs. Default: the principal that opened the session. In a tool, pass + * the user of the tool call: `ctx.context.run.principal`. */ principal?: Principal } @@ -215,6 +215,8 @@ export interface SessionDependencies { generationRuns: GenerationRunStore }> principal?: Principal + /** The tenant of the session. Without it, the default tenant. */ + tenantId?: string /** Identifies this host on run leases. */ hostId: string /** Writes the session index entry of the thread. */ @@ -229,14 +231,26 @@ export interface SessionDependencies { onFailure: (error: unknown) => void, ) => Promise } - /** The log of the session. Default: the thread id. */ + /** + * The log of the session: the `logId` of `host.open`, else the thread id. + * The records of that log with no `thread` field belong to the thread + * with this id. + */ logId: string + /** The id of that log in the store, with the tenant. See `logIdOf`. */ + storeLogId: string lease?: LeaseOptions /** The `promptCache` of `host.open()`. It overrides the harness value. */ promptCache?: PromptCacheOptions /** Gets each event the session publishes, after the feed has it. */ onEvent: (event: StreamChunk) => void onClose: () => void + /** + * The host checks access (`access` on the host). Then an entry with no + * owner stays without one, and provider keys read env vars only with + * `envKeys`. + */ + access?: { envKeys: boolean } } export interface QueuedTurn { diff --git a/packages/ai-harness/src/tenant-stores.ts b/packages/ai-harness/src/tenant-stores.ts new file mode 100644 index 0000000000..b2a7dc18d0 --- /dev/null +++ b/packages/ai-harness/src/tenant-stores.ts @@ -0,0 +1,98 @@ +import { logIdOf, parseLogId } from './log' +import type { + InterruptStore, + MessageStore, + RunStore, +} from '@tanstack/ai-persistence' + +/** The stores that `withPersistence` keys by the thread id of `chat()`. */ +interface ChatThreadStores { + messages?: MessageStore + runs?: RunStore + interrupts?: InterruptStore +} + +/** + * `stores` in the tenant `tenantId`. `withPersistence`, the checkpoints, and + * the session read and write these stores with the thread id of the app, + * which is also the thread id on the events. Here each thread id becomes its + * store key, `logIdOf({ tenantId, logId: threadId })`, and each record that + * comes back has the thread id of the app again. A record of another tenant + * reads as not found. So one thread id in two tenants is two threads. + * + * Only the stores in `stores` come back. + */ +export function tenantStores( + stores: ChatThreadStores, + tenantId: string | undefined, +) { + const { messages, runs, interrupts } = stores + const keyOf = (threadId: string) => logIdOf({ tenantId, logId: threadId }) + /** `record` with the thread id of the app, or `null` in another tenant. */ + const own = (record: T | null) => { + if (record === null) return null + const key = parseLogId(record.threadId) + return key.tenantId === tenantId ? { ...record, threadId: key.logId } : null + } + const ownAll = ( + records: ReadonlyArray | undefined, + ) => (records ?? []).flatMap((record) => own(record) ?? []) + + const messageView: MessageStore | undefined = messages && { + loadThread: (threadId: string) => messages.loadThread(keyOf(threadId)), + saveThread: (threadId, list) => messages.saveThread(keyOf(threadId), list), + } + const runView: RunStore | undefined = runs && { + createOrResume: async (input) => { + const record = own( + await runs.createOrResume({ + ...input, + threadId: keyOf(input.threadId), + }), + ) + // A run id is one key for all tenants: a run of another tenant does + // not go on here. + if (!record) throw new Error(`Run ${input.runId} is in another tenant.`) + return record + }, + update: (runId, patch) => runs.update(runId, patch), + get: async (runId) => own(await runs.get(runId)), + findActiveRun: async (threadId) => + own(await runs.findActiveRun(keyOf(threadId))), + listByThread: + runs.listByThread && + (async (threadId) => ownAll(await runs.listByThread?.(keyOf(threadId)))), + listByParentRun: + runs.listByParentRun && + (async (parentRunId) => + ownAll(await runs.listByParentRun?.(parentRunId))), + listReclaimable: + runs.listReclaimable && + (async (options) => ownAll(await runs.listReclaimable?.(options))), + } + const interruptView: InterruptStore | undefined = interrupts && { + create: (record) => + interrupts.create({ ...record, threadId: keyOf(record.threadId) }), + resolve: (interruptId, response) => + interrupts.resolve(interruptId, response), + cancel: (interruptId) => interrupts.cancel(interruptId), + commitBatch: + interrupts.commitBatch && + (async (entries) => { + await interrupts.commitBatch?.(entries) + }), + get: async (interruptId) => own(await interrupts.get(interruptId)), + list: async (threadId) => ownAll(await interrupts.list(keyOf(threadId))), + listPending: async (threadId) => + ownAll(await interrupts.listPending(keyOf(threadId))), + listByRun: async (runId) => ownAll(await interrupts.listByRun(runId)), + listPendingByRun: async (runId) => + ownAll(await interrupts.listPendingByRun(runId)), + } + const view: ChatThreadStores = { + ...(messageView && { messages: messageView }), + ...(runView && { runs: runView }), + ...(interruptView && { interrupts: interruptView }), + } + return view +} diff --git a/packages/ai-harness/src/types.ts b/packages/ai-harness/src/types.ts index 664534efc5..5befd077f7 100644 --- a/packages/ai-harness/src/types.ts +++ b/packages/ai-harness/src/types.ts @@ -6,11 +6,13 @@ import type { KeyedAdapter, ModelMessage, PromptCacheOptions, + ProviderKeys, ReasoningOption, RunAgentResumeItem, StreamChunk, } from '@tanstack/ai' import type { LogRecord } from '@tanstack/ai-persistence' +import type { CredentialsAccess } from './auth' import type { UsageCounts } from './usage' /** A user message: plain text, or content parts (text, images, files). */ @@ -414,6 +416,57 @@ export interface Principal { tenantId?: string } +/** + * One run of a session: a chat turn, an agent run, or a command. Each run has + * its own user, so a tool of one run never gets the credentials of another + * run that runs at the same time. + * + * Tools and middleware read it as `ctx.context.run`. Type them with + * {@link HarnessRunContext}. A turn plugin reads it as `ctx.turn.run`, and a + * command as `ctx.run`. + * + * @example + * ```ts + * const listIssues = toolDefinition({ name: 'list_issues', description: 'List issues' }) + * .server(async (_args, ctx) => { + * const credential = await ctx.context.run.credentials.require('github') + * return fetchIssues(credential) + * }) + * ``` + */ +export interface RunContext { + /** + * The tenant of the session: the tenant of the principal that opened it. + * `undefined` in the default tenant. + */ + tenantId?: string + /** The user who started this run. `undefined` in a session with no user. */ + principal: { id: string; tenantId?: string } | undefined + /** + * The id of the log of the session in the `LogStore`: the tenant and the + * `logId` of `host.open`, from `logIdOf`. + */ + logId: string + threadId: string + /** The operation id of this run. */ + runId: string + /** The run that this run continues or that started it, when there is one. */ + parentRunId?: string + /** Aborts when this run stops. */ + signal: AbortSignal + /** The credentials of the user of this run. */ + credentials: CredentialsAccess + /** The model provider keys of the user of this run. */ + keys: ProviderKeys +} + +/** + * The part of the runtime context that the harness adds to every chat call + * of a run. Use it as the context type of a tool or a middleware: + * `.server(...)`, `ChatMiddleware`. + */ +export type HarnessRunContext = { run: RunContext } + /** * A chat turn as plugins and the router see it: the input that started it, * and who sent it. @@ -431,6 +484,8 @@ export interface TurnInfo { principal?: Principal /** The settings of this turn only. Not kept after a restart. */ overrides?: TurnOverrides + /** The run of this turn: its user, and that user's credentials and keys. */ + run: RunContext } /** What a chat turn operation resolves to. */ diff --git a/packages/ai-harness/tests/access.test.ts b/packages/ai-harness/tests/access.test.ts new file mode 100644 index 0000000000..54d688c47a --- /dev/null +++ b/packages/ai-harness/tests/access.test.ts @@ -0,0 +1,573 @@ +import { afterEach, describe, expect, it, vi } from 'vitest' +import { z } from 'zod' +import { toolDefinition } from '@tanstack/ai' +import { defineByokProvider, keyedAdapter } from '@tanstack/ai/byok' +import { memoryLogStore, memoryPersistence } from '@tanstack/ai-persistence' +import { + AccessDeniedError, + createHarnessHost, + defineHarness, + definePlugin, + membersOnly, +} from '../src' +import { after, gate, mockAdapter, text, toolCall } from './helpers' +import type { AnyTextAdapter } from '@tanstack/ai' +import type { + Authorize, + HarnessHost, + HarnessRunContext, + HostEvent, +} from '../src' +import type { Reply } from './helpers' + +afterEach(() => { + vi.unstubAllEnvs() +}) + +interface User { + id: string + tenantId: string + role: 'user' | 'admin' +} + +const olga: User = { id: 'olga', tenantId: 'acme', role: 'user' } +const bea: User = { id: 'bea', tenantId: 'acme', role: 'user' } +const vic: User = { id: 'vic', tenantId: 'acme', role: 'user' } +const cid: User = { id: 'cid', tenantId: 'acme', role: 'user' } +const gus: User = { id: 'gus', tenantId: 'globex', role: 'user' } + +const THREAD = 't1' + +/** A host with access, and a harness whose turns are `replies`. */ +function accessHost({ + replies = () => text('ok'), + authorize = membersOnly(), + envKeys, + adapter, + durable = false, +}: { + replies?: Array | Reply + authorize?: Authorize + envKeys?: boolean + adapter?: AnyTextAdapter | ReturnType['adapter'] + /** A log store: the log id of each session has its tenant. */ + durable?: boolean +} = {}) { + const memory = memoryPersistence() + const { runs, metadata, sessions } = memory.stores + const persistence = durable + ? { stores: { log: memoryLogStore(), runs, metadata, sessions } } + : memory + const host = createHarnessHost({ + persistence, + access: { + identify: (): User | null => null, + authorize, + ...(envKeys !== undefined ? { envKeys } : {}), + }, + }) + const harness = defineHarness({ + name: 'test/access', + adapter: adapter ?? mockAdapter(replies).adapter, + }) + const open = (principal: User, threadId = THREAD) => + host.open(harness, { threadId, principal }) + return { host, harness, sessions, open } +} + +/** A tool that keeps the user of each run that calls it. */ +function whoami() { + const seen: Array = [] + const tool = toolDefinition({ + name: 'whoami', + description: 'Who runs this?', + inputSchema: z.object({}), + }) + const plugin = definePlugin({ + name: 'test/whoami', + setup: () => ({ + tools: [ + tool.server((_input, { context }) => { + seen.push(context.run.principal?.id) + return Promise.resolve({}) + }), + ], + }), + }) + return { plugin, seen } +} + +const acme = defineByokProvider({ + id: 'acme', + label: 'Acme', + env: 'ACME_TEST_KEY', +}) + +/** A keyed adapter for `acme`. It keeps each key it got. */ +function acmeAdapter() { + const keys: Array = [] + const adapter = keyedAdapter(acme, (key) => { + keys.push(key) + return mockAdapter(() => text('keyed answer')).adapter + }) + return { adapter, keys } +} + +/** Read `host.events()` into a list until `stop` runs. */ +function readEvents(host: Pick, 'events'>) { + const events: Array = [] + const reader = new AbortController() + const reading = (async () => { + for await (const event of host.events({ signal: reader.signal })) { + events.push(event) + } + })() + const stop = async () => { + reader.abort() + await reading + } + return { events, stop } +} + +describe('host access with membersOnly()', () => { + it('lets a member that the owner shared with send, and refuses a stranger', async () => { + const { host, harness, open } = accessHost({ + replies: [() => text('one'), () => text('two')], + }) + const owner = await open(olga) + await owner.send('Hi.') + await host.sessions.share({ + threadId: THREAD, + userId: 'bea', + role: 'member', + principal: olga, + }) + + const member = await open(bea) + + expect(await member.send('Me too.')).toEqual({ text: 'two', media: [] }) + await expect(open(cid)).rejects.toBeInstanceOf(AccessDeniedError) + await expect( + host.sessions.get({ threadId: THREAD, principal: cid }), + ).rejects.toBeInstanceOf(AccessDeniedError) + const listed = async (principal: User) => + (await host.sessions.list({ principal })).entries.map( + (entry) => entry.threadId, + ) + expect(await listed(bea)).toEqual([THREAD]) + expect(await listed(cid)).toEqual([]) + expect(host.isOpen(harness, { threadId: THREAD, principal: bea })).toBe( + true, + ) + await host.close() + }) + + it('opens a thread id of another tenant as a new session of the own tenant', async () => { + const { host, open } = accessHost({ durable: true }) + const owner = await open(olga) + await owner.send('The plan of acme.') + + const other = await open(gus) + + expect(await other.transcript()).toEqual([]) + expect( + await host.sessions.get({ threadId: THREAD, principal: gus }), + ).toMatchObject({ + tenantId: 'globex', + ownerId: 'gus', + members: [{ userId: 'gus', role: 'owner' }], + }) + expect( + await host.sessions.get({ threadId: THREAD, principal: olga }), + ).toMatchObject({ tenantId: 'acme', ownerId: 'olga' }) + await host.close() + }) + + it('lets a viewer read but not send, and a member send but not configure', async () => { + const { host, open } = accessHost() + await open(olga) + for (const [userId, role] of [ + ['bea', 'member'], + ['vic', 'viewer'], + ] as const) { + await host.sessions.share({ + threadId: THREAD, + userId, + role, + principal: olga, + }) + } + const viewer = await open(vic) + const member = await open(bea) + + expect(await viewer.transcript()).toEqual([]) + await expect(viewer.send('Hi.')).rejects.toBeInstanceOf(AccessDeniedError) + await expect( + member.configure({ settings: { instructions: 'Be short.' } }), + ).rejects.toBeInstanceOf(AccessDeniedError) + expect(member.settings()).toEqual({}) + await host.close() + }) + + it('lets a member cancel only a run that the member started', async () => { + const release = gate() + const { host, open } = accessHost({ + replies: [after(release.opened, 'owner done'), () => text('never')], + }) + const owner = await open(olga) + await host.sessions.share({ + threadId: THREAD, + userId: 'bea', + role: 'member', + principal: olga, + }) + const member = await open(bea) + const ownerTurn = owner.send('A long task.') + await ownerTurn.receipt + const memberTurn = member.send('After that.') + expect((await memberTurn.receipt).status).toBe('queued') + + await expect( + member.cancel({ operationId: ownerTurn.id }), + ).rejects.toBeInstanceOf(AccessDeniedError) + expect((await member.cancel({ operationId: memberTurn.id })).status).toBe( + 'accepted', + ) + + await expect(memberTurn).rejects.toThrow('Cancelled before it started.') + release.open() + expect(await ownerTurn).toEqual({ text: 'owner done', media: [] }) + await host.close() + }) + + it('removes a member with unshare: their queued turn stops, and their handle refuses every call', async () => { + const release = gate() + const { host, open } = accessHost({ + replies: [after(release.opened, 'owner done'), () => text('never')], + }) + const owner = await open(olga) + await host.sessions.share({ + threadId: THREAD, + userId: 'bea', + role: 'member', + principal: olga, + }) + const member = await open(bea) + const ownerTurn = owner.send('A long task.') + await ownerTurn.receipt + const memberTurn = member.send('After that.') + expect((await memberTurn.receipt).status).toBe('queued') + + const changed = await host.sessions.unshare({ + threadId: THREAD, + userId: 'bea', + principal: olga, + }) + + expect(changed?.members).toEqual([{ userId: 'olga', role: 'owner' }]) + await expect(memberTurn).rejects.toThrow('Cancelled before it started.') + expect(() => member.send('Again.')).toThrow(AccessDeniedError) + expect(() => member.snapshot()).toThrow(AccessDeniedError) + await expect(open(bea)).rejects.toBeInstanceOf(AccessDeniedError) + await expect( + host.sessions.unshare({ + threadId: THREAD, + userId: 'olga', + principal: olga, + }), + ).rejects.toThrow('the owner cannot be removed') + release.open() + expect(await ownerTurn).toEqual({ text: 'owner done', media: [] }) + await host.close() + }) + + it("forks for the forker alone by default, and with the source's members for 'all'", async () => { + const { host, harness, open } = accessHost({ + replies: () => text('answer'), + }) + const owner = await open(olga) + await owner.send('Hi.') + for (const [userId, role] of [ + ['bea', 'member'], + ['vic', 'viewer'], + ] as const) { + await host.sessions.share({ + threadId: THREAD, + userId, + role, + principal: olga, + }) + } + const last = (await owner.transcript()).at(-1)?.id ?? '' + const fork = (members?: 'self' | 'all') => + host.sessions.fork({ + harness, + threadId: THREAD, + at: { through: last }, + ...(members ? { members } : {}), + principal: bea, + }) + + expect(await fork()).toMatchObject({ + ownerId: 'bea', + members: [{ userId: 'bea', role: 'owner' }], + }) + expect(await fork('all')).toMatchObject({ + ownerId: 'bea', + members: [ + { userId: 'bea', role: 'owner' }, + { userId: 'olga', role: 'member' }, + { userId: 'vic', role: 'viewer' }, + ], + }) + await host.close() + }) + + it("forks with the source's members by default under membersOnly({ fork: { members: 'all' } })", async () => { + const { host, harness, open } = accessHost({ + authorize: membersOnly({ fork: { members: 'all' } }), + }) + const owner = await open(olga) + await owner.send('Hi.') + await host.sessions.share({ + threadId: THREAD, + userId: 'bea', + role: 'member', + principal: olga, + }) + const last = (await owner.transcript()).at(-1)?.id ?? '' + + const entry = await host.sessions.fork({ + harness, + threadId: THREAD, + at: { through: last }, + principal: olga, + }) + + expect(entry).toMatchObject({ + ownerId: 'olga', + members: [ + { userId: 'olga', role: 'owner' }, + { userId: 'bea', role: 'member' }, + ], + }) + await host.close() + }) + + it('does not give an entry with no owner to the next opener', async () => { + const { host, sessions, open } = accessHost() + await sessions.upsert({ + tenantId: 'acme', + threadId: 'orphan', + members: [], + createdAt: 1, + updatedAt: 1, + }) + + await expect(open(olga, 'orphan')).rejects.toBeInstanceOf(AccessDeniedError) + expect( + await sessions.get({ + tenantId: 'acme', + threadId: 'orphan', + }), + ).toEqual({ + tenantId: 'acme', + threadId: 'orphan', + members: [], + createdAt: 1, + updatedAt: 1, + }) + await host.close() + }) + + it('checks a child thread through its root thread', async () => { + const { host, sessions, open } = accessHost() + await open(olga) + await host.sessions.share({ + threadId: THREAD, + userId: 'bea', + role: 'viewer', + principal: olga, + }) + const child = { + tenantId: 'acme', + threadId: 'subagent:run-1', + parentThreadId: THREAD, + ownerId: 'olga', + members: [], + createdAt: 1, + updatedAt: 1, + } + await sessions.upsert(child) + + expect( + (await host.sessions.list({ parentThreadId: THREAD, principal: bea })) + .entries, + ).toEqual([child]) + expect( + await host.sessions.get({ threadId: child.threadId, principal: bea }), + ).toEqual(child) + await expect( + host.sessions.list({ parentThreadId: THREAD, principal: cid }), + ).rejects.toBeInstanceOf(AccessDeniedError) + await host.close() + }) +}) + +describe('host access with your own access tables', () => { + /** An authorize that allows the users in `allowed`, for every action. */ + function tableAccess(...users: Array) { + const allowed = new Set(users) + const authorize: Authorize = ({ principal }) => + allowed.has(principal.id) + return { allowed, authorize } + } + + it('closes the handles of a revoked user and tells host.events() readers', async () => { + const table = tableAccess('olga', 'bea') + const { host, open } = accessHost({ authorize: table.authorize }) + await open(olga) + const member = await open(bea) + const reader = readEvents(host) + + table.allowed.delete('bea') + host.revoke({ threadId: THREAD, principal: bea }) + + expect(() => member.snapshot()).toThrow(AccessDeniedError) + await vi.waitFor(() => + expect(reader.events.filter((event) => event.type === 'revoked')).toEqual( + [ + { + type: 'revoked', + tenantId: 'acme', + threadId: THREAD, + principalId: 'bea', + at: expect.any(Number), + }, + ], + ), + ) + // The app gives the access back: a new open works. + table.allowed.add('bea') + expect((await open(bea)).snapshot().threadId).toBe(THREAD) + await reader.stop() + await host.close() + }) + + it('refuses share without membersOnly()', async () => { + const { host, open } = accessHost({ + authorize: tableAccess('olga').authorize, + }) + await open(olga) + + await expect( + host.sessions.share({ + threadId: THREAD, + userId: 'bea', + role: 'member', + principal: olga, + }), + ).rejects.toThrow( + 'host.sessions.share needs membersOnly() as the authorize of the host access. With your own access tables, change them and call host.revoke.', + ) + await host.close() + }) +}) + +describe('host access and principals', () => { + it('needs a principal on each public host call', async () => { + const { host, harness } = accessHost() + + await expect(host.open(harness, { threadId: THREAD })).rejects.toThrow( + 'host.open needs a principal: the host has access.', + ) + await expect(host.sessions.list()).rejects.toThrow( + 'host.sessions.list needs a principal: the host has access.', + ) + expect(() => host.isOpen(harness, { threadId: THREAD })).toThrow( + 'host.isOpen needs a principal: the host has access.', + ) + await host.close() + }) + + it('refuses a call that acts as another principal than the handle', async () => { + const { host, open } = accessHost() + const owner = await open(olga) + + expect(() => owner.send('Hi.', { principal: bea })).toThrow( + 'This handle acts as "olga". Open another handle for "bea".', + ) + await host.close() + }) + + it('runs the send of a second opener as that opener, without access', async () => { + const { plugin, seen } = whoami() + const host = createHarnessHost({ persistence: memoryPersistence() }) + const harness = defineHarness({ + name: 'test/two-openers', + adapter: mockAdapter([ + () => toolCall('whoami', {}, 'c1'), + () => text('ok'), + () => toolCall('whoami', {}, 'c2'), + () => text('ok'), + ]).adapter, + plugins: () => [plugin], + }) + const alice = await host.open(harness, { + threadId: THREAD, + principal: { id: 'alice' }, + }) + const bob = await host.open(harness, { + threadId: THREAD, + principal: { id: 'bob' }, + }) + + await bob.send('Who am I?') + await alice.send('And me?') + + expect(seen).toEqual(['bob', 'alice']) + await host.close() + }) +}) + +describe('provider keys with host access', () => { + it('does not read env vars without envKeys', async () => { + vi.stubEnv('ACME_TEST_KEY', 'sk-acme-env-0000') + const keyed = acmeAdapter() + const { host, open } = accessHost({ adapter: keyed.adapter }) + const owner = await open(olga) + + await expect(owner.send('Hi.')).rejects.toThrow( + 'Sign in to acme first. Run /connect acme.', + ) + expect(keyed.keys).toEqual([]) + await host.close() + }) + + it('reads env vars with envKeys: true', async () => { + vi.stubEnv('ACME_TEST_KEY', 'sk-acme-env-0000') + const keyed = acmeAdapter() + const { host, open } = accessHost({ adapter: keyed.adapter, envKeys: true }) + const owner = await open(olga) + + expect(await owner.send('Hi.')).toEqual({ + text: 'keyed answer', + media: [], + }) + expect(keyed.keys).toEqual(['sk-acme-env-0000']) + await host.close() + }) +}) + +it('needs stores.sessions for host access', () => { + expect(() => + createHarnessHost({ + persistence: { + stores: { messages: memoryPersistence().stores.messages }, + }, + access: { identify: () => null, authorize: membersOnly() }, + }), + ).toThrow( + 'Host access needs stores.sessions: each session keeps its owner and members there.', + ) +}) diff --git a/packages/ai-harness/tests/agent-resume.test.ts b/packages/ai-harness/tests/agent-resume.test.ts index aea33966f6..9dfe7fedb6 100644 --- a/packages/ai-harness/tests/agent-resume.test.ts +++ b/packages/ai-harness/tests/agent-resume.test.ts @@ -2,7 +2,7 @@ import { describe, expect, it, vi } from 'vitest' import { z } from 'zod' import { defineAgent, toolDefinition } from '@tanstack/ai' import { memoryLogStore, memoryPersistence } from '@tanstack/ai-persistence' -import { createHarnessHost, defineHarness } from '../src' +import { createHarnessHost, defineHarness, logIdOf } from '../src' import { messageTexts, mockAdapter, @@ -14,6 +14,8 @@ import type { AnyTextAdapter } from '@tanstack/ai' import type { HarnessDurability } from '../src' const THREAD = 't1' +/** The id of the log of `THREAD` in the store. */ +const LOG = logIdOf({ logId: THREAD }) const alice = { id: 'alice' } /** @@ -91,7 +93,7 @@ async function stopDuringAgent(options: { refuse.mockRestore() await stores.runs.update(operation.id, { leaseExpiresAt: Date.now() - 1000 }) - const record = (await log.read({ logId: THREAD })) + const record = (await log.read({ logId: LOG })) .map((entry) => entry.record) .find( (entry) => diff --git a/packages/ai-harness/tests/agent-runs.test.ts b/packages/ai-harness/tests/agent-runs.test.ts index 9a71048502..134164d51f 100644 --- a/packages/ai-harness/tests/agent-runs.test.ts +++ b/packages/ai-harness/tests/agent-runs.test.ts @@ -6,7 +6,7 @@ import { memoryLogStore, memoryPersistence, } from '@tanstack/ai-persistence' -import { createHarnessHost, defineHarness } from '../src' +import { createHarnessHost, defineHarness, logIdOf } from '../src' import { after, gate, @@ -21,6 +21,8 @@ import type { AnyTextAdapter } from '@tanstack/ai' import type { AnyHarness, HarnessSession, Operation } from '../src' const THREAD = 't1' +/** The id of the log of `THREAD` in the store. */ +const LOG = logIdOf({ logId: THREAD }) const alice = { id: 'alice' } /** Open `harness` for alice on a host without a log. */ @@ -108,8 +110,9 @@ describe('the thread of an agent run', () => { await session.agents.drafter.start() expect(thread).toMatch(/^t1:drafter:in-/) + // The chain thread is in the tenant of the session: the default one. const saved = JSON.stringify( - await persistence.stores.messages.loadThread(thread), + await persistence.stores.messages.loadThread(logIdOf({ logId: thread })), ) expect(saved).toContain('Draft it.') expect(saved).toContain('Drafted.') @@ -206,9 +209,9 @@ describe('follow-ups to an agent run', () => { await vi.waitFor(() => expect(model.calls).toHaveLength(1)) const receipt = await run.send('Add a title.', { mode: 'followUp' }) const followUp = session.operation(receipt.operationId ?? '') - const [last] = (await inner.read({ logId: THREAD })).slice(-1) + const [last] = (await inner.read({ logId: LOG })).slice(-1) await inner.append({ - logId: THREAD, + logId: LOG, seq: (last?.seq ?? 0) + 1, epoch: UNOWNED_EPOCH, records: [{ type: 'app.from-another-host' }], diff --git a/packages/ai-harness/tests/auth-wait.test.ts b/packages/ai-harness/tests/auth-wait.test.ts index 0f3fba92fe..776b1ef274 100644 --- a/packages/ai-harness/tests/auth-wait.test.ts +++ b/packages/ai-harness/tests/auth-wait.test.ts @@ -9,11 +9,12 @@ import { defineCommand, defineHarness, definePlugin, + logIdOf, oauthConnector, } from '../src' import { createHarnessHandler } from '../src/server' import { after, gate, mockAdapter, text, toolCall } from './helpers' -import type { HarnessSession } from '../src' +import type { HarnessRunContext, HarnessSession } from '../src' import type { OAuthConfig } from '../src/oauth' import type { Reply } from './helpers' @@ -32,13 +33,13 @@ const oauth: OAuthConfig = { function github(seen: Array, wait: boolean) { return definePlugin({ name: 'test/github', - setup: (ctx) => ({ + setup: () => ({ tools: [ toolDefinition({ name: 'list_issues', description: 'List issues', - }).server(async () => { - const credential = await ctx.credentials.require( + }).server(async (_args, ctx) => { + const credential = await ctx.context.run.credentials.require( 'github', wait ? { wait: true } : undefined, ) @@ -61,11 +62,11 @@ function github(seen: Array, wait: boolean) { return 'Connected.' }, }), - // A save of another credential, by the session's own access. + // A save of another credential, for the user who runs the command. note: defineCommand({ description: 'Save a note credential', - run: () => - ctx.credentials.set({ + run: (_input, command) => + command.credentials.set({ id: 'note', credential: { type: 'api_key', value: 'n' }, }), @@ -124,8 +125,8 @@ describe('require with wait', () => { toolDefinition({ name: 'list_issues', description: 'List issues', - }).server(async () => { - seenTokens.push(await token({ wait: true })) + }).server(async (_args, ctx) => { + seenTokens.push(await token({ run: ctx.context.run, wait: true })) return ['#1'] }), ], @@ -396,19 +397,23 @@ describe('sign-ins in a shared thread', () => { ]) const persistence = memoryPersistence() const credentials = memoryCredentialStore() - const host = createHarnessHost({ persistence, credentials }) + const host = createHarnessHost({ + persistence, + credentials, + access: { + identify: ({ request }) => ({ + id: request.headers.get('x-user') ?? '', + }), + authorize: () => true, + }, + }) const harness = defineHarness({ name: 'test/auth-wait-shared-http', adapter, plugins: () => [github(seen, true)], expose: { commands: ['connect:github'] }, }) - const handler = createHarnessHandler({ - host, - harness, - canAccess: () => true, - authorize: (request) => ({ id: request.headers.get('x-user') ?? '' }), - }) + const handler = createHarnessHandler({ host, harness }) const control = async (user: string, input: object) => ( await handler( @@ -480,8 +485,7 @@ describe('sign-ins in a shared thread', () => { const handler = createHarnessHandler({ host, harness, - canAccess: () => true, - authorize: () => bob, + openAccess: { principal: bob }, media: { secret: 'test-secret' }, }) const session = await host.open(harness, { @@ -501,7 +505,7 @@ describe('sign-ins in a shared thread', () => { ).json() await expect(session.operation(receipt.operationId)).resolves.toBe('bob') expect(ran).toEqual(['bob']) - const commands = (await stores.log.read({ logId: 't' })) + const commands = (await stores.log.read({ logId: logIdOf({ logId: 't' }) })) .map((entry) => entry.record) .filter( (record) => @@ -534,11 +538,12 @@ describe('sign-ins in a shared thread', () => { await (await a.open()).send('List my issues.') await a.host.close() // A copy without `principal`, as an older version wrote it. - const stored = await stores.metadata.get('harness:interrupted', 't') + const key = logIdOf({ logId: 't' }) + const stored = await stores.metadata.get('harness:interrupted', key) if (typeof stored !== 'object' || stored === null) throw new Error('No stored interrupted turn.') const { principal: _principal, ...old } = stored as Record - await stores.metadata.set('harness:interrupted', 't', old) + await stores.metadata.set('harness:interrupted', key, old) const b = start([() => text('One issue.')]) const session = await b.open() diff --git a/packages/ai-harness/tests/auth.test.ts b/packages/ai-harness/tests/auth.test.ts index a51e69803d..19b19a53c9 100644 --- a/packages/ai-harness/tests/auth.test.ts +++ b/packages/ai-harness/tests/auth.test.ts @@ -8,6 +8,7 @@ import { defineCommand, defineHarness, definePlugin, + logIdOf, oauthConnector, scrubSecrets, } from '../src' @@ -20,7 +21,7 @@ import { refreshCredential, } from '../src/oauth' import { mockAdapter, text, toolCall } from './helpers' -import type { SessionEvent } from '../src' +import type { HarnessRunContext, SessionEvent } from '../src' import type { OAuthConfig } from '../src/oauth' const oauth: OAuthConfig = { @@ -199,8 +200,8 @@ describe('oauthConnector', () => { toolDefinition({ name: 'list_issues', description: 'List issues', - }).server(async () => { - seenTokens.push(await token()) + }).server(async (_args, ctx) => { + seenTokens.push(await token({ run: ctx.context.run })) return ['#1'] }), ], @@ -257,8 +258,9 @@ describe('oauthConnector', () => { expect(seenTokens).toEqual(['gh-token']) // The model never sees the token. const transcript = JSON.stringify( - await persistence.stores.messages.loadThread('t'), + await persistence.stores.messages.loadThread(logIdOf({ logId: 't' })), ) + expect(transcript).toContain('list my issues again') expect(transcript).not.toContain('gh-token') await session.command({ name: 'disconnect:github' }) diff --git a/packages/ai-harness/tests/build.test.ts b/packages/ai-harness/tests/build.test.ts index 184a803c02..e324904e23 100644 --- a/packages/ai-harness/tests/build.test.ts +++ b/packages/ai-harness/tests/build.test.ts @@ -64,15 +64,19 @@ describe('worker mode', () => { describe('harnessText({ url })', () => { it('uses a harness served over HTTP as a model', async () => { const { adapter } = mockAdapter([() => text('from far away')]) - const host = createHarnessHost({ persistence: memoryPersistence() }) + const host = createHarnessHost({ + persistence: memoryPersistence(), + access: { + identify: ({ request }) => + request.headers.get('authorization') === 'Bearer t' + ? { id: 'u' } + : null, + authorize: () => true, + }, + }) const handler = createHarnessHandler({ host, harness: defineHarness({ name: 'test/remote', adapter }), - canAccess: () => true, - authorize: (request) => - request.headers.get('authorization') === 'Bearer t' - ? { id: 'u' } - : null, }) const remote = harnessText({ url: 'http://remote.test/api', diff --git a/packages/ai-harness/tests/client-reads.test.ts b/packages/ai-harness/tests/client-reads.test.ts index 057dcb8b07..97d7601f9b 100644 --- a/packages/ai-harness/tests/client-reads.test.ts +++ b/packages/ai-harness/tests/client-reads.test.ts @@ -26,6 +26,18 @@ const asker = definePlugin({ }), }) +const user = { id: 'u' } + +/** A host where `user` may use every thread but `someone-else`. */ +const accessHost = () => + createHarnessHost({ + persistence: memoryPersistence(), + access: { + identify: () => user, + authorize: ({ session }) => session.threadId !== 'someone-else', + }, + }) + function setup( fetchWrap?: (fetch: typeof globalThis.fetch) => typeof globalThis.fetch, ) { @@ -36,13 +48,8 @@ function setup( plugins: () => [asker], expose: { config: ['tone'], commands: ['confirm'] }, }) - const host = createHarnessHost({ persistence: memoryPersistence() }) - const handler = createHarnessHandler({ - host, - harness, - authorize: () => ({ id: 'u' }), - canAccess: (_principal, threadId) => threadId !== 'someone-else', - }) + const host = accessHost() + const handler = createHarnessHandler({ host, harness }) const direct: typeof globalThis.fetch = (input, init) => handler(new Request(input, init)) const client = createHarnessClient({ @@ -51,7 +58,7 @@ function setup( fetch: fetchWrap ? fetchWrap(direct) : direct, reconnectDelayMs: 1, }) - return { host, client, handler } + return { host, harness, client, handler } } describe('client reads and actions', () => { @@ -131,7 +138,11 @@ describe('client reads and actions', () => { }) it('refuses a read without a thread id, or for a thread the user may not open', async () => { - const { host, handler } = setup() + const { host, harness, handler } = setup() + await host.open(harness, { + threadId: 'someone-else', + principal: { id: 'other' }, + }) const statusOf = async (path: string) => (await handler(new Request(`http://local/api/harness/${path}`))).status @@ -153,17 +164,17 @@ describe('client reads and actions', () => { () => toolCall('writer', { prompt: 'Take notes' }, 'call-1'), () => text('Done.'), ]) - const host = createHarnessHost({ persistence: memoryPersistence() }) - const handler = createHarnessHandler({ - host, - harness: defineHarness({ - name: 'test/client-reads-agents', - adapter, - subagents: { agents: [writer] }, - }), - authorize: () => ({ id: 'u' }), - canAccess: (_principal, threadId) => threadId !== 'someone-else', + const host = accessHost() + const harness = defineHarness({ + name: 'test/client-reads-agents', + adapter, + subagents: { agents: [writer] }, + }) + await host.open(harness, { + threadId: 'someone-else', + principal: { id: 'other' }, }) + const handler = createHarnessHandler({ host, harness }) const client = createHarnessClient({ url: 'http://local/api/harness', threadId: 'thread-a', @@ -210,13 +221,15 @@ describe('client reads and actions', () => { it('reports the connection state of the event stream', async () => { let calls = 0 - const { host, client } = setup((fetch) => async (input, init) => { + const { host, harness, client } = setup((fetch) => async (input, init) => { const url = String(input instanceof Request ? input.url : input) if (url.includes('/events') && calls++ === 0) { return new Response('down', { status: 500 }) } return fetch(input, init) }) + // With host access, a thread with no session has no events to read. + await host.open(harness, { threadId: 'thread-r', principal: user }) const states: Array = [] const reader = new AbortController() const reading = (async () => { diff --git a/packages/ai-harness/tests/client.test.ts b/packages/ai-harness/tests/client.test.ts index b378823471..a9e6de6744 100644 --- a/packages/ai-harness/tests/client.test.ts +++ b/packages/ai-harness/tests/client.test.ts @@ -2,7 +2,7 @@ import { describe, expect, it, vi } from 'vitest' import { z } from 'zod' import { defineAgent, toolDefinition } from '@tanstack/ai' import { memoryPersistence } from '@tanstack/ai-persistence' -import { createHarnessHost, defineHarness } from '../src' +import { createHarnessHost, defineHarness, logIdOf } from '../src' import { createHarnessHandler } from '../src/server' import { createHarnessClient } from '../src/client' import { gate, messageTexts, mockAdapter, text } from './helpers' @@ -23,16 +23,17 @@ describe('createHarnessClient', () => { agents: [pricer], expose: { agents: ['pricer'] }, }) - const host = createHarnessHost({ persistence: memoryPersistence() }) - const handler = createHarnessHandler({ - host, - harness: studio, - canAccess: () => true, - authorize: (request) => - request.headers.get('authorization') === 'Bearer t' - ? { id: 'u' } - : null, + const host = createHarnessHost({ + persistence: memoryPersistence(), + access: { + identify: ({ request }) => + request.headers.get('authorization') === 'Bearer t' + ? { id: 'u' } + : null, + authorize: () => true, + }, }) + const handler = createHarnessHandler({ host, harness: studio }) const client = createHarnessClient({ url: 'http://local/api/harness/', threadId: 'thread-c', @@ -73,8 +74,7 @@ describe('createHarnessClient', () => { const handler = createHarnessHandler({ host, harness: studio, - canAccess: () => true, - authorize: () => ({ id: 'u' }), + openAccess: { principal: { id: 'u' } }, }) const client = createHarnessClient({ url: 'http://local/api/harness', @@ -101,15 +101,15 @@ describe('createHarnessClient', () => { const { adapter, calls } = mockAdapter([() => text('went')]) const studio = defineHarness({ name: 'test/client-continue', adapter }) const persistence = memoryPersistence() - await persistence.stores.messages.saveThread('thread-go', [ - { id: 'u1', role: 'user', content: 'go' }, - ]) + await persistence.stores.messages.saveThread( + logIdOf({ logId: 'thread-go' }), + [{ id: 'u1', role: 'user', content: 'go' }], + ) const host = createHarnessHost({ persistence }) const handler = createHarnessHandler({ host, harness: studio, - canAccess: () => true, - authorize: () => ({ id: 'u' }), + openAccess: { principal: { id: 'u' } }, }) const client = createHarnessClient({ url: 'http://local/api/harness', @@ -143,8 +143,7 @@ describe('createHarnessClient', () => { const handler = createHarnessHandler({ host, harness: studio, - canAccess: () => true, - authorize: () => ({ id: 'u' }), + openAccess: { principal: { id: 'u' } }, }) const client = createHarnessClient({ url: 'http://local/api/harness', @@ -169,7 +168,9 @@ describe('createHarnessClient', () => { await vi.waitFor(async () => expect( JSON.stringify( - await persistence.stores.messages.loadThread('thread-tool'), + await persistence.stores.messages.loadThread( + logIdOf({ logId: 'thread-tool' }), + ), ), ).toContain('billing'), ) @@ -189,15 +190,14 @@ describe('createHarnessClient', () => { const studio = defineHarness({ name: 'test/client-ephemeral', adapter }) const persistence = memoryPersistence() const { messages } = persistence.stores - await messages.saveThread('thread-eph', [ - { id: 'u1', role: 'user', content: 'go' }, - ]) + // The key of the thread in the store, in the default tenant. + const key = logIdOf({ logId: 'thread-eph' }) + await messages.saveThread(key, [{ id: 'u1', role: 'user', content: 'go' }]) const host = createHarnessHost({ persistence }) const handler = createHarnessHandler({ host, harness: studio, - canAccess: () => true, - authorize: () => ({ id: 'u' }), + openAccess: { principal: { id: 'u' } }, }) const client = createHarnessClient({ url: 'http://local/api/harness', @@ -210,27 +210,23 @@ describe('createHarnessClient', () => { ephemeral: [{ role: 'user', content: 'note 2' }], }) await vi.waitFor(async () => - expect(await messages.loadThread('thread-eph')).toHaveLength(4), + expect(await messages.loadThread(key)).toHaveLength(4), ) expect(messageTexts(calls[0])).toEqual(['go', 'note 1']) expect(messageTexts(calls[1])).toEqual(['go', 'one', 'again', 'note 2']) - expect( - JSON.stringify(await messages.loadThread('thread-eph')), - ).not.toContain('note') + expect(JSON.stringify(await messages.loadThread(key))).not.toContain('note') await host.close() }) it('throws on a refused request instead of retrying', async () => { const { adapter } = mockAdapter([]) const studio = defineHarness({ name: 'test/client-denied', adapter }) - const host = createHarnessHost({ persistence: memoryPersistence() }) - const handler = createHarnessHandler({ - host, - harness: studio, - canAccess: () => true, - authorize: () => null, + const host = createHarnessHost({ + persistence: memoryPersistence(), + access: { identify: () => null, authorize: () => true }, }) + const handler = createHarnessHandler({ host, harness: studio }) const client = createHarnessClient({ url: 'http://local/api/harness', threadId: 't', @@ -266,8 +262,7 @@ describe('client.send to an agent run', () => { const handler = createHarnessHandler({ host, harness: studio, - canAccess: () => true, - authorize: () => ({ id: 'u' }), + openAccess: { principal: { id: 'u' } }, }) const client = createHarnessClient({ url: 'http://local/api/harness', diff --git a/packages/ai-harness/tests/coding-snapshots.test.ts b/packages/ai-harness/tests/coding-snapshots.test.ts index a5621f17bc..3dcfaae5eb 100644 --- a/packages/ai-harness/tests/coding-snapshots.test.ts +++ b/packages/ai-harness/tests/coding-snapshots.test.ts @@ -11,7 +11,7 @@ import { dirname, join } from 'node:path' import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest' import { toolDefinition } from '@tanstack/ai' import { memoryLogStore, memoryPersistence } from '@tanstack/ai-persistence' -import { createHarnessHost, defineHarness } from '../src' +import { createHarnessHost, defineHarness, logIdOf } from '../src' import { hostBackend } from '../src/first-party/coding/backend' import { quoteArg } from '../src/first-party/coding/search' import { snapshots } from '../src/first-party/coding/snapshots' @@ -28,6 +28,8 @@ import type { WorkspaceBackend } from '../src/first-party/coding/backend' import type { Reply } from './helpers' const THREAD = 't1' +/** The id of the log of `THREAD` in the store. */ +const LOG = logIdOf({ logId: THREAD }) const hasGit = (await hostBackend.exec('git --version')).exitCode === 0 let base = '' @@ -202,7 +204,7 @@ describe.skipIf(!hasGit)('snapshots', () => { await session.send('Edit a.txt') - const entries = await log.read({ logId: THREAD, after: 0 }) + const entries = await log.read({ logId: LOG, after: 0 }) const records = entries .map((entry) => entry.record) .filter((record) => record.type === 'tanstack/snapshots:step') diff --git a/packages/ai-harness/tests/durable-inputs.test.ts b/packages/ai-harness/tests/durable-inputs.test.ts index da2d5c64d2..0e472d11f7 100644 --- a/packages/ai-harness/tests/durable-inputs.test.ts +++ b/packages/ai-harness/tests/durable-inputs.test.ts @@ -6,7 +6,12 @@ import { memoryLogStore, memoryPersistence, } from '@tanstack/ai-persistence' -import { InputRejectedError, createHarnessHost, defineHarness } from '../src' +import { + InputRejectedError, + createHarnessHost, + defineHarness, + logIdOf, +} from '../src' import { after, gate, @@ -24,6 +29,8 @@ import type { HarnessDurability, HarnessInput, RecoverDecision } from '../src' import type { Reply } from './helpers' const THREAD = 't1' +/** The id of the log of `THREAD` in the store, and its key in the other stores. */ +const LOG = logIdOf({ logId: THREAD }) function durablePersistence() { const { runs, metadata } = memoryPersistence().stores @@ -93,14 +100,14 @@ async function seedStoppedTurn( : []), ] await persistence.stores.log.append({ - logId: THREAD, + logId: LOG, seq: 1, epoch: UNOWNED_EPOCH, records, }) await persistence.stores.runs.createOrResume({ runId: 'op-stopped', - threadId: THREAD, + threadId: LOG, startedAt: Date.now() - 60_000, }) await persistence.stores.runs.update('op-stopped', { @@ -133,7 +140,7 @@ async function openDurable(options: { } const logRecords = async (persistence: Durable) => - (await persistence.stores.log.read({ logId: THREAD })).map( + (await persistence.stores.log.read({ logId: LOG })).map( (entry) => entry.record, ) @@ -230,7 +237,7 @@ describe('recovery of a durable input', () => { it('settles aborted a pending input that was cancelled before it ran', async () => { const persistence = durablePersistence() await persistence.stores.log.append({ - logId: THREAD, + logId: LOG, seq: 1, epoch: UNOWNED_EPOCH, records: [ @@ -426,13 +433,13 @@ describe('recovery of a durable input', () => { await first.host.close() // The host stopped after it saved the answer, before the turn settled. // The live run gave the call its result. A log from before that has none. - const entries = await persistence.stores.log.read({ logId: THREAD }) + const entries = await persistence.stores.log.read({ logId: LOG }) const settledAt = entries.findIndex( (entry) => entry.record.type === 'harness.input.settled', ) const stopped = durablePersistence() await stopped.stores.log.append({ - logId: THREAD, + logId: LOG, seq: 1, epoch: UNOWNED_EPOCH, records: entries.slice(0, settledAt).map(({ record }) => @@ -532,7 +539,7 @@ describe('recovery of a durable input', () => { it('fails a background agent whose host stopped, and wakes the thread', async () => { const persistence = durablePersistence() await persistence.stores.log.append({ - logId: THREAD, + logId: LOG, seq: 1, epoch: UNOWNED_EPOCH, records: [ @@ -552,7 +559,7 @@ describe('recovery of a durable input', () => { }) await persistence.stores.runs.createOrResume({ runId: 'op-agent', - threadId: THREAD, + threadId: LOG, startedAt: Date.now() - 60_000, kind: 'agent', agent: 'waiter', @@ -596,7 +603,7 @@ describe('recovery of a durable input', () => { it('keeps the wake mark of a logged follow-up that never ran', async () => { const persistence = durablePersistence() await persistence.stores.log.append({ - logId: THREAD, + logId: LOG, seq: 1, epoch: UNOWNED_EPOCH, records: [ @@ -689,7 +696,7 @@ describe('recovery of a durable input', () => { it('asks the recover hook about an input that never ran', async () => { const persistence = durablePersistence() await persistence.stores.log.append({ - logId: THREAD, + logId: LOG, seq: 1, epoch: UNOWNED_EPOCH, records: [pendingInput('in-new', 'later')], @@ -714,8 +721,8 @@ describe('recovery of a durable input', () => { records: Array, ) => persistence.stores.log.append({ - logId: THREAD, - seq: (await persistence.stores.log.read({ logId: THREAD })).length + 1, + logId: LOG, + seq: (await persistence.stores.log.read({ logId: LOG })).length + 1, epoch: UNOWNED_EPOCH, records, }) @@ -881,9 +888,9 @@ describe('session.recover()', () => { expect(calls).toHaveLength(0) // The other host stored one more input, then stopped. - const head = (await log.read({ logId: THREAD })).length + const head = (await log.read({ logId: LOG })).length await log.append({ - logId: THREAD, + logId: LOG, seq: head + 1, epoch: UNOWNED_EPOCH, records: [pendingInput('in-2', 'next')], @@ -1199,7 +1206,7 @@ describe('waiting inputs', () => { async (_when, landed) => { const persistence = durablePersistence() await persistence.stores.log.append({ - logId: THREAD, + logId: LOG, seq: 1, epoch: UNOWNED_EPOCH, records: [ @@ -1235,7 +1242,7 @@ describe('waiting inputs', () => { async (_when, landed) => { const persistence = durablePersistence() await persistence.stores.log.append({ - logId: THREAD, + logId: LOG, seq: 1, epoch: UNOWNED_EPOCH, records: [ @@ -1431,7 +1438,7 @@ describe('session.continue()', () => { async (_name, messages) => { const persistence = durablePersistence() await persistence.stores.log.append({ - logId: THREAD, + logId: LOG, seq: 1, epoch: UNOWNED_EPOCH, records: [transcript(...messages)], @@ -1459,7 +1466,7 @@ describe('session.continue()', () => { it('returns the first operation for a duplicate inputId, and runs once', async () => { const persistence = durablePersistence() await persistence.stores.log.append({ - logId: THREAD, + logId: LOG, seq: 1, epoch: UNOWNED_EPOCH, records: [transcript(user)], @@ -1481,7 +1488,7 @@ describe('session.continue()', () => { it('runs a continue that a stopped host left pending', async () => { const persistence = durablePersistence() await persistence.stores.log.append({ - logId: THREAD, + logId: LOG, seq: 1, epoch: UNOWNED_EPOCH, records: [ @@ -1519,7 +1526,7 @@ describe('session.continue()', () => { const persistence = durablePersistence() if (records.length > 0) { await persistence.stores.log.append({ - logId: THREAD, + logId: LOG, seq: 1, epoch: UNOWNED_EPOCH, records, @@ -1757,9 +1764,9 @@ describe('durable approval recovery', () => { const approval = result.interrupts?.[0] if (!approval) throw new Error('The turn did not stop for approval.') await first.close() - const entries = await log.read({ logId: THREAD }) + const entries = await log.read({ logId: LOG }) await log.append({ - logId: THREAD, + logId: LOG, seq: entries.length + 1, epoch: UNOWNED_EPOCH, records: [ @@ -1784,7 +1791,7 @@ describe('durable approval recovery', () => { }) await stores.runs.createOrResume({ runId: 'resume-stopped', - threadId: THREAD, + threadId: LOG, startedAt: Date.now() - 60000, }) await stores.runs.update('resume-stopped', { @@ -1849,7 +1856,7 @@ describe('durable approval recovery', () => { const result = await session.send('remove') const approval = result.interrupts?.[0] if (!approval) throw new Error('The turn did not stop for approval.') - const cached = await stores.metadata.get('harness:interrupted', THREAD) + const cached = await stores.metadata.get('harness:interrupted', LOG) const deletion = vi .spyOn(stores.metadata, 'delete') .mockRejectedValueOnce(new Error('metadata delete failed')) @@ -1862,7 +1869,7 @@ describe('durable approval recovery', () => { deletion.mockRestore() expect(session.snapshot().pendingInterrupts).toEqual([]) await host.close() - await stores.metadata.set('harness:interrupted', THREAD, cached) + await stores.metadata.set('harness:interrupted', LOG, cached) const next = createHarnessHost({ persistence }) const reopened = await next.open(config, { threadId: THREAD }) expect(reopened.snapshot().pendingInterrupts).toEqual([]) diff --git a/packages/ai-harness/tests/durable-joins.test.ts b/packages/ai-harness/tests/durable-joins.test.ts index cecb508abf..da10edd443 100644 --- a/packages/ai-harness/tests/durable-joins.test.ts +++ b/packages/ai-harness/tests/durable-joins.test.ts @@ -5,7 +5,7 @@ import { memoryLogStore, memoryPersistence, } from '@tanstack/ai-persistence' -import { createHarnessHost, defineHarness } from '../src' +import { createHarnessHost, defineHarness, logIdOf } from '../src' import { after, gate, @@ -21,6 +21,8 @@ import type { HarnessTurnOptions } from '../src' import type { Reply } from './helpers' const THREAD = 't1' +/** The id of the log of `THREAD` in the store, and its key in the other stores. */ +const LOG = logIdOf({ logId: THREAD }) /** * A durable store set. `batches` has the record types of each append. @@ -117,14 +119,14 @@ async function seedStoppedTurn( : []), ] await persistence.stores.log.append({ - logId: THREAD, + logId: LOG, seq: 1, epoch: UNOWNED_EPOCH, records, }) await persistence.stores.runs.createOrResume({ runId: 'op-stopped', - threadId: THREAD, + threadId: LOG, startedAt: Date.now() - 60_000, }) await persistence.stores.runs.update('op-stopped', { @@ -428,7 +430,7 @@ describe('the join rule', () => { // It did not join, so it ends with the session that stopped. await expect(join).rejects.toThrow('Session closed.') - const types = (await inner.read({ logId: THREAD })).map( + const types = (await inner.read({ logId: LOG })).map( (entry) => entry.record.type, ) expect(types).not.toContain('harness.input.joined') @@ -516,7 +518,7 @@ describe('the join rule', () => { await turn expect(messageTexts(calls[1]).slice(-2)).toEqual(['more', 'only now']) - expect(JSON.stringify(await inner.read({ logId: THREAD }))).not.toContain( + expect(JSON.stringify(await inner.read({ logId: LOG }))).not.toContain( 'only now', ) await host.close() @@ -556,7 +558,7 @@ describe('the join rule', () => { outcome: 'completed', operationId: turn.id, }) - const types = (await inner.read({ logId: THREAD })) + const types = (await inner.read({ logId: LOG })) .map((entry) => entry.record) .filter((record) => record.inputId === 's1') .map((record) => record.type) diff --git a/packages/ai-harness/tests/durable-session.test.ts b/packages/ai-harness/tests/durable-session.test.ts index 7367ab7bd4..7788d7fca6 100644 --- a/packages/ai-harness/tests/durable-session.test.ts +++ b/packages/ai-harness/tests/durable-session.test.ts @@ -5,7 +5,7 @@ import { memoryLogStore, memoryPersistence, } from '@tanstack/ai-persistence' -import { createHarnessHost, defineHarness, definePlugin } from '../src' +import { createHarnessHost, defineHarness, definePlugin, logIdOf } from '../src' import { loadLogState, sessionOf } from '../src/log' import { gate, @@ -22,6 +22,8 @@ import type { HarnessSession, ProjectOptions, SessionEvent } from '../src' import type { Reply } from './helpers' const THREAD = 't1' +/** The id of the log of `THREAD` in the store. */ +const LOG = logIdOf({ logId: THREAD }) /** The stores of a durable host: a log, runs, and metadata. */ function durablePersistence(log: LogStore = memoryLogStore()) { @@ -275,7 +277,7 @@ describe('durable session log', () => { // The deltas reach the log after the coalescing window. await vi.waitFor(async () => { - const entries = await persistence.stores.log.read({ logId: THREAD }) + const entries = await persistence.stores.log.read({ logId: LOG }) const events = entries.map((entry) => entry.record.event) expect(JSON.stringify(events)).toContain('wor') }) @@ -446,7 +448,8 @@ describe('durable session log', () => { const rebuilt = sessionOf( await loadLogState({ store: persistence.stores.log, - logId: THREAD, + logId: LOG, + rootThread: THREAD, project, }), THREAD, @@ -493,7 +496,8 @@ describe('durable session log', () => { const rebuilt = sessionOf( await loadLogState({ store: persistence.stores.log, - logId: THREAD, + logId: LOG, + rootThread: THREAD, project, }), THREAD, @@ -517,7 +521,7 @@ describe('durable session log', () => { const harness = defineHarness({ name: 'test/conflict', adapter }) const session = await host.open(harness, { threadId: THREAD }) await inner.append({ - logId: THREAD, + logId: LOG, seq: 1, epoch: UNOWNED_EPOCH, records: [{ type: 'app.from-another-host' }], @@ -525,9 +529,10 @@ describe('durable session log', () => { await expect(session.send('hi')).rejects.toThrow('Another host wrote') - // The host opens a new session, which folds the log again. - await vi.waitFor(async () => - expect(await host.open(harness, { threadId: THREAD })).not.toBe(session), + // The host drops the stopped session, so the next open folds the log + // again in a new one. + await vi.waitFor(() => + expect(host.isOpen(harness, { threadId: THREAD })).toBe(false), ) await host.close() }) @@ -595,7 +600,11 @@ describe('durable session log', () => { // A fold of the log gives the same messages. const rebuilt = sessionOf( - await loadLogState({ store: persistence.stores.log, logId: THREAD }), + await loadLogState({ + store: persistence.stores.log, + logId: LOG, + rootThread: THREAD, + }), THREAD, ) expect(rebuilt.messages).toEqual(transcript) @@ -642,7 +651,8 @@ describe('LogRecordsCapability', () => { const rebuilt = sessionOf( await loadLogState({ store: persistence.stores.log, - logId: THREAD, + logId: LOG, + rootThread: THREAD, project, }), THREAD, @@ -697,7 +707,7 @@ describe('LogRecordsCapability', () => { '[signal] saved', ) const entries = await persistence.stores.log.read({ - logId: THREAD, + logId: LOG, after: 0, }) expect(entries.some((entry) => entry.record.type === 'app.signal')).toBe( @@ -706,7 +716,8 @@ describe('LogRecordsCapability', () => { const rebuilt = sessionOf( await loadLogState({ store: persistence.stores.log, - logId: THREAD, + logId: LOG, + rootThread: THREAD, project, }), THREAD, @@ -819,7 +830,8 @@ describe('compaction-style host records from a middleware', () => { async function rebuild(persistence: ReturnType) { const state = await loadLogState({ store: persistence.stores.log, - logId: THREAD, + logId: LOG, + rootThread: THREAD, project, }) return sessionOf(state, THREAD).messages @@ -860,7 +872,7 @@ describe('compaction-style host records from a middleware', () => { // (a) The record comes after the transcript record that saved the reply. const entries = await persistence.stores.log.read({ - logId: THREAD, + logId: LOG, after: 0, }) const replyAt = entries.findIndex( diff --git a/packages/ai-harness/tests/durable-tools.test.ts b/packages/ai-harness/tests/durable-tools.test.ts index febc2b03d9..37766271ea 100644 --- a/packages/ai-harness/tests/durable-tools.test.ts +++ b/packages/ai-harness/tests/durable-tools.test.ts @@ -2,13 +2,15 @@ import { describe, expect, it, vi } from 'vitest' import { z } from 'zod' import { EventType, getLogRecords, toolDefinition } from '@tanstack/ai' import { memoryLogStore, memoryPersistence } from '@tanstack/ai-persistence' -import { createHarnessHost, defineHarness, durableTool } from '../src' +import { createHarnessHost, defineHarness, durableTool, logIdOf } from '../src' import { INTERRUPTED_TOOL_RESULT } from '../src/resume' import { gate, mockAdapter, text, toolCall } from './helpers' import type { AnyChatMiddleware, AnyTool, StreamChunk } from '@tanstack/ai' import type { Reply } from './helpers' const THREAD = 't1' +/** The id of the log of `THREAD` in the store. */ +const LOG = logIdOf({ logId: THREAD }) function durablePersistence() { const { runs, metadata } = memoryPersistence().stores @@ -35,7 +37,7 @@ async function openDurable( /** The records of the log, of one type. */ async function recordsOf(persistence: Durable, type: string) { - return (await persistence.stores.log.read({ logId: THREAD })) + return (await persistence.stores.log.read({ logId: LOG })) .map((entry) => entry.record) .filter((record) => record.type === type) } diff --git a/packages/ai-harness/tests/extensions.test.ts b/packages/ai-harness/tests/extensions.test.ts index 2651e8139c..d2f9dae4f1 100644 --- a/packages/ai-harness/tests/extensions.test.ts +++ b/packages/ai-harness/tests/extensions.test.ts @@ -11,6 +11,7 @@ import { defineCommand, defineHarness, definePlugin, + logIdOf, } from '../src' import { mockAdapter, text } from './helpers' import type { StreamChunk } from '@tanstack/ai' @@ -182,7 +183,10 @@ describe('events and state', () => { ).rejects.toThrow('Input validation failed') expect(heard).toEqual(['milk', 'eggs']) expect( - await persistence.stores.metadata.get('plugin:test/todos', 't'), + await persistence.stores.metadata.get( + 'plugin:test/todos', + logIdOf({ logId: 't' }), + ), ).toEqual({ items: ['milk', 'eggs'], }) @@ -376,8 +380,8 @@ describe('adapter override and credentials', () => { commands: { token: defineCommand({ description: 'Read the GitHub token', - run: async () => { - const credential = await ctx.credentials.require('github') + run: async (_input, { credentials }) => { + const credential = await credentials.require('github') token = credential.type === 'oauth' ? credential.accessToken diff --git a/packages/ai-harness/tests/first-party.test.ts b/packages/ai-harness/tests/first-party.test.ts index 2dbcf546ac..089a54c17c 100644 --- a/packages/ai-harness/tests/first-party.test.ts +++ b/packages/ai-harness/tests/first-party.test.ts @@ -3,7 +3,7 @@ import { tmpdir } from 'node:os' import { join } from 'node:path' import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest' import { memoryPersistence } from '@tanstack/ai-persistence' -import { createHarnessHost, defineHarness } from '../src' +import { createHarnessHost, defineHarness, logIdOf } from '../src' import { compact, decidePermission, @@ -19,6 +19,9 @@ import { messageTexts, mockAdapter, text, toolCall } from './helpers' import type { HarnessPlugin } from '../src' import type { AnyTextAdapter } from '@tanstack/ai' +/** The key of the thread `t` in the stores, in the default tenant. */ +const T_KEY = logIdOf({ logId: 't' }) + let root = '' beforeEach(async () => { root = await mkdtemp(join(tmpdir(), 'harness-')) @@ -170,7 +173,7 @@ describe('todos', () => { await session.send('continue') expect(JSON.stringify(calls[2].systemPrompts)).toContain('[~] write tests') expect( - await persistence.stores.metadata.get('plugin:tanstack/todos', 't'), + await persistence.stores.metadata.get('plugin:tanstack/todos', T_KEY), ).toEqual({ items: [{ text: 'write tests', status: 'in_progress' }], }) @@ -238,7 +241,7 @@ describe('compact and usage', () => { expect(await session.command({ name: 'compact' })).toBe( 'Compacted 4 messages into a summary.', ) - const transcript = await persistence.stores.messages.loadThread('t') + const transcript = await persistence.stores.messages.loadThread(T_KEY) expect(transcript).toHaveLength(2) expect(transcript[0]?.content).toContain('We planned a trip.') expect(await session.command({ name: 'usage' })).toBe( diff --git a/packages/ai-harness/tests/fork.test.ts b/packages/ai-harness/tests/fork.test.ts index da3711d767..fb5e89ff29 100644 --- a/packages/ai-harness/tests/fork.test.ts +++ b/packages/ai-harness/tests/fork.test.ts @@ -113,12 +113,11 @@ describe.each([ await source.send('one') await host.fork(harness, { threadId: SOURCE, newThreadId: FORK }) // The app opened the source: the fork left it open. - expect(await host.open(harness, { threadId: SOURCE })).toBe(source) + expect(host.isOpen(harness, { threadId: SOURCE })).toBe(true) await source.close() await host.fork(harness, { threadId: SOURCE, newThreadId: 'other' }) // The fork opened the source itself, so it closed it again. - const reopened = await host.open(harness, { threadId: SOURCE }) - expect(reopened).not.toBe(source) + expect(host.isOpen(harness, { threadId: SOURCE })).toBe(false) await host.close() }) diff --git a/packages/ai-harness/tests/forwarding.test.ts b/packages/ai-harness/tests/forwarding.test.ts index b3ba4d7776..50aad7da15 100644 --- a/packages/ai-harness/tests/forwarding.test.ts +++ b/packages/ai-harness/tests/forwarding.test.ts @@ -7,6 +7,7 @@ import { defineCommand, defineHarness, definePlugin, + logIdOf, } from '../src' import { SessionImpl } from '../src/session/harness-session' import { @@ -22,6 +23,8 @@ import type { AnyTextAdapter } from '@tanstack/ai' import type { ChatTurnResult, Operation } from '../src' const THREAD = 't1' +/** The id of the log of `THREAD` in the store, and its key in the other stores. */ +const LOG = logIdOf({ logId: THREAD }) const NAME = 'test/forwarding' /** Stores that several hosts share, as on a serverless platform. */ @@ -143,7 +146,7 @@ describe('an input that reaches a host that does not own the thread', () => { const harness = harnessWith(adapter) // A host that stopped holds a claim that is about to expire. await persistence.stores.workClaims.claim({ - threadId: THREAD, + threadId: LOG, harness: NAME, ownerId: 'stopped-host', until: Date.now() + 50, @@ -294,12 +297,12 @@ describe('an input that reaches a host that does not own the thread', () => { await vi.waitFor(async () => expect( await persistence.stores.inbox!.get({ - threadId: THREAD, + threadId: LOG, inputId: receipt.inputId, }), ).toBeNull(), ) - const log = await persistence.stores.log.read({ logId: THREAD }) + const log = await persistence.stores.log.read({ logId: LOG }) expect(JSON.stringify(log)).not.toContain('hunter2') }) @@ -350,7 +353,7 @@ describe('an input that reaches a host that does not own the thread', () => { await turn.receipt await other.command({ name: 'count' }) await vi.waitFor( - async () => expect(await inbox.listPending(THREAD)).toEqual([]), + async () => expect(await inbox.listPending(LOG)).toEqual([]), { timeout: 4_000 }, ) expect(runs).toBe(1) @@ -395,7 +398,7 @@ describe('an input that reaches a host that does not own the thread', () => { it('keeps a bounded number of forwarded ids', async () => { const persistence = sharedStores() await persistence.stores.workClaims.claim({ - threadId: THREAD, + threadId: LOG, harness: NAME, ownerId: 'busy-host', until: Date.now() + 60_000, @@ -456,7 +459,7 @@ describe('an input that reaches a host that does not own the thread', () => { const { adapter, calls } = mockAdapter([() => text('Picked up.')]) // A host that stopped holds a claim that is about to expire. await persistence.stores.workClaims.claim({ - threadId: THREAD, + threadId: LOG, harness: NAME, ownerId: 'stopped-host', until: Date.now() + 50, @@ -511,7 +514,7 @@ describe('an input that reaches a host that does not own the thread', () => { await new Promise((resolve) => setTimeout(resolve, 1_500)) // The input waits for the next host. expect( - (await persistence.stores.inbox!.listPending(THREAD)).map( + (await persistence.stores.inbox!.listPending(LOG)).map( (entry) => entry.inputId, ), ).toEqual([receipt.inputId]) @@ -576,7 +579,7 @@ describe('an input that reaches a host that does not own the thread', () => { // comes (a turn that still ends, for example), then stops too. await owner.close() await persistence.stores.workClaims.claim({ - threadId: THREAD, + threadId: LOG, harness: NAME, ownerId: 'busy-host', until: Date.now() + 100, diff --git a/packages/ai-harness/tests/goal.test.ts b/packages/ai-harness/tests/goal.test.ts index c79bc2dc35..cdbef87ceb 100644 --- a/packages/ai-harness/tests/goal.test.ts +++ b/packages/ai-harness/tests/goal.test.ts @@ -2,7 +2,7 @@ import { describe, expect, it, vi } from 'vitest' import { z } from 'zod' import { toolDefinition } from '@tanstack/ai' import { memoryPersistence } from '@tanstack/ai-persistence' -import { createHarnessHost, defineHarness, definePlugin } from '../src' +import { createHarnessHost, defineHarness, definePlugin, logIdOf } from '../src' import { GoalMet, goal, selectGoal } from '../src/first-party' import { after, @@ -328,7 +328,11 @@ describe('goal', () => { note: '', queued: '', } - await persistence.stores.metadata.set('plugin:tanstack/goal', 't', active) + await persistence.stores.metadata.set( + 'plugin:tanstack/goal', + logIdOf({ logId: 't' }), + active, + ) const { host, session, calls, read } = await start({ persistence, replies: [() => text('Hello.')], diff --git a/packages/ai-harness/tests/helpers.ts b/packages/ai-harness/tests/helpers.ts index c5a491e7ae..e6735677d7 100644 --- a/packages/ai-harness/tests/helpers.ts +++ b/packages/ai-harness/tests/helpers.ts @@ -4,8 +4,10 @@ import type { AnyTextAdapter, ModelMessage, StreamChunk, + WebSocketLike, } from '@tanstack/ai' import type { LogStore } from '@tanstack/ai-persistence' +import type { HostFrame } from '../src/protocol' /** * A log store that calls `overrides` for the methods it has, and `inner` for @@ -204,6 +206,75 @@ export const wakeMessages = (messages: ReadonlyArray) => ) .map((message) => message.content) +/** + * A WebSocket for `handleHarnessSocket`. `sent` has each frame it sent, + * `closed` each close code. `receive` gives it a frame from the client, and + * `fail` fires its error event. With `failSend`, each send throws. + */ +export function fakeSocket(options: { failSend?: boolean } = {}) { + const handlers: Record void>> = {} + const sent: Array = [] + const closed: Array = [] + const socket: WebSocketLike = { + send: (data) => { + if (options.failSend) throw new Error('socket gone') + sent.push(JSON.parse(data)) + }, + close: (code?: number) => { + closed.push(code) + handlers.close?.forEach((handler) => handler()) + }, + addEventListener: (type: string, handler: (event?: any) => void) => { + ;(handlers[type] ??= []).push(handler) + }, + } + const receive = (frame: unknown) => + handlers.message?.forEach((handler) => + handler({ + data: typeof frame === 'string' ? frame : JSON.stringify(frame), + }), + ) + const fail = () => handlers.error?.forEach((handler) => handler()) + return { socket, sent, closed, receive, fail } +} + +/** + * The JSON `data` of each SSE event of `response`, until `until` says stop + * (then the read is cancelled) or the stream ends. + */ +export async function readSse( + response: Response, + until: (frame: T) => boolean, +) { + if (!response.body) throw new Error('The response has no body.') + const reader = response.body.getReader() + const decoder = new TextDecoder() + const frames: Array = [] + let buffer = '' + while (true) { + const { value, done } = await reader.read() + if (done) break + buffer += decoder.decode(value, { stream: true }) + const blocks = buffer.split('\n\n') + buffer = blocks.pop() ?? '' + for (const block of blocks) { + const data = block + .split('\n') + .find((line) => line.startsWith('data: ')) + ?.slice(6) + if (!data) continue + // The handler writes each frame as JSON. + const parsed: T = JSON.parse(data) + frames.push(parsed) + if (until(parsed)) { + await reader.cancel() + return frames + } + } + } + return frames +} + /** * A middleware that sends the model `[hidden]` in place of each message * `secret`, at each model call. The transcript keeps `secret`. diff --git a/packages/ai-harness/tests/http-access.test.ts b/packages/ai-harness/tests/http-access.test.ts new file mode 100644 index 0000000000..807b623f1c --- /dev/null +++ b/packages/ai-harness/tests/http-access.test.ts @@ -0,0 +1,309 @@ +import { describe, expect, it, vi } from 'vitest' +import { memoryPersistence } from '@tanstack/ai-persistence' +import { createHarnessHost, defineHarness, membersOnly } from '../src' +import { createHarnessHandler, handleHarnessSocket } from '../src/server' +import { createHarnessClient } from '../src/client' +import { isMediaRecord } from '../src/media-entry' +import { fakeSocket, mockAdapter, readSse, text } from './helpers' +import type { HostEvent } from '../src' + +const BASE = 'http://x/api/harness' + +const olga = { id: 'olga', tenantId: 'acme' } +const bea = { id: 'bea', tenantId: 'acme' } +const cid = { id: 'cid', tenantId: 'acme' } +const gus = { id: 'gus', tenantId: 'globex' } +const users = new Map([olga, bea, cid, gus].map((user) => [user.id, user])) + +/** A request of the user named by the `x-user` header. */ +const requestAs = (userId: string, path = 'socket', init: RequestInit = {}) => + new Request(`${BASE}/${path}`, { + ...init, + headers: { + 'x-user': userId, + 'content-type': 'application/json', + ...init.headers, + }, + }) + +/** A membersOnly() host whose users come from the `x-user` header. */ +function setup() { + const harness = defineHarness({ + name: 'test/http-access', + adapter: mockAdapter(() => text('ok')).adapter, + }) + const host = createHarnessHost({ + persistence: memoryPersistence(), + access: { + identify: ({ request }) => + users.get(request.headers.get('x-user') ?? '') ?? null, + authorize: membersOnly(), + }, + }) + const handler = createHarnessHandler({ host, harness }) + const call = (userId: string, path: string, init?: RequestInit) => + handler(requestAs(userId, path, init)) + const prompt = (userId: string, threadId: string) => + call(userId, 'control', { + method: 'POST', + body: JSON.stringify({ + threadId, + input: { op: 'prompt', message: 'Hi.' }, + }), + }) + const share = (userId: string, role: 'member' | 'viewer') => + host.sessions.share({ threadId: 't1', userId, role, principal: olga }) + return { host, harness, handler, call, prompt, share } +} + +/** The thread of a host event, from its own field or its entry. */ +const threadOf = (event: HostEvent) => + event.type === 'session' || event.type === 'session-deleted' + ? event.entry.threadId + : event.threadId + +describe('handlers and sockets fail closed', () => { + it('throws at the start for a host without access and without openAccess', () => { + const host = createHarnessHost({ persistence: memoryPersistence() }) + const harness = defineHarness({ name: 'test/open' }) + + expect(() => createHarnessHandler({ host, harness })).toThrow( + 'createHarnessHandler needs host access ({ identify, authorize }) or openAccess: true', + ) + expect(() => + handleHarnessSocket({ + host, + harness, + socket: fakeSocket().socket, + request: requestAs('olga'), + }), + ).toThrow( + 'handleHarnessSocket needs host access ({ identify, authorize }) or openAccess: true', + ) + }) + + it('throws at the start for openAccess on a host with access', async () => { + const { host, harness } = setup() + + expect(() => + createHarnessHandler({ host, harness, openAccess: true }), + ).toThrow( + 'createHarnessHandler: the host has access, so access.identify reads the user of each request. Remove openAccess.', + ) + expect(() => + handleHarnessSocket({ + host, + harness, + socket: fakeSocket().socket, + request: requestAs('olga'), + openAccess: { principal: olga }, + }), + ).toThrow('handleHarnessSocket: the host has access') + await host.close() + }) +}) + +describe('the handler on a host with access', () => { + it('refuses a stranger with 403, and a read of a missing thread is 404 and starts nothing', async () => { + const { host, call, prompt } = setup() + expect((await prompt('olga', 't1')).status).toBe(200) + + expect((await call('cid', 'transcript?threadId=t1')).status).toBe(403) + expect((await prompt('cid', 't1')).status).toBe(403) + expect((await call('nobody', 'transcript?threadId=t1')).status).toBe(401) + expect((await call('cid', 'events?threadId=t-none')).status).toBe(404) + expect((await call('cid', 'snapshot?threadId=t-none')).status).toBe(404) + expect( + await host.sessions.get({ threadId: 't-none', principal: cid }), + ).toBeUndefined() + await host.close() + }) + + it('shares and unshares a thread through POST sessions', async () => { + const { host, call, prompt } = setup() + await prompt('olga', 't1') + const sessionsOp = (userId: string, body: object) => + call(userId, 'sessions', { method: 'POST', body: JSON.stringify(body) }) + + const shared = await sessionsOp('olga', { + op: 'share', + threadId: 't1', + userId: 'bea', + role: 'viewer', + }) + expect((await shared.json()).members).toEqual([ + { userId: 'olga', role: 'owner' }, + { userId: 'bea', role: 'viewer' }, + ]) + expect((await call('bea', 'transcript?threadId=t1')).status).toBe(200) + expect((await prompt('bea', 't1')).status).toBe(403) + expect( + ( + await sessionsOp('bea', { + op: 'share', + threadId: 't1', + userId: 'cid', + role: 'member', + }) + ).status, + ).toBe(403) + + await sessionsOp('olga', { op: 'unshare', threadId: 't1', userId: 'bea' }) + + expect((await call('bea', 'transcript?threadId=t1')).status).toBe(403) + await host.close() + }) + + it('refuses a signed media URL after unshare removed its user', async () => { + const { host, handler, call, prompt, share } = setup() + await prompt('olga', 't1') + const uploaded = await call('olga', 'media?threadId=t1&name=cat.png', { + method: 'POST', + body: 'hello', + headers: { 'content-type': 'image/png' }, + }) + const record: unknown = await uploaded.json() + if (!isMediaRecord(record)) throw new Error('The upload failed.') + await share('bea', 'viewer') + const signed = await call('bea', `media-url?threadId=t1&id=${record.id}`) + const { path }: { path: string } = await signed.json() + const fetchSigned = () => handler(new Request(`${BASE}/${path}`)) + expect(await (await fetchSigned()).text()).toBe('hello') + + await host.sessions.unshare({ + threadId: 't1', + userId: 'bea', + principal: olga, + }) + + expect((await fetchSigned()).status).toBe(403) + await host.close() + }) + + it('closes the open events stream of a user on revoke', async () => { + const { host, call, prompt, share } = setup() + await prompt('olga', 't1') + await share('bea', 'member') + const response = await call('bea', 'events?threadId=t1') + const reading = readSse(response, () => false) + + host.revoke({ threadId: 't1', principal: bea }) + + // The read ends only when the stream closes. + await expect(reading).resolves.toBeInstanceOf(Array) + await host.close() + }) + + it('never shows the threads of another tenant in host-events', async () => { + const { host, call, prompt } = setup() + const reading = readSse( + await call('gus', 'host-events'), + (event: HostEvent) => threadOf(event) === 'g1', + ) + + await prompt('olga', 't1') + await prompt('gus', 'g1') + + const seen = (await reading).map(threadOf) + expect(seen).toContain('g1') + expect(seen).not.toContain('t1') + await host.close() + }) +}) + +describe('a new thread and a deleted thread', () => { + it('lets the client stream the events of a new thread before its first message', async () => { + const { handler, prompt } = setup() + const client = createHarnessClient({ + url: BASE, + threadId: 't-new', + reconnectDelayMs: 5, + headers: { 'x-user': 'olga' }, + fetch: (input, init) => handler(new Request(input, init)), + }) + const reading = new AbortController() + const first = (async () => { + for await (const entry of client.events({ signal: reading.signal })) { + return entry + } + return undefined + })() + // The thread has no entry yet, so the first reads answer 404. + await new Promise((resolve) => setTimeout(resolve, 20)) + expect((await prompt('olga', 't-new')).status).toBe(200) + + expect(await first).toBeDefined() + reading.abort() + }) + + it('deletes the transcript with the entry, so the next opener reads nothing', async () => { + const { host, harness } = setup() + const session = await host.open(harness, { + threadId: 't1', + principal: olga, + }) + await session.send('A secret.') + expect(await session.transcript()).not.toEqual([]) + + await host.sessions.delete({ threadId: 't1', principal: olga }) + + const reopened = await host.open(harness, { + threadId: 't1', + principal: bea, + }) + expect(await reopened.transcript()).toEqual([]) + await host.close() + }) +}) + +describe('the socket on a host with access', () => { + it('refuses a request that identify refuses, and a stranger to a thread', async () => { + const { host, harness, prompt } = setup() + await prompt('olga', 't1') + const nobody = fakeSocket() + const stranger = fakeSocket() + handleHarnessSocket({ + host, + harness, + socket: nobody.socket, + request: requestAs('nobody'), + }) + handleHarnessSocket({ + host, + harness, + socket: stranger.socket, + request: requestAs('cid'), + }) + + nobody.receive({ type: 'harness.subscribe', threadId: 't1' }) + stranger.receive({ type: 'harness.subscribe', threadId: 't1' }) + + await vi.waitFor(() => expect(nobody.closed).toEqual([4401])) + expect(nobody.sent).toEqual([ + { type: 'harness.error', message: 'unauthorized' }, + ]) + await vi.waitFor(() => expect(stranger.closed).toEqual([4403])) + expect(stranger.sent).toEqual([ + { type: 'harness.error', message: 'forbidden' }, + ]) + await host.close() + }) + + it('sends forbidden and closes when the user is revoked', async () => { + const { host, harness, prompt, share } = setup() + await prompt('olga', 't1') + await share('bea', 'member') + const { socket, sent, closed, receive } = fakeSocket() + handleHarnessSocket({ host, harness, socket, request: requestAs('bea') }) + receive({ type: 'harness.subscribe', threadId: 't1' }) + await vi.waitFor(() => + expect(sent.some((frame) => frame.type === 'harness.hello')).toBe(true), + ) + + host.revoke({ threadId: 't1', principal: bea }) + + await vi.waitFor(() => expect(closed).toEqual([4403])) + expect(sent.at(-1)).toEqual({ type: 'harness.error', message: 'forbidden' }) + await host.close() + }) +}) diff --git a/packages/ai-harness/tests/http-media.test.ts b/packages/ai-harness/tests/http-media.test.ts index 0f42598c30..55ef0fe9aa 100644 --- a/packages/ai-harness/tests/http-media.test.ts +++ b/packages/ai-harness/tests/http-media.test.ts @@ -6,21 +6,21 @@ import { isMediaRecord, mediaPart } from '../src/media-entry' import { mockAdapter, text } from './helpers' import type { ContentPart } from '@tanstack/ai' import type { MediaRecord } from '../src' -import type { Authorize } from '../src/server' import type { MediaOptions } from '../src/media-entry' const BASE = 'http://x/api/harness' const HOUR = 60 * 60 * 1000 const NOW = 1_000_000 -const allowUser: Authorize = (request) => - request.headers.get('authorization') === 'Bearer good' - ? { id: 'user-1' } - : null +const user = { id: 'user-1' } + +const allowUser = (request: Request) => + request.headers.get('authorization') === 'Bearer good' ? user : null /** * A handler for user-1, who can open only threads that start with `user-1`. - * `handlerWith` makes another handler on the same host. + * `handlerWith` makes another handler, on the same host or on `hostWith`: a + * second host on the same stores, as another server instance. */ function setup(options: { media?: MediaOptions; secret?: string } = {}) { const { adapter, calls } = mockAdapter(() => text('ok')) @@ -29,16 +29,25 @@ function setup(options: { media?: MediaOptions; secret?: string } = {}) { adapter, media: options.media, }) - const host = createHarnessHost({ persistence: memoryPersistence() }) + const persistence = memoryPersistence() + /** A host on the shared stores. `identify` reads the user of a request. */ + const hostWith = (identify: (request: Request) => typeof user | null) => + createHarnessHost({ + persistence, + access: { + identify: ({ request }) => identify(request), + authorize: ({ principal, session }) => + session.threadId.startsWith(principal.id), + }, + }) + const host = hostWith(allowUser) /** `secret` signs the media URLs. Without it, a random key per handler. */ const handlerWith = ( - handlerOptions: { authorize?: Authorize; secret?: string } = {}, + handlerOptions: { host?: typeof host; secret?: string } = {}, ) => createHarnessHandler({ - host, + host: handlerOptions.host ?? host, harness, - authorize: handlerOptions.authorize ?? allowUser, - canAccess: (principal, threadId) => threadId.startsWith(principal.id), ...(handlerOptions.secret ? { media: { secret: handlerOptions.secret } } : {}), @@ -90,7 +99,9 @@ function setup(options: { media?: MediaOptions; secret?: string } = {}) { return { host, + harness, calls, + hostWith, handlerWith, call, post, @@ -139,12 +150,7 @@ describe('media upload and signed URLs', () => { }, }, }) - const options = { - host, - harness, - authorize: allowUser, - canAccess: () => true, - } + const options = { host, harness, openAccess: true as const } expect(() => createHarnessHandler(options)).toThrow(/media\.secret/) expect(() => createHarnessHandler({ ...options, media: { secret: 'shared' } }), @@ -175,7 +181,7 @@ describe('media upload and signed URLs', () => { expect(expiresAt).toBe(NOW + HOUR) expect(path).toMatch( new RegExp( - `^media\\?threadId=user-1-a&id=${record.id}&exp=${NOW + HOUR}&sig=[\\w-]+$`, + `^media\\?threadId=user-1-a&id=${record.id}&userId=user-1&exp=${NOW + HOUR}&sig=[\\w-]+$`, ), ) expect(response.status).toBe(200) @@ -183,16 +189,13 @@ describe('media upload and signed URLs', () => { await host.close() }) - it('serves a signed URL through a handler whose authorize refuses everything', async () => { - const { host, upload, signedPath, fetchSigned, handlerWith } = setup({ - secret: 'shared', - }) + it('serves a signed URL on another instance whose identify refuses everyone', async () => { + const { host, upload, signedPath, fetchSigned, hostWith, handlerWith } = + setup({ secret: 'shared' }) const record = await upload('user-1-a') const { path } = await signedPath(record) - const refuseAll = handlerWith({ - authorize: () => null, - secret: 'shared', - }) + const other = hostWith(() => null) + const refuseAll = handlerWith({ host: other, secret: 'shared' }) const signed = await fetchSigned(path, refuseAll) const unsigned = await fetchSigned( @@ -203,6 +206,7 @@ describe('media upload and signed URLs', () => { expect(signed.status).toBe(200) expect(await signed.text()).toBe('hello') expect(unsigned.status).toBe(401) + await other.close() await host.close() }) @@ -222,6 +226,14 @@ describe('media upload and signed URLs', () => { label: 'a later expiry', change: (query) => query.set('exp', String(Date.now() + 2 * HOUR)), }, + { + label: 'another user', + change: (query) => query.set('userId', 'user-1-admin'), + }, + { + label: 'another tenant', + change: (query) => query.set('tenantId', 'other-tenant'), + }, ] it.each(tampered)( @@ -270,7 +282,11 @@ describe('media upload and signed URLs', () => { }) it('serves an unsigned GET only to a principal that can open the thread', async () => { - const { host, call, upload, fetchSigned } = setup() + const { host, harness, call, upload, fetchSigned } = setup() + await host.open(harness, { + threadId: 'someone-else', + principal: { id: 'someone' }, + }) const record = await upload('user-1-a') const path = `media?threadId=user-1-a&id=${record.id}` diff --git a/packages/ai-harness/tests/input-sender.test.ts b/packages/ai-harness/tests/input-sender.test.ts index 5738ca77fb..90d4e7613f 100644 --- a/packages/ai-harness/tests/input-sender.test.ts +++ b/packages/ai-harness/tests/input-sender.test.ts @@ -12,6 +12,7 @@ import { defineCommand, defineHarness, definePlugin, + logIdOf, } from '../src' import { credentialsFor } from '../src/auth' import { after, gate, mockAdapter, text, toolCall } from './helpers' @@ -19,16 +20,22 @@ import type { LogRecord } from '@tanstack/ai-persistence' import type { HarnessPlugin, HarnessRouterContext, + HarnessRunContext, JoinCandidate, Operation, TurnInfo, } from '../src' const THREAD = 't1' +/** The id of the log of `THREAD` in the store, and its key in the other stores. */ +const LOG = logIdOf({ logId: THREAD }) const alice = { id: 'alice' } const bob = { id: 'bob' } -/** A tool that keeps who runs it and the context it gets. */ +/** + * A tool that keeps the user of its run, and the context it gets without + * the run. + */ function whoami() { const seen: Array<{ principal: string | undefined; context: unknown }> = [] const tool = toolDefinition({ @@ -38,13 +45,11 @@ function whoami() { }) const plugin = definePlugin({ name: 'test/whoami', - setup: (ctx) => ({ + setup: () => ({ tools: [ - tool.server((_input, toolContext) => { - seen.push({ - principal: ctx.session.principal?.id, - context: toolContext?.context, - }) + tool.server((_input, toolContext) => { + const { run, ...context } = toolContext.context + seen.push({ principal: run.principal?.id, context }) return Promise.resolve({}) }), ], @@ -78,16 +83,16 @@ describe('the sender of each input', () => { threadId: THREAD, principal: alice, }) - expect(await host.open(harness, { threadId: THREAD, principal: bob })).toBe( - session, - ) + // Bob's open reaches the same live session. Without host access, a call + // may name another sender. + await host.open(harness, { threadId: THREAD, principal: bob }) await session.send('Who am I?', { inputId: 'in-bob', principal: bob }) // Without a principal, the input is the opener's. await session.send('And now?', { inputId: 'in-default' }) expect(seen.map((entry) => entry.principal)).toEqual(['bob', 'alice']) - const inputs = (await persistence.stores.log.read({ logId: THREAD })) + const inputs = (await persistence.stores.log.read({ logId: LOG })) .map((entry): LogRecord => entry.record) .filter((record) => record.type === 'harness.input') .map((record) => [record['inputId'], record['principal']]) @@ -350,8 +355,11 @@ describe('the sender of each input', () => { name: 'test/starter', setup: (ctx) => ({ tools: [ - startWriter.server(() => { - ctx.agents.start(writer, undefined, { wake: true }) + startWriter.server((_input, toolContext) => { + ctx.agents.start(writer, undefined, { + wake: true, + principal: toolContext.context.run.principal, + }) return Promise.resolve({}) }), ], @@ -411,8 +419,10 @@ describe('the sender of each input', () => { name: 'test/starter', setup: (ctx) => ({ tools: [ - startWriter.server(() => { - started = ctx.agents.start(writer) + startWriter.server((_input, toolContext) => { + started = ctx.agents.start(writer, undefined, { + principal: toolContext.context.run.principal, + }) return Promise.resolve({}) }), ], @@ -482,7 +492,7 @@ describe('recovery keeps the sender and the context', () => { it('runs an input that never ran with its sender and context', async () => { const persistence = durablePersistence() await persistence.stores.log.append({ - logId: THREAD, + logId: LOG, seq: 1, epoch: UNOWNED_EPOCH, records: [ @@ -531,7 +541,7 @@ describe('recovery keeps the sender and the context', () => { it('runs a turn whose host stopped again with its message, sender, and context', async () => { const persistence = durablePersistence() await persistence.stores.log.append({ - logId: THREAD, + logId: LOG, seq: 1, epoch: UNOWNED_EPOCH, records: [ @@ -557,7 +567,7 @@ describe('recovery keeps the sender and the context', () => { }) await persistence.stores.runs.createOrResume({ runId: 'op-stopped', - threadId: THREAD, + threadId: LOG, startedAt: Date.now() - 60_000, }) await persistence.stores.runs.update('op-stopped', { @@ -650,6 +660,11 @@ describe('the turn where the model is picked', () => { message: 'Pick for me.', context: { plan: 'pro' }, principal: bob, + run: expect.objectContaining({ + principal: { id: 'bob' }, + runId: turn.id, + threadId: THREAD, + }), } expect(setups).toEqual([expected]) expect(picks).toEqual([expected]) @@ -710,10 +725,11 @@ describe('the tenant of a principal', () => { }) const plugin = definePlugin({ name: 'test/token', - setup: (ctx) => ({ + setup: () => ({ tools: [ - read.server(async () => { - const credential = await ctx.credentials.get('github') + read.server(async (_input, toolContext) => { + const credential = + await toolContext.context.run.credentials.get('github') seen.push(credential?.type === 'api_key' ? credential.value : null) return {} }), diff --git a/packages/ai-harness/tests/interrupt-resolve.test.ts b/packages/ai-harness/tests/interrupt-resolve.test.ts index 7c4fded224..0b42b5b9f5 100644 --- a/packages/ai-harness/tests/interrupt-resolve.test.ts +++ b/packages/ai-harness/tests/interrupt-resolve.test.ts @@ -129,7 +129,10 @@ describe.each([ expect(await session.settled('approve')).toMatchObject({ outcome: 'completed', }) - expect(seen).toEqual([{ screen: 'settings' }]) + // The harness adds the run of the turn to the context. + expect(seen).toEqual([ + { screen: 'settings', run: expect.objectContaining({ threadId: 't1' }) }, + ]) await second.host.close() }) diff --git a/packages/ai-harness/tests/interrupt-restart.test.ts b/packages/ai-harness/tests/interrupt-restart.test.ts index defca3cea6..9e2a9352da 100644 --- a/packages/ai-harness/tests/interrupt-restart.test.ts +++ b/packages/ai-harness/tests/interrupt-restart.test.ts @@ -2,13 +2,15 @@ import { describe, expect, it, vi } from 'vitest' import { z } from 'zod' import { defineAgent, toolDefinition } from '@tanstack/ai' import { memoryLogStore, memoryPersistence } from '@tanstack/ai-persistence' -import { createHarnessHost, defineHarness } from '../src' +import { createHarnessHost, defineHarness, logIdOf } from '../src' import { gate, mockAdapter, text, toolCall } from './helpers' import type { AnyTool } from '@tanstack/ai' import type { AnyAgent, HarnessRouting, HarnessSubagents } from '../src' import type { Reply } from './helpers' const THREAD = 't1' +/** The key of `THREAD` in the stores, in the default tenant. */ +const THREAD_KEY = logIdOf({ logId: THREAD }) /** A tool that asks approval before it removes a file. */ function removeTool() { @@ -290,7 +292,7 @@ describe.each([false, true])( const initial = await session.send('remove') const approval = initial.interrupts?.[0] if (!approval) throw new Error('The turn did not stop for approval.') - const records = await stores.interrupts.listPending(THREAD) + const records = await stores.interrupts.listPending(THREAD_KEY) const record = records[0] if (!record) throw new Error('The pending record is missing.') const read = vi @@ -630,7 +632,7 @@ describe.each([false, true])( 'pending', ) expect( - await stores.metadata.get('harness:interrupted', THREAD), + await stores.metadata.get('harness:interrupted', THREAD_KEY), ).toMatchObject({ interrupts: [expect.objectContaining({ id: approval.id })], }) @@ -797,7 +799,7 @@ describe.each([false, true])( await stores.interrupts.create({ interruptId: 'external', runId: 'external-run', - threadId: THREAD, + threadId: THREAD_KEY, requestedAt: Date.now(), payload: { id: 'external', diff --git a/packages/ai-harness/tests/leases.test.ts b/packages/ai-harness/tests/leases.test.ts index adbf641c38..1fcac7d6d8 100644 --- a/packages/ai-harness/tests/leases.test.ts +++ b/packages/ai-harness/tests/leases.test.ts @@ -1,7 +1,7 @@ import { describe, expect, it, vi } from 'vitest' import { memoryLogStore, memoryPersistence } from '@tanstack/ai-persistence' -import { createHarnessHost, defineHarness } from '../src' -import { defineAgent, toolDefinition } from '@tanstack/ai' +import { createHarnessHost, defineHarness, logIdOf } from '../src' +import { EventType, defineAgent, toolDefinition } from '@tanstack/ai' import { z } from 'zod' import { CUT_OFF_NOTE, INTERRUPTED_TOOL_RESULT } from '../src/resume' import { @@ -15,7 +15,6 @@ import { wrapLogStore, } from './helpers' import type { AnyTool, StreamChunk } from '@tanstack/ai' -import { EventType } from '@tanstack/ai' import type { LeaseStore, TurnLease, @@ -25,6 +24,8 @@ import type { AnyAgent, HarnessConfig } from '../src' import type { Reply } from './helpers' const THREAD = 't1' +/** The id of the log of `THREAD` in the store. */ +const LOG = logIdOf({ logId: THREAD }) /** A background agent that runs until it is cancelled. */ const waiter = defineAgent({ @@ -340,7 +341,7 @@ describe('a durable host with a lease store and no run store', () => { await mailStarted.opened await vi.waitFor(async () => expect( - (await persistence.stores.log.read({ logId: THREAD })).filter( + (await persistence.stores.log.read({ logId: LOG })).filter( (entry) => entry.record.type === 'harness.tool.result', ), ).toHaveLength(1), @@ -456,7 +457,7 @@ describe('close({ recoverable: true })', () => { } const logRecords = async (persistence: ReturnType) => - (await persistence.stores.log.read({ logId: THREAD })).map( + (await persistence.stores.log.read({ logId: LOG })).map( (entry) => entry.record, ) @@ -638,7 +639,7 @@ describe('durability.continueCutOff', () => { ) await vi.waitFor(async () => expect( - JSON.stringify(await persistence.stores.log.read({ logId: THREAD })), + JSON.stringify(await persistence.stores.log.read({ logId: LOG })), ).toContain(streamed), ) leases.expire() @@ -681,7 +682,7 @@ describe('durability.continueCutOff', () => { ['assistant', '42.'], ]) // One append, so the transcript never ends with the cut answer. - const added = (await persistence.stores.log.read({ logId: THREAD })).find( + const added = (await persistence.stores.log.read({ logId: LOG })).find( (entry) => entry.record.type === 'harness.transcript' && JSON.stringify(entry.record).includes(note), @@ -715,7 +716,7 @@ describe('durability.continueCutOff', () => { ['user', two], ] expect(roles(next.calls[0].messages)).toEqual([['user', 'go'], ...added]) - const record = (await persistence.stores.log.read({ logId: THREAD })).find( + const record = (await persistence.stores.log.read({ logId: LOG })).find( (entry) => entry.record.type === 'harness.transcript' && JSON.stringify(entry.record).includes(one), diff --git a/packages/ai-harness/tests/log.test.ts b/packages/ai-harness/tests/log.test.ts index 630b2eb9bd..5da87c7804 100644 --- a/packages/ai-harness/tests/log.test.ts +++ b/packages/ai-harness/tests/log.test.ts @@ -12,6 +12,7 @@ import { emptySharedLogState, engineMessageStore, loadLogState, + logIdOf, logMessageStore, sessionMessageStore, sessionOf, @@ -27,6 +28,8 @@ import type { ProjectOptions, ReduceOptions } from '../src/log' import type { SessionEvent } from '../src/types' const THREAD = 't1' +/** The id of the log of `THREAD` in the store. */ +const LOG = logIdOf({ logId: THREAD }) function newWriter( store: LogStore, @@ -40,7 +43,8 @@ function newWriter( ) { const log = new SharedLog({ store, - logId: THREAD, + logId: LOG, + rootThread: THREAD, state: options.state ?? emptySharedLogState(), coalesceMs: options.coalesceMs ?? 0, ...(options.project ? { project: options.project } : {}), @@ -51,8 +55,12 @@ function newWriter( /** The fold of the session `THREAD`. */ const loadSession = async ( - options: Omit[0], 'logId'>, -) => sessionOf(await loadLogState({ ...options, logId: THREAD }), THREAD) + options: Omit[0], 'logId' | 'rootThread'>, +) => + sessionOf( + await loadLogState({ ...options, logId: LOG, rootThread: THREAD }), + THREAD, + ) /** A store that records the record types of each append batch. */ function recordingStore(inner: LogStore = memoryLogStore()) { @@ -71,7 +79,7 @@ function silentStore(inner: LogStore) { return wrapLogStore(inner, { subscribe: () => () => {} }) } -const records = async (store: LogStore, logId = THREAD) => +const records = async (store: LogStore, logId = LOG) => (await store.read({ logId })).map((entry) => entry.record) const user = (id: string, content: string): ModelMessage => ({ @@ -248,7 +256,7 @@ describe('fold robustness', () => { }) for (let seq = 1; seq <= 5; seq += 1) { await inner.append({ - logId: THREAD, + logId: LOG, seq, epoch: UNOWNED_EPOCH, records: [{ type: 'app.signal', text: `s${seq}` }], @@ -338,7 +346,7 @@ describe('input fold', () => { retries, }) await store.append({ - logId: THREAD, + logId: LOG, seq: 1, epoch: UNOWNED_EPOCH, records: [ @@ -355,14 +363,14 @@ describe('input fold', () => { ) await store.append({ - logId: THREAD, + logId: LOG, seq: 2, epoch: UNOWNED_EPOCH, records: [retry(1), retry(2)], }) expect((await loadSession({ store })).inputs.get('in-1')?.retries).toBe(2) await store.append({ - logId: THREAD, + logId: LOG, seq: 4, epoch: UNOWNED_EPOCH, records: [retry(0)], @@ -376,7 +384,8 @@ describe('a closed view', () => { const store = memoryLogStore() const log = new SharedLog({ store, - logId: THREAD, + logId: LOG, + rootThread: THREAD, state: emptySharedLogState(), coalesceMs: 0, }) @@ -448,7 +457,7 @@ describe('events and batches', () => { writer.flush(), ]) - const entries = await store.read({ logId: THREAD }) + const entries = await store.read({ logId: LOG }) expect(entries.map((entry) => [entry.seq, entry.record.type])).toEqual([ [1, 'harness.event'], [2, 'app.a'], @@ -469,7 +478,7 @@ describe('events and batches', () => { first.close() // A new session: the first events are only in the store. - const state = await loadLogState({ store, logId: THREAD }) + const state = await loadLogState({ store, logId: LOG, rootThread: THREAD }) const writer = newWriter(store, { state }) let done = false const seen: Array = [] @@ -516,7 +525,7 @@ describe('write failures', () => { const failures: Array = [] const first = newWriter(store, { onFailure: (e) => failures.push(e) }) await first.append([{ type: 'app.mine' }]) - const state = await loadLogState({ store, logId: THREAD }) + const state = await loadLogState({ store, logId: LOG, rootThread: THREAD }) const next = newWriter(store, { state }) // The newer host writes. The first host is notified and stops. @@ -539,12 +548,12 @@ describe('write failures', () => { await old.commit({ messages: [user('u1', 'one')] }) // A newer host appends a harness record type that this version does not know. await store.append({ - logId: THREAD, + logId: LOG, seq: 2, epoch: UNOWNED_EPOCH, records: [{ type: 'harness.later', value: 1 }], }) - const state = await loadLogState({ store, logId: THREAD }) + const state = await loadLogState({ store, logId: LOG, rootThread: THREAD }) expect(state.newerRecord).toBe('harness.later') expect(sessionOf(state, THREAD).messages).toEqual([user('u1', 'one')]) @@ -573,7 +582,7 @@ describe('write failures', () => { onFailure: (error) => failures.push(error), }) await store.append({ - logId: THREAD, + logId: LOG, seq: 1, epoch: UNOWNED_EPOCH, records: [{ type: 'app.from-another-host' }], @@ -627,7 +636,7 @@ describe('write failures', () => { isDown = true await inner.append({ - logId: THREAD, + logId: LOG, seq: 1, epoch: UNOWNED_EPOCH, records: [{ type: 'app.other' }], @@ -718,7 +727,7 @@ describe('fold checkpoints', () => { const store = memoryLogStore() await seedLog(store) const { metadata } = memoryPersistence().stores - await metadata.set(NAMESPACE, THREAD, checkpoint({})) + await metadata.set(NAMESPACE, LOG, checkpoint({})) const state = await loadSession({ store, metadata, project }) @@ -743,7 +752,7 @@ describe('fold checkpoints', () => { const store = memoryLogStore() await seedLog(store) const { metadata } = memoryPersistence().stores - await metadata.set(NAMESPACE, THREAD, value) + await metadata.set(NAMESPACE, LOG, value) const state = await loadSession({ store, metadata, project }) @@ -767,7 +776,7 @@ describe('fold checkpoints', () => { } for (let index = 0; index < 51; index += 1) await writer.append([usage]) await expect - .poll(() => metadata.get(NAMESPACE, THREAD)) + .poll(() => metadata.get(NAMESPACE, LOG)) .toMatchObject({ seq: 50 }) await loadSession({ store, metadata }) @@ -786,7 +795,7 @@ describe('fold checkpoints', () => { } await expect - .poll(() => metadata.get(NAMESPACE, THREAD)) + .poll(() => metadata.get(NAMESPACE, LOG)) .toMatchObject({ v: 2, seq: 50, version: 'v1' }) }) @@ -799,14 +808,14 @@ describe('fold checkpoints', () => { await writer.append([{ type: 'app.tick', index }]) } await expect - .poll(() => metadata.get(NAMESPACE, THREAD)) + .poll(() => metadata.get(NAMESPACE, LOG)) .toMatchObject({ seq: 50 }) expect( (await loadSession({ store, metadata, project })).transcriptSeq, ).toBe(1) // A checkpoint from before the field does not know it. - await metadata.set(NAMESPACE, THREAD, checkpoint({ seq: 2 })) + await metadata.set(NAMESPACE, LOG, checkpoint({ seq: 2 })) expect( (await loadSession({ store, metadata, project })).transcriptSeq, ).toBeUndefined() @@ -830,23 +839,24 @@ describe('fold checkpoints', () => { async (_name, value, reduce) => { const store = memoryLogStore() await store.append({ - logId: THREAD, + logId: LOG, seq: 1, epoch: UNOWNED_EPOCH, records: [{ type: 'app.tick' }], }) await store.append({ - logId: THREAD, + logId: LOG, seq: 2, epoch: UNOWNED_EPOCH, records: [{ type: 'app.tick' }], }) const { metadata } = memoryPersistence().stores - await metadata.set(NAMESPACE, THREAD, value) + await metadata.set(NAMESPACE, LOG, value) const state = await loadLogState({ store, - logId: THREAD, + logId: LOG, + rootThread: THREAD, metadata, project, reduce, @@ -870,7 +880,8 @@ describe('fold checkpoints', () => { } const log = new SharedLog({ store, - logId: THREAD, + logId: LOG, + rootThread: THREAD, state: emptySharedLogState(reduce), coalesceMs: 0, reduce, @@ -881,11 +892,17 @@ describe('fold checkpoints', () => { await writer.append([{ type: 'app.tick', index }]) } await expect - .poll(() => metadata.get(NAMESPACE, THREAD)) + .poll(() => metadata.get(NAMESPACE, LOG)) .toMatchObject({ seq: 50, reduceVersion: '' }) calls = 0 - const state = await loadLogState({ store, logId: THREAD, metadata, reduce }) + const state = await loadLogState({ + store, + logId: LOG, + rootThread: THREAD, + metadata, + reduce, + }) expect(state.reduced).toBe(50) expect(state.reduced).toEqual(log.state.reduced) @@ -1008,6 +1025,15 @@ describe('engine message store', () => { }) }) +describe('logIdOf', () => { + it('encodes the tenant and the log id, so no two pairs give one id', () => { + expect(logIdOf({ logId: 't1' })).toBe(':t1') + expect(logIdOf({ tenantId: 'org/1', logId: 'a:b' })).toBe('org%2F1:a%3Ab') + expect(logIdOf({ tenantId: 'a:b', logId: 'c' })).toBe('a%3Ab:c') + expect(logIdOf({ tenantId: 'a', logId: 'b:c' })).toBe('a:b%3Ac') + }) +}) + describe('message stores', () => { it('logMessageStore writes a change as one record and reads the fold', async () => { const store = memoryLogStore() @@ -1019,7 +1045,7 @@ describe('message stores', () => { await messages.saveThread('t9', [u1, a1]) expect(await messages.loadThread('t9')).toEqual([u1, a1]) - expect(await records(store, 't9')).toEqual([ + expect(await records(store, logIdOf({ logId: 't9' }))).toEqual([ { type: 'harness.transcript', keep: 0, add: [u1] }, { type: 'harness.transcript', keep: 1, add: [a1] }, ]) @@ -1043,10 +1069,11 @@ describe('message stores', () => { user('u1', 'hi'), ]) - expect((await records(inner, 't9')).map((record) => record.type)).toEqual([ - 'app.other', - 'harness.transcript', - ]) + expect( + (await records(inner, logIdOf({ logId: 't9' }))).map( + (record) => record.type, + ), + ).toEqual(['app.other', 'harness.transcript']) }) it('sessionMessageStore uses the writer for its thread and the log for others', async () => { @@ -1062,16 +1089,22 @@ describe('message stores', () => { expect(await messages.loadThread('subagent:1')).toEqual([ user('c1', 'child'), ]) - const child: Array = await records(store, 'subagent:1') + const child: Array = await records( + store, + logIdOf({ logId: 'subagent:1' }), + ) expect(child).toHaveLength(1) }) }) describe('a log shared by sessions', () => { + /** The id in the store of the log whose root thread is `log-1`. */ + const SHARED = logIdOf({ logId: 'log-1' }) const sharedWriter = (store: LogStore) => { const log = new SharedLog({ store, - logId: 'log-1', + logId: SHARED, + rootThread: 'log-1', state: emptySharedLogState(), coalesceMs: 0, }) @@ -1088,7 +1121,7 @@ describe('a log shared by sessions', () => { await root.append([{ type: 'app.note', text: 'one' }]) - expect(await records(store, 'log-1')).toEqual([ + expect(await records(store, SHARED)).toEqual([ { type: 'app.note', text: 'one' }, ]) }) @@ -1103,7 +1136,7 @@ describe('a log shared by sessions', () => { { type: 'app.child_linked', child: 'child' }, ]) - expect(await records(store, 'log-1')).toEqual([ + expect(await records(store, SHARED)).toEqual([ { type: 'app.note', text: 'from the child', thread: 'child' }, { type: 'app.child_created', thread: 'child' }, { type: 'app.child_linked', child: 'child' }, @@ -1119,7 +1152,11 @@ describe('a log shared by sessions', () => { expect(root.state.messages).toEqual([user('r1', 'root message')]) expect(child.state.messages).toEqual([user('c1', 'child message')]) - const state = await loadLogState({ store, logId: 'log-1' }) + const state = await loadLogState({ + store, + logId: SHARED, + rootThread: 'log-1', + }) expect(sessionOf(state, 'child').messages).toEqual([ user('c1', 'child message'), ]) @@ -1142,7 +1179,8 @@ describe('a log shared by sessions', () => { } const log = new SharedLog({ store, - logId: 'log-1', + logId: SHARED, + rootThread: 'log-1', state: emptySharedLogState(reduce), coalesceMs: 0, reduce, @@ -1155,7 +1193,14 @@ describe('a log shared by sessions', () => { expect(log.state.reduced).toEqual({ x: 2 }) expect( - (await loadLogState({ store, logId: 'log-1', reduce })).reduced, + ( + await loadLogState({ + store, + logId: SHARED, + rootThread: 'log-1', + reduce, + }) + ).reduced, ).toEqual({ x: 2, }) @@ -1190,7 +1235,7 @@ describe('a log shared by sessions', () => { await vi.waitFor(() => expect(log.isOpen).toBe(false)) expect(() => log.view('log-1', () => {})).toThrow( - 'The log "log-1" is closed.', + `The log ${JSON.stringify(SHARED)} is closed.`, ) }) @@ -1199,7 +1244,8 @@ describe('a log shared by sessions', () => { const onIdle = vi.fn() const log = new SharedLog({ store, - logId: 'log-1', + logId: SHARED, + rootThread: 'log-1', state: emptySharedLogState(), coalesceMs: 0, onIdle, @@ -1212,7 +1258,7 @@ describe('a log shared by sessions', () => { // Another writer takes the next position. await store.append({ - logId: 'log-1', + logId: SHARED, seq: 2, epoch: UNOWNED_EPOCH, records: [{ type: 'app.other' }], @@ -1245,7 +1291,7 @@ describe('a log shared by sessions', () => { // The child view keeps the log open, so the thread opens again. const again = log.view('log-1', () => {}) await again.append([{ type: 'app.note' }]) - expect(await records(store, 'log-1')).toEqual([{ type: 'app.note' }]) + expect(await records(store, SHARED)).toEqual([{ type: 'app.note' }]) }) it('stays open until the last append of a closed view lands', async () => { @@ -1262,7 +1308,8 @@ describe('a log shared by sessions', () => { }) const log = new SharedLog({ store, - logId: THREAD, + logId: LOG, + rootThread: THREAD, state: emptySharedLogState(), coalesceMs: 0, }) @@ -1290,8 +1337,8 @@ describe('a log shared by sessions', () => { const next = log.view('child', () => {}) await next.commit({ messages: [user('c1', 'hi')] }) - expect( - (await records(store, 'log-1')).map((record) => record.type), - ).toEqual(['harness.transcript']) + expect((await records(store, SHARED)).map((record) => record.type)).toEqual( + ['harness.transcript'], + ) }) }) diff --git a/packages/ai-harness/tests/media.test.ts b/packages/ai-harness/tests/media.test.ts index e9354b787b..b21e29123c 100644 --- a/packages/ai-harness/tests/media.test.ts +++ b/packages/ai-harness/tests/media.test.ts @@ -9,6 +9,7 @@ import { toBase64, } from '../src/media' import { mediaPart } from '../src/media-ref' +import { logIdOf } from '../src/log' import { mockAdapter, text } from './helpers' import type { ContentPart, @@ -138,7 +139,11 @@ describe('createMediaStore', () => { await store .put(new Blob(['hello']).stream(), { mimeType: 'image/png', name: 'a' }) .catch(() => undefined) - expect(await persistence.stores.artifacts.listForThread('t1')).toEqual([]) + expect( + await persistence.stores.artifacts.listForThread( + logIdOf({ logId: 't1' }), + ), + ).toEqual([]) expect((await persistence.stores.blobs.list()).objects).toEqual([]) }) @@ -153,10 +158,13 @@ describe('createMediaStore', () => { ).rejects.toMatchObject({ status: 415 }) }) - it('does not show media of another thread', async () => { + it.each([ + ['another thread', { threadId: 'b' }], + ['the same thread id in another tenant', { threadId: 'a', tenantId: 'x' }], + ])('does not show media of %s', async (_label, other) => { const persistence = memoryPersistence() const a = createMediaStore({ persistence, threadId: 'a' }) - const b = createMediaStore({ persistence, threadId: 'b' }) + const b = createMediaStore({ persistence, ...other }) const record = await a.put(encoder.encode('secret'), { mimeType: 'image/png', name: 'a.png', @@ -173,7 +181,22 @@ describe('createMediaStore', () => { name: 'cat.png', }) const blobKey = (await persistence.stores.artifacts.get(record.id))?.blobKey - expect(blobKey).toBe(`artifacts/upload:..%2Fx/${record.id}`) + expect(blobKey).toBe(`artifacts/upload::..%2Fx/${record.id}`) + }) + + it('puts the tenant in the upload folder of the blob key', async () => { + const persistence = memoryPersistence() + const store = createMediaStore({ + persistence, + threadId: 't1', + tenantId: 'org/1', + }) + const record = await store.put(encoder.encode('hello'), { + mimeType: 'image/png', + name: 'cat.png', + }) + const blobKey = (await persistence.stores.artifacts.get(record.id))?.blobKey + expect(blobKey).toBe(`artifacts/upload:org%2F1:t1/${record.id}`) }) it('throws 404 for an id that does not exist', async () => { diff --git a/packages/ai-harness/tests/members-only.test.ts b/packages/ai-harness/tests/members-only.test.ts new file mode 100644 index 0000000000..355d81862a --- /dev/null +++ b/packages/ai-harness/tests/members-only.test.ts @@ -0,0 +1,252 @@ +import { describe, expect, it } from 'vitest' +import { membersOnly, membersOnlyPolicyOf } from '../src/access' +import type { SessionMember } from '@tanstack/ai-persistence' +import type { + AccessAction, + AccessRun, + AccessSession, + HarnessAccess, + MembersOnlyOptions, +} from '../src/access' +import type { Principal } from '../src/types' + +const authorize = membersOnly() + +const members: ReadonlyArray = [ + { userId: 'mia', role: 'member' }, + { userId: 'vic', role: 'viewer' }, +] + +/** `olga` owns it through `ownerId` only: `members` has no entry for her. */ +const session: AccessSession = { + threadId: 'thread-1', + ownerId: 'olga', + members, +} + +/** `otto` owns it through `members` only: it has no `ownerId`. */ +const sessionOwnedByMembers: AccessSession = { + threadId: 'thread-2', + members: [{ userId: 'otto', role: 'owner' }, ...members], +} + +type Column = 'owner' | 'member' | 'viewer' | 'stranger' + +/** The roles table of the design. */ +const table: Array<{ action: AccessAction } & Record> = [ + { action: 'read', owner: true, member: true, viewer: true, stranger: false }, + { action: 'list', owner: true, member: true, viewer: true, stranger: false }, + { action: 'send', owner: true, member: true, viewer: false, stranger: false }, + { + action: 'answer', + owner: true, + member: true, + viewer: false, + stranger: false, + }, + { + action: 'cancel', + owner: true, + member: true, + viewer: false, + stranger: false, + }, + { + action: 'configure', + owner: true, + member: false, + viewer: false, + stranger: false, + }, + { + action: 'share', + owner: true, + member: false, + viewer: false, + stranger: false, + }, + { + action: 'delete', + owner: true, + member: false, + viewer: false, + stranger: false, + }, + { action: 'fork', owner: true, member: true, viewer: false, stranger: false }, + { + action: 'transfer', + owner: true, + member: false, + viewer: false, + stranger: false, + }, +] + +const callers: Array<{ + who: string + id: string + target: AccessSession + column: Column +}> = [ + { who: 'owner by ownerId', id: 'olga', target: session, column: 'owner' }, + { + who: 'owner by members', + id: 'otto', + target: sessionOwnedByMembers, + column: 'owner', + }, + { who: 'member', id: 'mia', target: session, column: 'member' }, + { who: 'viewer', id: 'vic', target: session, column: 'viewer' }, + { who: 'stranger', id: 'sam', target: session, column: 'stranger' }, +] + +const cases = table.flatMap((row) => + callers.map((caller) => ({ + ...caller, + action: row.action, + allowed: row[caller.column], + })), +) + +describe('membersOnly roles', () => { + // Each caller targets a run that the caller started, so `answer` and + // `cancel` follow the role alone here. + it.each(cases)( + '$who: $action is $allowed', + async ({ id, target, action, allowed }) => { + const run = { runId: 'run-1', principalId: id } + const decision = await authorize({ + principal: { id }, + action, + session: target, + run, + }) + expect(decision).toBe(allowed) + }, + ) +}) + +describe('membersOnly runs for answer and cancel', () => { + const runCases: Array<{ + who: string + id: string + target: string + run: AccessRun | undefined + allowed: boolean + }> = [ + { + who: 'member', + id: 'mia', + target: 'a run of another user', + run: { runId: 'run-2', principalId: 'olga' }, + allowed: false, + }, + { + who: 'member', + id: 'mia', + target: 'a run with no principal', + run: { runId: 'run-3' }, + allowed: false, + }, + { + who: 'member', + id: 'mia', + target: 'no run', + run: undefined, + allowed: false, + }, + { + who: 'owner', + id: 'olga', + target: 'a run of another user', + run: { runId: 'run-4', principalId: 'mia' }, + allowed: true, + }, + { + who: 'owner', + id: 'olga', + target: 'no run', + run: undefined, + allowed: true, + }, + ] + const actions: Array = ['answer', 'cancel'] + const rows = actions.flatMap((action) => + runCases.map((runCase) => ({ ...runCase, action })), + ) + + it.each(rows)( + '$who: $action on $target is $allowed', + async ({ id, action, run, allowed }) => { + const decision = await authorize({ + principal: { id }, + action, + session, + run, + }) + expect(decision).toBe(allowed) + }, + ) +}) + +describe('membersOnly marker', () => { + const defaults: Array<{ + name: string + options: MembersOnlyOptions | undefined + }> = [ + { name: 'no options', options: undefined }, + { name: 'no fork option', options: {} }, + { name: 'a fork option without members', options: { fork: {} } }, + ] + + it.each(defaults)('sets fork members to self with $name', ({ options }) => { + expect(membersOnly(options).membersOnly).toEqual({ + fork: { members: 'self' }, + }) + }) + + it('sets fork members to all when the option is all', () => { + expect(membersOnly({ fork: { members: 'all' } }).membersOnly).toEqual({ + fork: { members: 'all' }, + }) + }) + + it('keeps a separate marker on each authorize', () => { + const self = membersOnly() + membersOnly({ fork: { members: 'all' } }) + expect(self.membersOnly).toEqual({ fork: { members: 'self' } }) + }) +}) + +describe('membersOnlyPolicyOf', () => { + it('reads the marker of a membersOnly authorize', () => { + const policy = membersOnlyPolicyOf( + membersOnly({ fork: { members: 'all' } }), + ) + expect(policy).toEqual({ fork: { members: 'all' } }) + }) + + it('is undefined for a plain authorize function', () => { + expect(membersOnlyPolicyOf(() => true)).toBeUndefined() + }) +}) + +interface AppUser extends Principal { + plan: 'free' | 'pro' +} + +describe('membersOnly with a wider principal type', () => { + it('works as the authorize of an access object', async () => { + const access: HarnessAccess = { + identify: () => null, + authorize: membersOnly(), + } + const principal: AppUser = { id: 'mia', plan: 'pro' } + const decision = await access.authorize({ + principal, + action: 'send', + session, + }) + expect(decision).toBe(true) + }) +}) diff --git a/packages/ai-harness/tests/oauth-edges.test.ts b/packages/ai-harness/tests/oauth-edges.test.ts index aa5c0cec33..58c13735dc 100644 --- a/packages/ai-harness/tests/oauth-edges.test.ts +++ b/packages/ai-harness/tests/oauth-edges.test.ts @@ -18,6 +18,7 @@ import { refreshCredential, } from '../src/oauth' import { mockAdapter, text, toolCall } from './helpers' +import type { HarnessRunContext } from '../src' import type { OAuthConfig } from '../src/oauth' const oauth: OAuthConfig = { @@ -215,12 +216,13 @@ describe('oauthConnector edges', () => { oauth, fetch: fetchFn, tools: (token) => [ - toolDefinition({ name: 'whoami', description: 'Who am I' }).server( - async () => { - seen.push(await token()) - return 'me' - }, - ), + toolDefinition({ + name: 'whoami', + description: 'Who am I', + }).server(async (_args, ctx) => { + seen.push(await token({ run: ctx.context.run })) + return 'me' + }), ], }) } diff --git a/packages/ai-harness/tests/protocol.test.ts b/packages/ai-harness/tests/protocol.test.ts index b44ec5fe53..ae7f843671 100644 --- a/packages/ai-harness/tests/protocol.test.ts +++ b/packages/ai-harness/tests/protocol.test.ts @@ -9,8 +9,15 @@ import { parseControlFrame, parseHarnessInput, } from '../src/server' -import { after, gate, mockAdapter, text, wakeMessages } from './helpers' -import type { WebSocketLike } from '@tanstack/ai' +import { + after, + fakeSocket, + gate, + mockAdapter, + readSse, + text, + wakeMessages, +} from './helpers' import type { HostFrame } from '../src/server' import type { Reply } from './helpers' @@ -36,20 +43,22 @@ function setup( agents: [pricer, secret], expose: { agents: ['pricer'] }, }) - const host = createHarnessHost({ persistence: memoryPersistence() }) - const handler = createHarnessHandler({ - host, - harness, - authorize: (request) => - request.headers.get('authorization') === 'Bearer good' - ? { id: 'user-1' } - : null, - canAccess: (principal, threadId) => threadId.startsWith(principal.id), + const host = createHarnessHost({ + persistence: memoryPersistence(), + access: { + identify: ({ request }) => + request.headers.get('authorization') === 'Bearer good' ? user : null, + // A user may use the threads whose id starts with the user id. + authorize: ({ principal, session }) => + session.threadId.startsWith(principal.id), + }, }) + const handler = createHarnessHandler({ host, harness }) return { handler, host, harness, calls } } const auth = { authorization: 'Bearer good' } +const user = { id: 'user-1' } type Handler = ReturnType['handler'] @@ -72,39 +81,6 @@ const isTurnEnd = (frame: HostFrame) => frame.event.type === 'CUSTOM' && frame.event.name === 'harness.operation.finished' -async function readSse( - response: Response, - until: (frame: HostFrame) => boolean, -) { - if (!response.body) throw new Error('The response has no body.') - const reader = response.body.getReader() - const decoder = new TextDecoder() - const frames: Array = [] - let buffer = '' - while (true) { - const { value, done } = await reader.read() - if (done) break - buffer += decoder.decode(value, { stream: true }) - const blocks = buffer.split('\n\n') - buffer = blocks.pop() ?? '' - for (const block of blocks) { - const data = block - .split('\n') - .find((line) => line.startsWith('data: ')) - ?.slice(6) - if (!data) continue - // The handler writes each frame as JSON. - const parsed: HostFrame = JSON.parse(data) - frames.push(parsed) - if (until(parsed)) { - await reader.cancel() - return frames - } - } - } - return frames -} - describe('HTTP handler', () => { it('refuses requests that authorize rejects', async () => { const { handler, host } = setup() @@ -116,7 +92,11 @@ describe('HTTP handler', () => { }) it('refuses threads the principal cannot access', async () => { - const { handler, host } = setup() + const { handler, host, harness } = setup() + await host.open(harness, { + threadId: 'other-thread', + principal: { id: 'user-2' }, + }) const response = await handler( new Request('http://x/api/harness/snapshot?threadId=other-thread', { headers: auth, @@ -271,7 +251,10 @@ describe('HTTP handler', () => { }), ).toMatchObject({ status: expect.stringMatching(/^(accepted|queued)$/) }) - const session = await host.open(harness, { threadId: 'user-1-m' }) + const session = await host.open(harness, { + threadId: 'user-1-m', + principal: user, + }) const hidden = session.agent('secret')?.start() expect( await control({ @@ -335,7 +318,10 @@ describe('HTTP handler', () => { await control({ op: 'followUp', message: 'I am real', wake: true }) await vi.waitFor(() => expect(calls).toHaveLength(1)) - const session = await host.open(harness, { threadId: 'user-1-wake' }) + const session = await host.open(harness, { + threadId: 'user-1-wake', + principal: user, + }) await vi.waitFor(() => expect(session.snapshot().status).toBe('idle')) expect(wakeMessages(await session.transcript())).toEqual([]) await host.close() @@ -395,23 +381,6 @@ describe('parseHarnessInput', () => { }) describe('WebSocket', () => { - function fakeSocket() { - const handlers: Record void>> = {} - const sent: Array = [] - const socket: WebSocketLike = { - send: (data) => sent.push(JSON.parse(data)), - close: () => handlers.close?.forEach((handler) => handler()), - addEventListener: (type: string, handler: (event?: any) => void) => { - ;(handlers[type] ??= []).push(handler) - }, - } - const receive = (frame: unknown) => - handlers.message?.forEach((handler) => - handler({ data: JSON.stringify(frame) }), - ) - return { socket, sent, receive } - } - it('subscribes, applies inputs with receipts, and streams events', async () => { const { host, harness } = setup() const { socket, sent, receive } = fakeSocket() @@ -419,8 +388,7 @@ describe('WebSocket', () => { host, harness, socket, - principal: { id: 'user-1' }, - canAccess: () => true, + request: new Request('http://x/api/harness/socket', { headers: auth }), }) receive({ diff --git a/packages/ai-harness/tests/provider-keys.test.ts b/packages/ai-harness/tests/provider-keys.test.ts index 76ecf68086..78a77d7f15 100644 --- a/packages/ai-harness/tests/provider-keys.test.ts +++ b/packages/ai-harness/tests/provider-keys.test.ts @@ -3,7 +3,12 @@ import { afterEach, describe, expect, it, vi } from 'vitest' import { EventType, defineAgent } from '@tanstack/ai' import { defineByokProvider, keyedAdapter } from '@tanstack/ai/byok' import { memoryPersistence } from '@tanstack/ai-persistence' -import { HARNESS_EVENTS, createHarnessHost, defineHarness } from '../src' +import { + HARNESS_EVENTS, + createHarnessHost, + defineHarness, + logIdOf, +} from '../src' import { compact, modelPicker, providerKeys } from '../src/first-party' import { mockAdapter, text, toolCall } from './helpers' import type { AnyTextAdapter } from '@tanstack/ai' @@ -446,11 +451,13 @@ describe('providerKeys', () => { await session.send('hi'), ] expect(keyed.keys).toEqual([PASTED_KEY]) + // The key of the thread `t` in the stores, in the default tenant. + const key = logIdOf({ logId: 't' }) const inbox = await persistence.stores.inbox.get({ - threadId: 't', + threadId: key, inputId: receipt.inputId, }) - const transcript = await persistence.stores.messages.loadThread('t') + const transcript = await persistence.stores.messages.loadThread(key) const everything = JSON.stringify({ events, results, inbox, transcript }) expect(everything).not.toContain(PASTED_KEY) expect(everything).toContain('...5678') diff --git a/packages/ai-harness/tests/remote-worker.test.ts b/packages/ai-harness/tests/remote-worker.test.ts index 1898d41b76..5b3142e8b2 100644 --- a/packages/ai-harness/tests/remote-worker.test.ts +++ b/packages/ai-harness/tests/remote-worker.test.ts @@ -24,15 +24,19 @@ const textOf = (chunks: Array) => describe('harnessText for a remote harness', () => { it('sends the token and the last user message, and streams the answer', async () => { const { adapter, calls } = mockAdapter([() => text('remote answer')]) - const host = createHarnessHost({ persistence: memoryPersistence() }) + const host = createHarnessHost({ + persistence: memoryPersistence(), + access: { + identify: ({ request }) => + request.headers.get('authorization') === 'Bearer secret' + ? { id: 'u' } + : null, + authorize: () => true, + }, + }) const handler = createHarnessHandler({ host, harness: defineHarness({ name: 'test/remote', adapter }), - canAccess: () => true, - authorize: (request) => - request.headers.get('authorization') === 'Bearer secret' - ? { id: 'u' } - : null, }) const remote = harnessText({ url: 'http://remote.test/api/harness/', diff --git a/packages/ai-harness/tests/replay-adapter-parity.test.ts b/packages/ai-harness/tests/replay-adapter-parity.test.ts index 4bf661c296..eae3e2a577 100644 --- a/packages/ai-harness/tests/replay-adapter-parity.test.ts +++ b/packages/ai-harness/tests/replay-adapter-parity.test.ts @@ -2,7 +2,7 @@ import { describe, expect, it, vi } from 'vitest' import { EventType, toolDefinition } from '@tanstack/ai' import { memoryLogStore, memoryPersistence } from '@tanstack/ai-persistence' import { z } from 'zod' -import { createHarnessHost, defineHarness } from '../src' +import { createHarnessHost, defineHarness, logIdOf } from '../src' import { gate, mockAdapter, text, toolCall } from './helpers' import type { MessageSource, StreamChunk } from '@tanstack/ai' @@ -107,7 +107,9 @@ describe('saved harness replay across turn adapters', () => { JSON.stringify((await session.transcript()).slice(0, before.length)), ), ).toEqual(before) - const saved = await persistence.stores.messages.loadThread('replay') + const saved = await persistence.stores.messages.loadThread( + logIdOf({ logId: 'replay' }), + ) expect(saved.at(-1)?.metadata?.tanstack).toMatchObject({ source: secondSource, responseId: 'generation-second', diff --git a/packages/ai-harness/tests/reset.test.ts b/packages/ai-harness/tests/reset.test.ts index ccb081c37e..10f852e9bc 100644 --- a/packages/ai-harness/tests/reset.test.ts +++ b/packages/ai-harness/tests/reset.test.ts @@ -6,7 +6,7 @@ import { memoryLogStore, memoryPersistence, } from '@tanstack/ai-persistence' -import { createHarnessHost, defineHarness } from '../src' +import { createHarnessHost, defineHarness, logIdOf } from '../src' import { after, gate, @@ -294,7 +294,7 @@ describe('reset on a durable host', () => { const { runs, metadata } = memoryPersistence().stores const persistence = { stores: { log, runs, metadata } } await log.append({ - logId: 't1', + logId: logIdOf({ logId: 't1' }), seq: 1, epoch: UNOWNED_EPOCH, records: [ diff --git a/packages/ai-harness/tests/resume-edges.test.ts b/packages/ai-harness/tests/resume-edges.test.ts index ec58747a74..8e33bc9a95 100644 --- a/packages/ai-harness/tests/resume-edges.test.ts +++ b/packages/ai-harness/tests/resume-edges.test.ts @@ -2,7 +2,7 @@ import { describe, expect, it, vi } from 'vitest' import { z } from 'zod' import { EventType, chat, defineAgent } from '@tanstack/ai' import { memoryPersistence } from '@tanstack/ai-persistence' -import { createHarnessHost, defineHarness, harnessText } from '../src' +import { createHarnessHost, defineHarness, harnessText, logIdOf } from '../src' import { INTERRUPTED_TOOL_RESULT, findCrashedRuns, @@ -11,6 +11,9 @@ import { import { mockAdapter, text, untilAborted } from './helpers' import type { RunRecord, RunStore, StreamChunk } from '@tanstack/ai' +/** The key of the thread `t1` in the stores of a session, in the default tenant. */ +const T1_KEY = logIdOf({ logId: 't1' }) + async function crash( persistence: ReturnType, runId: string, @@ -19,7 +22,7 @@ async function crash( ) { await persistence.stores.runs.createOrResume({ runId, - threadId: 't1', + threadId: T1_KEY, startedAt, }) await persistence.stores.runs.update(runId, { @@ -45,7 +48,7 @@ describe('crash recovery edges', () => { it('continues only the newest crashed turn and fails the older ones', async () => { const persistence = memoryPersistence() - await persistence.stores.messages.saveThread('t1', [ + await persistence.stores.messages.saveThread(T1_KEY, [ { id: 'u1', role: 'user', content: 'hello' }, ]) await crash(persistence, 'old-run', Date.now() - 120_000) diff --git a/packages/ai-harness/tests/resume.test.ts b/packages/ai-harness/tests/resume.test.ts index 3db7f4d9f6..dcf476f025 100644 --- a/packages/ai-harness/tests/resume.test.ts +++ b/packages/ai-harness/tests/resume.test.ts @@ -2,7 +2,12 @@ import { describe, expect, it, vi } from 'vitest' import { z } from 'zod' import { EventType, chat, defineAgent, toolDefinition } from '@tanstack/ai' import { memoryPersistence } from '@tanstack/ai-persistence' -import { HARNESS_EVENTS, createHarnessHost, defineHarness } from '../src' +import { + HARNESS_EVENTS, + createHarnessHost, + defineHarness, + logIdOf, +} from '../src' import { INTERRUPTED_TOOL_RESULT, LEASE, @@ -16,6 +21,9 @@ import { gate, mockAdapter, text, toolCall } from './helpers' import type { AnyChatMiddleware, AnyTool, ModelMessage } from '@tanstack/ai' import type { SessionEvent } from '../src' +/** The key of the thread `t1` in the stores of a session, in the default tenant. */ +const T1_KEY = logIdOf({ logId: 't1' }) + /** Run one chat() turn with the checkpoint middleware and a fresh run record. */ async function runWithCheckpoint(options: { replies: Parameters[0] @@ -109,7 +117,7 @@ describe('checkpoints and leases', () => { { toolCallId: 'call-1', name: 'lookup', replay: 'safe' }, ]) // The assistant message with the tool call is saved before the tool runs. - const saved = await persistence.stores.messages.loadThread('t1') + const saved = await persistence.stores.messages.loadThread(T1_KEY) expect(saved.at(-1)?.toolCalls?.[0]?.function.name).toBe('lookup') release.open() @@ -150,7 +158,7 @@ describe('checkpoints and leases', () => { describe('crash resume', () => { it('continues a crashed turn, runs safe tools again, and notes the others', async () => { const persistence = memoryPersistence() - await persistence.stores.messages.saveThread('t1', [ + await persistence.stores.messages.saveThread(T1_KEY, [ { id: 'u1', role: 'user', content: 'refund and look up' }, { id: 'a1', @@ -172,7 +180,7 @@ describe('crash resume', () => { ]) await persistence.stores.runs.createOrResume({ runId: 'crashed-run', - threadId: 't1', + threadId: T1_KEY, startedAt: Date.now() - 60_000, }) await persistence.stores.runs.update('crashed-run', { @@ -228,7 +236,7 @@ describe('crash resume', () => { expect(lookup).toHaveBeenCalledTimes(1) expect(charge).not.toHaveBeenCalled() - const transcript = await persistence.stores.messages.loadThread('t1') + const transcript = await persistence.stores.messages.loadThread(T1_KEY) const chargeResult = transcript.find( (message) => message.role === 'tool' && message.toolCallId === 'call-charge', @@ -480,7 +488,7 @@ describe('crash resume', () => { const persistence = memoryPersistence() await persistence.stores.runs.createOrResume({ runId: 'crashed-agent', - threadId: 't1', + threadId: T1_KEY, startedAt: Date.now() - 60_000, kind: 'agent', agent: 'waiter', @@ -509,7 +517,7 @@ describe('crash resume', () => { status: 'failed', error: { message: 'The host stopped during this agent run.' }, }) - const transcript = await persistence.stores.messages.loadThread('t1') + const transcript = await persistence.stores.messages.loadThread(T1_KEY) expect(transcript.at(-1)?.content).toBe( '[waiter failed] The host stopped during this agent run.', ) @@ -522,12 +530,12 @@ describe('crash resume', () => { const persistence = memoryPersistence() await persistence.stores.runs.createOrResume({ runId: 'live-run', - threadId: 't1', + threadId: T1_KEY, startedAt: Date.now(), }) await persistence.stores.runs.createOrResume({ runId: 'live-agent', - threadId: 't1', + threadId: T1_KEY, startedAt: Date.now(), kind: 'agent', agent: 'waiter', diff --git a/packages/ai-harness/tests/routing.test.ts b/packages/ai-harness/tests/routing.test.ts index 37c4aaaea9..92ade40a3b 100644 --- a/packages/ai-harness/tests/routing.test.ts +++ b/packages/ai-harness/tests/routing.test.ts @@ -304,7 +304,7 @@ describe('routing.router', () => { await turn const ctx = seen[1] - expect(ctx?.session).toBe(session) + expect(ctx?.session.threadId).toBe(session.threadId) expect(ctx?.input).toEqual([{ type: 'text', content: 'second' }]) expect(ctx?.operationId).toBe(turn.id) expect(ctx?.inputId).toBe('in-2') diff --git a/packages/ai-harness/tests/run-context.test.ts b/packages/ai-harness/tests/run-context.test.ts new file mode 100644 index 0000000000..df9d7029fa --- /dev/null +++ b/packages/ai-harness/tests/run-context.test.ts @@ -0,0 +1,435 @@ +import { memoryCredentialStore } from '@tanstack/ai/byok/server' +import { describe, expect, it } from 'vitest' +import { z } from 'zod' +import { defineAgent, toolDefinition } from '@tanstack/ai' +import { memoryLogStore, memoryPersistence } from '@tanstack/ai-persistence' +import { createHarnessHost, defineHarness, definePlugin, logIdOf } from '../src' +import { applyInput } from '../src/server' +import { after, gate, mockAdapter, text, toolCall } from './helpers' +import type { AnyTool, AnyTextAdapter } from '@tanstack/ai' +import type { Credential } from '@tanstack/ai/byok' +import type { LogRecord } from '@tanstack/ai-persistence' +import type { + AgentRun, + HarnessInput, + HarnessPlugin, + HarnessRunContext, +} from '../src' + +const THREAD = 't1' +/** The id of the log of `THREAD` in the store. */ +const LOG = logIdOf({ logId: THREAD }) +/** Opens the session. She has a token too, so a fall back to her shows. */ +const olga = { id: 'olga' } +const alice = { id: 'alice' } +const bob = { id: 'bob' } + +const SESSION_CREDENTIALS = + 'ctx.credentials of a session plugin has no user. In a tool or a middleware hook, read ctx.context.run.credentials. In a command, use the credentials of its run context.' +const SESSION_KEYS = + 'ctx.keys of a session plugin has no user. In a tool or a middleware hook, read ctx.context.run.keys. In a command, use the keys of its run context.' + +/** A credential store where each of olga, alice, and bob has a GitHub token. */ +async function tokens() { + const store = memoryCredentialStore() + for (const user of ['olga', 'alice', 'bob']) { + await store.set({ + scope: { threadId: THREAD, userId: user }, + id: 'github', + credential: { type: 'api_key', value: `${user}-token` }, + }) + } + return store +} + +/** The value of an `api_key` credential, else `null`. */ +const valueOf = (credential: Credential | null) => + credential?.type === 'api_key' ? credential.value : null + +/** A tool that keeps the user of its run and that user's GitHub token. */ +function tokenReader() { + const seen: Array<{ user: string | undefined; token: string | null }> = [] + const tool = toolDefinition({ + name: 'read_token', + description: 'Read the GitHub token', + inputSchema: z.object({}), + }).server(async (_input, { context }) => { + const credential = await context.run.credentials.get('github') + seen.push({ user: context.run.principal?.id, token: valueOf(credential) }) + return {} + }) + return { tool, seen } +} + +/** A host with the tokens of every user, and a session that olga opens. */ +async function open(options: { + adapter: AnyTextAdapter + plugins: () => Array +}) { + const host = createHarnessHost({ + persistence: memoryPersistence(), + credentials: await tokens(), + }) + const harness = defineHarness({ + name: 'test/run-context', + adapter: options.adapter, + plugins: options.plugins, + }) + const session = await host.open(harness, { + threadId: THREAD, + principal: olga, + }) + return { host, harness, session } +} + +describe('the run context of each run', () => { + it("gives a background agent's tool the credential of the user who started it, while another user's turn runs", async () => { + const bobRunning = gate() + const agentFinished = gate() + const reader = tokenReader() + const agentTools: Array = [reader.tool] + const agentModel = mockAdapter([ + () => toolCall('read_token', {}, 'a1'), + () => text('agent done'), + ]) + const worker = defineAgent({ + name: 'worker', + description: 'Reads the token later', + run: async (ctx) => { + // The tool runs only while bob's turn runs. + await bobRunning.opened + return ctx.chat({ + adapter: agentModel.adapter, + tools: agentTools, + stream: false, + }) + }, + }) + let started: AgentRun | undefined + const starter = definePlugin({ + name: 'test/starter', + setup: (ctx) => ({ + tools: [ + toolDefinition({ + name: 'start_worker', + description: 'Start the worker', + inputSchema: z.object({}), + }).server((_input, { context }) => { + started = ctx.agents.start(worker, undefined, { + principal: context.run.principal, + }) + return Promise.resolve({}) + }), + ], + }), + }) + const { adapter } = mockAdapter([ + () => toolCall('start_worker', {}, 'c1'), + () => text('started'), + // Bob's turn waits until the agent ended. + (request) => ( + bobRunning.open(), + after(agentFinished.opened, 'bob done')(request) + ), + ]) + const { host, session } = await open({ adapter, plugins: () => [starter] }) + + await session.send('Start the worker.', { principal: alice }) + const bobTurn = session.send('Hi.', { principal: bob }) + expect(await started).toBe('agent done') + agentFinished.open() + + expect(await bobTurn).toEqual({ text: 'bob done', media: [] }) + expect(reader.seen).toEqual([{ user: 'alice', token: 'alice-token' }]) + await host.close() + }) + + it("gives a tool of bob's turn bob's credential", async () => { + const reader = tokenReader() + const tools = definePlugin({ + name: 'test/tools', + setup: () => ({ tools: [reader.tool] }), + }) + const { adapter } = mockAdapter([ + () => toolCall('read_token', {}, 'c1'), + () => text('ok'), + ]) + const { host, session } = await open({ adapter, plugins: () => [tools] }) + + await session.send('Read it.', { principal: bob }) + + expect(reader.seen).toEqual([{ user: 'bob', token: 'bob-token' }]) + await host.close() + }) + + it("throws for a session plugin's credentials and keys, and its session principal is the opener", async () => { + const seen: Array<{ + principal: string | undefined + errors: Array + }> = [] + const errorOf = (error: unknown) => + error instanceof Error ? error.message : String(error) + const probe = definePlugin({ + name: 'test/probe', + setup: (ctx) => ({ + tools: [ + toolDefinition({ + name: 'probe', + description: 'Use the session services', + inputSchema: z.object({}), + }).server(async () => { + const errors = await Promise.all([ + ctx.credentials.get('github').then(() => 'read', errorOf), + ctx.keys.require('acme').then(() => 'read', errorOf), + ]) + seen.push({ principal: ctx.session.principal?.id, errors }) + return {} + }), + ], + }), + }) + const { adapter } = mockAdapter([ + () => toolCall('probe', {}, 'c1'), + () => text('ok'), + ]) + const { host, session } = await open({ adapter, plugins: () => [probe] }) + + await session.send('Probe.', { principal: bob }) + + expect(seen).toEqual([ + { principal: 'olga', errors: [SESSION_CREDENTIALS, SESSION_KEYS] }, + ]) + await host.close() + }) + + it('gives a session plugin the owner of the entry as its session principal, also when another user opens the session first', async () => { + const seen: Array = [] + const probe = definePlugin({ + name: 'test/owner-probe', + setup: (ctx) => ({ + tools: [ + toolDefinition({ + name: 'probe', + description: 'Read the session principal', + inputSchema: z.object({}), + }).server(() => { + seen.push(ctx.session.principal?.id) + return {} + }), + ], + }), + }) + const { adapter } = mockAdapter([ + () => text('ok'), + () => toolCall('probe', {}, 'c1'), + () => text('ok'), + ]) + const { host, harness, session } = await open({ + adapter, + plugins: () => [probe], + }) + await session.send('Hi.') + await session.close() + + // Bob opens the thread first on this host now. Olga still owns it. + const reopened = await host.open(harness, { + threadId: THREAD, + principal: bob, + }) + await reopened.send('Probe.') + + expect(seen).toEqual(['olga']) + await host.close() + }) + + it("binds a turn plugin's credentials and session to the sender of its turn", async () => { + const seen: Array<{ principal: string | undefined; token: string | null }> = + [] + const perTurn = definePlugin({ + name: 'test/per-turn', + lifetime: 'turn', + setup: (ctx) => ({ + tools: [ + toolDefinition({ + name: 'read_token', + description: 'Read the GitHub token', + inputSchema: z.object({}), + }).server(async () => { + const credential = await ctx.credentials.get('github') + seen.push({ + principal: ctx.session.principal?.id, + token: valueOf(credential), + }) + return {} + }), + ], + }), + }) + const { adapter } = mockAdapter([ + () => toolCall('read_token', {}, 'c1'), + () => text('ok'), + () => toolCall('read_token', {}, 'c2'), + () => text('ok'), + ]) + const { host, session } = await open({ adapter, plugins: () => [perTurn] }) + + await session.send('Read it.', { principal: bob }) + await session.send('Read it.', { principal: alice }) + + expect(seen).toEqual([ + { principal: 'bob', token: 'bob-token' }, + { principal: 'alice', token: 'alice-token' }, + ]) + await host.close() + }) + + it('replaces a run that a client sends in the context with the run of the server', async () => { + const seen: Array<{ + user: unknown + token: string | null + screen: unknown + }> = [] + const peek = definePlugin({ + name: 'test/peek', + setup: () => ({ + tools: [ + toolDefinition({ + name: 'peek', + description: 'Read the context', + inputSchema: z.object({}), + }).server( + async (_input, { context }) => { + const credential = await context.run.credentials.get('github') + seen.push({ + user: context.run.principal, + token: valueOf(credential), + screen: context.screen, + }) + return {} + }, + ), + ], + }), + }) + const { adapter } = mockAdapter([ + () => toolCall('peek', {}, 'c1'), + () => text('ok'), + ]) + const { host, harness, session } = await open({ + adapter, + plugins: () => [peek], + }) + + await applyInput({ + harness, + session, + input: { + op: 'prompt', + message: 'Peek.', + inputId: 'in-1', + context: { screen: 'settings', run: { principal: { id: 'mallory' } } }, + }, + principal: bob, + }) + expect(await session.settled('in-1')).toMatchObject({ + outcome: 'completed', + }) + + expect(seen).toEqual([ + { user: { id: 'bob' }, token: 'bob-token', screen: 'settings' }, + ]) + await host.close() + }) +}) + +describe('the sender of each control input', () => { + /** The `op` of a logged input. */ + const opOf = (record: LogRecord) => { + const { input } = record + return typeof input === 'object' && input !== null && 'op' in input + ? input.op + : undefined + } + + it('logs the principal that sent a cancel, an answer, a delivery change, a config change, and an agent run', async () => { + const log = memoryLogStore() + const { runs, metadata } = memoryPersistence().stores + const host = createHarnessHost({ + persistence: { stores: { log, runs, metadata } }, + }) + const noop = defineAgent({ + name: 'noop', + description: 'Does nothing', + run: () => Promise.resolve('ok'), + }) + const harness = defineHarness({ + name: 'test/control-sender', + adapter: mockAdapter([]).adapter, + agents: [noop], + expose: { agents: ['noop'], config: ['missing'] }, + }) + const session = await host.open(harness, { + threadId: THREAD, + principal: olga, + }) + const inputs: Array = [ + { op: 'cancel' }, + { op: 'cancelInput', inputId: 'none' }, + { op: 'setDelivery', inputId: 'none', delivery: 'steer' }, + { op: 'answer', questionId: 'none', value: 'yes' }, + { op: 'config', key: 'missing', value: 1 }, + { op: 'agent', agent: 'noop' }, + ] + + for (const input of inputs) { + await applyInput({ harness, session, input, principal: bob }) + } + await session.whenIdle() + + const logged = (await log.read({ logId: LOG })) + .map((entry) => entry.record) + .filter((record) => record.type === 'harness.input') + .map((record) => [opOf(record), record['principal']]) + expect(logged).toEqual([ + ['cancel', { id: 'bob' }], + ['cancelInput', { id: 'bob' }], + ['setDelivery', { id: 'bob' }], + ['answer', { id: 'bob' }], + ['config', { id: 'bob' }], + ['agent', { id: 'bob' }], + ]) + await host.close() + }) +}) + +describe('the tenant of a run', () => { + it('names the tenant of the session and the id of its log in the store', async () => { + const seen: Array<{ tenantId: string | undefined; logId: string }> = [] + const probe = toolDefinition({ + name: 'probe', + description: 'Read the run context', + inputSchema: z.object({}), + }).server(async (_input, { context }) => { + seen.push({ tenantId: context.run.tenantId, logId: context.run.logId }) + return {} + }) + const host = createHarnessHost({ persistence: memoryPersistence() }) + const harness = defineHarness({ + name: 'test/run-tenant', + adapter: mockAdapter([ + () => toolCall('probe', {}, 'p1'), + () => text('ok'), + ]).adapter, + tools: [probe], + }) + const session = await host.open(harness, { + threadId: THREAD, + principal: { id: 'ada', tenantId: 'org-a' }, + }) + + await session.send('Read the run.') + + expect(seen).toEqual([{ tenantId: 'org-a', logId: 'org-a:t1' }]) + await host.close() + }) +}) diff --git a/packages/ai-harness/tests/run-route-hosts.test.ts b/packages/ai-harness/tests/run-route-hosts.test.ts index b7ee45beb7..49c3f6c6c8 100644 --- a/packages/ai-harness/tests/run-route-hosts.test.ts +++ b/packages/ai-harness/tests/run-route-hosts.test.ts @@ -10,8 +10,10 @@ import type { Reply } from './helpers' * that starts on host A must be joinable and resumable from host B. */ function twoHosts() { - const { runs, metadata } = memoryPersistence().stores - const persistence = { stores: { log: memoryLogStore(), runs, metadata } } + const { runs, metadata, sessions } = memoryPersistence().stores + const persistence = { + stores: { log: memoryLogStore(), runs, metadata, sessions }, + } const release = gate() // RUN_STARTED, then the rest after `release`. By RUN_STARTED, host A holds // the run lease, so host B only reads. @@ -25,14 +27,21 @@ function twoHosts() { const { adapter, calls } = mockAdapter([reply]) const harness = defineHarness({ name: 'test/run-hosts', adapter }) const handlerOf = () => { - const host = createHarnessHost({ persistence, log: { coalesceMs: 0 } }) + const host = createHarnessHost({ + persistence, + log: { coalesceMs: 0 }, + access: { + identify: ({ request }) => ({ + id: request.headers.get('x-user') ?? 'user-1', + }), + // A user may use the threads whose id starts with the user id. + authorize: ({ principal, session }) => + session.threadId.startsWith(principal.id), + }, + }) const handler = createHarnessHandler({ host, harness, - authorize: (request) => ({ - id: request.headers.get('x-user') ?? 'user-1', - }), - canAccess: (principal, threadId) => threadId.startsWith(principal.id), media: { secret: 'shared' }, }) const call = (path: string, init: RequestInit = {}) => diff --git a/packages/ai-harness/tests/run-route.test.ts b/packages/ai-harness/tests/run-route.test.ts index 3ed5af40d6..04eaaec0d5 100644 --- a/packages/ai-harness/tests/run-route.test.ts +++ b/packages/ai-harness/tests/run-route.test.ts @@ -31,18 +31,19 @@ const deploy = toolDefinition({ return { ok: true } }) +const user = { id: 'user-1' } + +/** A handler for `user`. `authorize` decides each action. */ function setup( harness: AnyHarness, - canAccess: () => boolean = () => true, + authorize: () => boolean = () => true, persistence = memoryPersistence(), ) { - const host = createHarnessHost({ persistence }) - const handler = createHarnessHandler({ - host, - harness, - authorize: () => ({ id: 'user-1' }), - canAccess, + const host = createHarnessHost({ + persistence, + access: { identify: () => user, authorize }, }) + const handler = createHarnessHandler({ host, harness }) const call = (path: string, body?: object) => handler( new Request(`http://h.test/api/${path}`, { @@ -222,6 +223,7 @@ describe('POST run', () => { expect(deployed).toEqual([]) const session = await host.open(approvalHarness([]), { threadId: 't-1', + principal: user, }) expect(session.snapshot().pendingInterrupts).toHaveLength(1) await host.close() @@ -229,12 +231,18 @@ describe('POST run', () => { it("refuses another person's request with the same runId, and gives them no stream", async () => { const { adapter, calls } = mockAdapter([() => text('hello')]) - const host = createHarnessHost({ persistence: memoryPersistence() }) + const host = createHarnessHost({ + persistence: memoryPersistence(), + access: { + identify: ({ request }) => ({ + id: request.headers.get('x-user') ?? '', + }), + authorize: () => true, + }, + }) const handler = createHarnessHandler({ host, harness: defineHarness({ name: 'test/run-route', adapter }), - canAccess: () => true, - authorize: (request) => ({ id: request.headers.get('x-user') ?? '' }), }) const post = (user: string) => handler( @@ -311,12 +319,20 @@ describe('GET run', () => { await host.close() }) - it('checks canAccess', async () => { + it('refuses a thread that authorize refuses, and hydrates a new thread with nothing and no entry', async () => { const { adapter } = mockAdapter([]) - const { host, call } = setup( - defineHarness({ name: 'test/run-route', adapter }), - () => false, - ) + const harness = defineHarness({ name: 'test/run-route', adapter }) + const { host, call } = setup(harness, () => false) + expect(await (await call('run?threadId=t-1')).json()).toEqual({ + messages: [], + activeRun: null, + interrupts: null, + }) + expect( + await host.sessions.get({ threadId: 't-1', principal: user }), + ).toBeUndefined() + // The user who creates a thread needs no check. + await host.open(harness, { threadId: 't-1', principal: user }) expect((await call('run?threadId=t-1')).status).toBe(403) await host.close() }) @@ -394,6 +410,7 @@ describe('questions', () => { defineHarness({ name: 'test/run-route', adapter }), { threadId: 't-1', + principal: user, }, ) let questions = 0 @@ -477,8 +494,8 @@ describe('join a running turn', () => { } /** Start a turn as `run-a` and wait until its tool asks. */ - async function startAsking(canAccess?: () => boolean) { - const route = setup(askingHarness(), canAccess) + async function startAsking(authorize?: () => boolean) { + const route = setup(askingHarness(), authorize) const posted = await readUntil( await route.run({ runId: 'run-a', diff --git a/packages/ai-harness/tests/session-plugins.test.ts b/packages/ai-harness/tests/session-plugins.test.ts index b9aebc4801..f4c06b7929 100644 --- a/packages/ai-harness/tests/session-plugins.test.ts +++ b/packages/ai-harness/tests/session-plugins.test.ts @@ -13,12 +13,16 @@ import { defineCommand, defineHarness, definePlugin, + logIdOf, } from '../src' import { compact, modelPicker, todos, usage } from '../src/first-party' import { fileCommands, projectInstructions } from '../src/first-party/node' import { gate, mockAdapter, text, toolCall } from './helpers' import type { PluginState, SessionEvent } from '../src' +/** The key of the thread `t` in the stores, in the default tenant. */ +const THREAD_KEY = logIdOf({ logId: 't' }) + const settings = definePlugin({ name: 'test/settings', setup: () => ({ @@ -54,7 +58,7 @@ describe('config values', () => { it('rejects unknown keys and bad values, and loads saved values in a new session', async () => { const persistence = memoryPersistence() - await persistence.stores.metadata.set('harness:config', 't', { + await persistence.stores.metadata.set('harness:config', THREAD_KEY, { verbose: 'not a boolean', gone: 'old key', }) @@ -600,7 +604,7 @@ describe('plugin state and events', () => { it('loads the saved state of a plugin whose name is an Object key', async () => { const persistence = memoryPersistence() - await persistence.stores.metadata.set('plugin:constructor', 't', { + await persistence.stores.metadata.set('plugin:constructor', THREAD_KEY, { saved: true, }) const host = createHarnessHost({ persistence }) diff --git a/packages/ai-harness/tests/session.test.ts b/packages/ai-harness/tests/session.test.ts index 0e5043b060..f79e462792 100644 --- a/packages/ai-harness/tests/session.test.ts +++ b/packages/ai-harness/tests/session.test.ts @@ -13,6 +13,7 @@ import { defineCommand, defineHarness, definePlugin, + logIdOf, } from '../src' import { after, @@ -29,6 +30,9 @@ import type { AnyChatMiddleware, ContentPart, StreamChunk } from '@tanstack/ai' import type { HarnessTurnOptions, Receipt, SessionEvent } from '../src' import type { Reply } from './helpers' +/** The key of the thread `t1` in the stores, in the default tenant. */ +const T1_KEY = logIdOf({ logId: 't1' }) + function setup() { const persistence = memoryPersistence() const host = createHarnessHost({ persistence }) @@ -64,7 +68,7 @@ describe('chat turns', () => { await expect(turn).resolves.toEqual({ text: 'Hello there', media: [] }) expect(turn.status()).toBe('completed') - const saved = await persistence.stores.messages.loadThread('t1') + const saved = await persistence.stores.messages.loadThread(T1_KEY) expect(saved.map((message) => message.role)).toEqual(['user', 'assistant']) const names = customNames(await collect(turn.events({ from: '0' }))) expect(names).toContain(HARNESS_EVENTS.inputApplied) @@ -460,7 +464,7 @@ describe('inbox', () => { const persistence = memoryPersistence() await persistence.stores.inbox.append({ inputId: 'in-1', - threadId: 't1', + threadId: T1_KEY, input: { op: 'prompt', message: 'left over' }, createdAt: Date.now(), }) @@ -475,8 +479,12 @@ describe('inbox', () => { await vi.waitFor(() => expect(session.snapshot().status).toBe('idle')) expect( - (await persistence.stores.inbox.get({ threadId: 't1', inputId: 'in-1' })) - ?.status, + ( + await persistence.stores.inbox.get({ + threadId: T1_KEY, + inputId: 'in-1', + }) + )?.status, ).toBe('applied') expect(messageTexts(calls[0])).toEqual(['left over']) await host.close() @@ -486,7 +494,7 @@ describe('inbox', () => { const persistence = memoryPersistence() await persistence.stores.inbox.append({ inputId: 'in-wake', - threadId: 't1', + threadId: T1_KEY, input: { op: 'followUp', message: 'Background job done.', wake: true }, createdAt: Date.now(), }) diff --git a/packages/ai-harness/tests/sessions-index.test.ts b/packages/ai-harness/tests/sessions-index.test.ts index bb72bbf3e5..8eefa44544 100644 --- a/packages/ai-harness/tests/sessions-index.test.ts +++ b/packages/ai-harness/tests/sessions-index.test.ts @@ -11,10 +11,16 @@ import { createHarnessHost, defineHarness, definePlugin, + logIdOf, + logMessageStore, } from '../src' import { mockAdapter, text, toolCall } from './helpers' import type { ModelMessage, StreamChunk, TokenUsage } from '@tanstack/ai' import type { + MessageStore, + SessionIndexEntry, + SessionIndexKey, + SessionIndexListOptions, SessionIndexPage, SessionIndexStore, } from '@tanstack/ai-persistence' @@ -42,6 +48,10 @@ function idOf(message: ModelMessage | undefined) { const turnsOf = (messages: ReadonlyArray) => messages.map((message) => [message.role, message.content]) +/** The stored transcript of `threadId` in the default tenant. */ +const transcriptOf = (messages: MessageStore, threadId: string) => + messages.loadThread(logIdOf({ logId: threadId })) + const threadIds = (page: SessionIndexPage) => page.entries.map((entry) => entry.threadId) @@ -55,6 +65,16 @@ function hostWith(replies: Array | Reply = () => text('answer')) { return { persistence, host, harness } } +/** A durable host (a log, runs, metadata, and a session index), and its log. */ +function durableHost() { + const log = memoryLogStore() + const { runs, metadata, sessions } = memoryPersistence().stores + const host = createHarnessHost({ + persistence: { stores: { log, runs, metadata, sessions } }, + }) + return { log, host } +} + /** The context items a user pinned to a thread. */ const chipsSchema = z.object({ chips: z.array(z.string()) }) @@ -72,18 +92,17 @@ const chipsHarness = (name = 'test/chips') => */ function storeWith( persistence: HarnessPersistence, - overrides: (store: Required) => Partial, + overrides: (store: SessionIndexStore) => Partial, ) { const store = persistence.stores.sessions - if (!store?.getVersioned || !store.upsertIf) - throw new Error('The persistence has no compare-and-set sessions store.') - const forwarding: Required = { + if (!store) throw new Error('The persistence has no sessions store.') + const forwarding: SessionIndexStore = { upsert: (entry) => store.upsert(entry), - get: (threadId) => store.get(threadId), + get: (key) => store.get(key), list: (options) => store.list(options), - delete: (threadId) => store.delete(threadId), - getVersioned: store.getVersioned.bind(store), - upsertIf: store.upsertIf.bind(store), + delete: (key) => store.delete(key), + getVersioned: (key) => store.getVersioned(key), + upsertIf: (input) => store.upsertIf(input), } return defineSessionIndexStore({ ...forwarding, ...overrides(forwarding) }) } @@ -135,18 +154,23 @@ describe('host.sessions', () => { principal: { id: 'u1', name: 'Ada', tenantId: 'org-1' }, }) - expect(await host.sessions.get('t1')).toEqual({ + const read = { threadId: 't1', principal: { id: 'u1', tenantId: 'org-1' } } + expect(await host.sessions.get(read)).toEqual({ + tenantId: 'org-1', threadId: 't1', + ownerId: 'u1', + members: [{ userId: 'u1', role: 'owner' }], harness: 'test/sessions', createdAt: expect.any(Number), updatedAt: expect.any(Number), - principal: { id: 'u1', tenantId: 'org-1' }, }) + // The entry is in the tenant of the principal only. + expect(await host.sessions.get({ threadId: 't1' })).toBeUndefined() await session.send('first') await session.send('second') - expect((await host.sessions.get('t1'))?.usage).toEqual({ + expect((await host.sessions.get(read))?.usage).toEqual({ turns: 2, promptTokens: 30, completionTokens: 12, @@ -164,7 +188,10 @@ describe('host.sessions', () => { await host.open(harness, { threadId: 't1', principal: { id: 'bob' } }) - expect((await host.sessions.get('t1'))?.principal).toEqual({ id: 'alice' }) + expect(await host.sessions.get({ threadId: 't1' })).toMatchObject({ + ownerId: 'alice', + members: [{ userId: 'alice', role: 'owner' }], + }) await host.close() }) @@ -179,6 +206,7 @@ describe('host.sessions', () => { await persistence.stores.sessions.upsert({ threadId: 'subagent:run-1', parentThreadId: 'older', + members: [], createdAt: 3000, updatedAt: 3000, }) @@ -200,11 +228,13 @@ describe('host.sessions', () => { }) expect(renamed?.title).toBe('Fix the build') - expect((await host.sessions.get('t1'))?.title).toBe('Fix the build') + expect((await host.sessions.get({ threadId: 't1' }))?.title).toBe( + 'Fix the build', + ) expect( await host.sessions.rename({ threadId: 'unknown', title: 'x' }), ).toBeUndefined() - expect(await host.sessions.get('unknown')).toBeUndefined() + expect(await host.sessions.get({ threadId: 'unknown' })).toBeUndefined() await host.close() }) @@ -235,8 +265,8 @@ describe('host.sessions', () => { title: 'Chips', metadata: { chips: ['a.ts', 'b.ts'] }, }) - expect(await host.sessions.get('t1')).toMatchObject({ - principal: { id: 'u1' }, + expect(await host.sessions.get({ threadId: 't1' })).toMatchObject({ + ownerId: 'u1', metadata: { chips: ['a.ts', 'b.ts'] }, }) expect( @@ -252,7 +282,7 @@ describe('host.sessions', () => { it('writes nothing when the change function returns undefined', async () => { const { host, harness } = hostWith() await host.open(harness, { threadId: 't1' }) - const before = await host.sessions.get('t1') + const before = await host.sessions.get({ threadId: 't1' }) expect( await host.sessions.update({ @@ -261,7 +291,7 @@ describe('host.sessions', () => { update: () => undefined, }), ).toBeUndefined() - expect(await host.sessions.get('t1')).toEqual(before) + expect(await host.sessions.get({ threadId: 't1' })).toEqual(before) await host.close() }) @@ -281,7 +311,7 @@ describe('host.sessions', () => { }) await expect(write).rejects.toBeInstanceOf(SessionMetadataError) - expect((await host.sessions.get('t1'))?.title).toBeUndefined() + expect((await host.sessions.get({ threadId: 't1' }))?.title).toBeUndefined() await host.close() }) @@ -304,8 +334,8 @@ describe('host.sessions', () => { // A store that lets another host write right after the next read. let other: (() => Promise) | undefined const racing = storeWith(persistence, (store) => ({ - getVersioned: async (threadId) => { - const read = await store.getVersioned(threadId) + getVersioned: async (key) => { + const read = await store.getVersioned(key) const write = other other = undefined await write?.() @@ -334,7 +364,7 @@ describe('host.sessions', () => { }) expect(runs).toBe(2) - expect(await racing.get('t1')).toMatchObject({ + expect(await racing.get({ threadId: 't1' })).toMatchObject({ title: 'From the other host', metadata: { chips: ['a.ts'] }, }) @@ -361,26 +391,81 @@ describe('host.sessions', () => { await opener.close() }) - it('deletes the index entry only, and a later turn does not bring it back', async () => { + it('deletes the entry and the transcript, and closes the open session', async () => { const { persistence, host, harness } = hostWith() const session = await host.open(harness, { threadId: 't1' }) await session.send('hi') - await host.sessions.delete('t1') - await session.send('still here') - - expect(await host.sessions.get('t1')).toBeUndefined() - expect(turnsOf(await persistence.stores.messages.loadThread('t1'))).toEqual( - [ - ['user', 'hi'], - ['assistant', 'answer'], - ['user', 'still here'], - ['assistant', 'answer'], - ], + await host.sessions.delete({ threadId: 't1' }) + + expect(await host.sessions.get({ threadId: 't1' })).toBeUndefined() + expect( + turnsOf(await transcriptOf(persistence.stores.messages, 't1')), + ).toEqual([]) + expect(host.isOpen(harness, { threadId: 't1' })).toBe(false) + await host.close() + }) + + it('deletes the entry for good, so the next open writes a new entry', async () => { + vi.useFakeTimers({ toFake: ['Date'] }) + const { host, harness } = hostWith() + vi.setSystemTime(1000) + await host.open(harness, { threadId: 't1', principal: { id: 'u1' } }) + await host.sessions.rename({ threadId: 't1', title: 'Old title' }) + + await host.sessions.delete({ threadId: 't1' }) + expect(await host.sessions.get({ threadId: 't1' })).toBeUndefined() + await host.close() + vi.setSystemTime(2000) + await host.open(harness, { threadId: 't1' }) + + expect(await host.sessions.get({ threadId: 't1' })).toEqual({ + threadId: 't1', + members: [], + harness: 'test/sessions', + createdAt: 2000, + updatedAt: 2000, + }) + await host.close() + }) + + it('refuses to open a child thread: its root thread reads it', async () => { + const { persistence, host, harness } = hostWith() + await persistence.stores.sessions.upsert({ + threadId: 'subagent:run-1', + parentThreadId: 't1', + members: [], + createdAt: 1, + updatedAt: 1, + }) + + await expect( + host.open(harness, { threadId: 'subagent:run-1' }), + ).rejects.toThrow( + '"subagent:run-1" is a child thread of "t1". Open the root thread and read the child through it.', ) + expect(host.isOpen(harness, { threadId: 'subagent:run-1' })).toBe(false) await host.close() }) + it('refuses a sessions store without compare-and-set at the start', () => { + const { messages, sessions } = memoryPersistence().stores + // A store in plain JS, with no getVersioned and no upsertIf. + const plain = { + get: (key: SessionIndexKey) => sessions.get(key), + upsert: (entry: SessionIndexEntry) => sessions.upsert(entry), + list: (options?: SessionIndexListOptions) => sessions.list(options), + delete: (key: SessionIndexKey) => sessions.delete(key), + } + + expect(() => + createHarnessHost({ + // @ts-expect-error The types need both methods. A JS caller can leave them out. + persistence: { stores: { messages, sessions: plain } }, + }), + ).toThrow('stores.sessions needs getVersioned and upsertIf.') + }) + it('forks the transcript before or through a message into a new thread, with the settings', async () => { const { persistence, host, harness } = hostWith() const session = await host.open(harness, { @@ -412,18 +497,26 @@ describe('host.sessions', () => { threadId: expect.any(String), harness: 'test/sessions', title: 'Chat (fork)', - principal: { id: 'u1' }, + ownerId: 'u1', + members: [{ userId: 'u1', role: 'owner' }], createdAt: expect.any(Number), updatedAt: expect.any(Number), }) expect(before.threadId).not.toBe(through.threadId) - expect(await host.sessions.get(before.threadId)).toEqual(before) - const { messages } = persistence.stores - expect(turnsOf(await messages.loadThread(before.threadId))).toEqual([ + expect(await host.sessions.get({ threadId: before.threadId })).toEqual( + before, + ) + expect( + turnsOf(await transcriptOf(persistence.stores.messages, before.threadId)), + ).toEqual([ ['user', 'first'], ['assistant', 'answer'], ]) - expect(turnsOf(await messages.loadThread(through.threadId))).toEqual([ + expect( + turnsOf( + await transcriptOf(persistence.stores.messages, through.threadId), + ), + ).toEqual([ ['user', 'first'], ['assistant', 'answer'], ['user', 'second'], @@ -450,20 +543,16 @@ describe('host.sessions', () => { }, }) - expect(await persistence.stores.messages.loadThread(fork.threadId)).toEqual( - [], - ) + expect( + await transcriptOf(persistence.stores.messages, fork.threadId), + ).toEqual([]) const forked = await host.open(harness, { threadId: fork.threadId }) expect(forked.settings()).toEqual({ instructions: 'Be brief.' }) await host.close() }) it('forks a durable thread into one transcript record in a new log', async () => { - const log = memoryLogStore() - const { runs, metadata, sessions } = memoryPersistence().stores - const host = createHarnessHost({ - persistence: { stores: { log, runs, metadata, sessions } }, - }) + const { log, host } = durableHost() const harness = defineHarness({ name: 'test/sessions-durable', adapter: mockAdapter(() => text('durable answer')).adapter, @@ -480,7 +569,9 @@ describe('host.sessions', () => { }, }) - const records = await log.read({ logId: fork.threadId }) + const records = await log.read({ + logId: logIdOf({ logId: fork.threadId }), + }) expect(records.map(({ record }) => [record.type, record.keep])).toEqual([ ['harness.transcript', 0], ]) @@ -528,7 +619,9 @@ describe('host.sessions', () => { 'Thread "t1" belongs to harness "test/sessions", not "test/other".', ) expect(threadIds(await host.sessions.list())).toEqual(['t1']) - expect((await host.sessions.get('t1'))?.harness).toBe('test/sessions') + expect((await host.sessions.get({ threadId: 't1' }))?.harness).toBe( + 'test/sessions', + ) await host.close() }) @@ -555,7 +648,7 @@ describe('host.sessions', () => { }) expect(fork.harness).toBe('test/other') - expect(turnsOf(await messages.loadThread(fork.threadId))).toEqual([ + expect(turnsOf(await transcriptOf(messages, fork.threadId))).toEqual([ ['user', 'hi'], ['assistant', 'answer'], ]) @@ -563,6 +656,54 @@ describe('host.sessions', () => { }) }) +describe('tenants', () => { + const ada = { id: 'ada', tenantId: 'org-a' } + const bo = { id: 'bo', tenantId: 'org-b' } + + it('keeps one thread id of two tenants apart: two sessions, two entries, and two logs', async () => { + const { log, host } = durableHost() + const harness = defineHarness({ + name: 'test/tenants', + adapter: mockAdapter(() => text('answer')).adapter, + }) + const a = await host.open(harness, { threadId: 't1', principal: ada }) + const b = await host.open(harness, { threadId: 't1', principal: bo }) + + await a.send('from a') + await b.send('from b') + + expect(a).not.toBe(b) + expect(turnsOf(await a.transcript())).toEqual([ + ['user', 'from a'], + ['assistant', 'answer'], + ]) + expect(turnsOf(await b.transcript())).toEqual([ + ['user', 'from b'], + ['assistant', 'answer'], + ]) + expect( + await host.sessions.get({ threadId: 't1', principal: ada }), + ).toMatchObject({ tenantId: 'org-a', ownerId: 'ada' }) + expect( + await host.sessions.get({ threadId: 't1', principal: bo }), + ).toMatchObject({ tenantId: 'org-b', ownerId: 'bo' }) + const listed = await host.sessions.list({ principal: ada }) + expect(listed.entries.map((entry) => entry.tenantId)).toEqual(['org-a']) + const logOf = (tenantId: string) => + logMessageStore({ store: log, tenantId }).loadThread('t1') + expect(turnsOf(await logOf('org-a'))).toEqual([ + ['user', 'from a'], + ['assistant', 'answer'], + ]) + expect(turnsOf(await logOf('org-b'))).toEqual([ + ['user', 'from b'], + ['assistant', 'answer'], + ]) + expect(await log.read({ logId: logIdOf({ logId: 't1' }) })).toEqual([]) + await host.close() + }) +}) + describe('child sessions', () => { it('gives an agent run child the harness of its parent', async () => { const writer = defineAgent({ @@ -621,7 +762,8 @@ describe('child sessions', () => { threadId: expect.stringMatching(/^subagent:/), parentThreadId: 'lead', parentToolCallId: 'call-1', - principal: { id: 'alice' }, + ownerId: 'alice', + members: [], harness: 'test/subagent-child', createdAt: expect.any(Number), updatedAt: expect.any(Number), @@ -648,11 +790,11 @@ describe('ctx.session.entry and updateEntry', () => { title: 'From a plugin', metadata: { model: 'm-1' }, }) - expect(await host.sessions.get('p1')).toMatchObject({ + expect(await host.sessions.get({ threadId: 'p1' })).toMatchObject({ threadId: 'p1', title: 'From a plugin', metadata: { model: 'm-1' }, - principal: { id: 'u1' }, + ownerId: 'u1', }) await host.close() }) @@ -690,7 +832,7 @@ describe('ctx.session.entry and updateEntry', () => { metadata: { chips: [...(entry.metadata?.chips ?? []), 'c.ts'] }, })), ).rejects.toBeInstanceOf(SessionMetadataError) - expect((await host.sessions.get('p1'))?.metadata).toEqual({ + expect((await host.sessions.get({ threadId: 'p1' }))?.metadata).toEqual({ chips: ['a.ts', 'b.ts'], }) await host.close() @@ -727,7 +869,7 @@ describe('ctx.session.entry and updateEntry', () => { metadata: { ...entry.metadata, chips: ['a.ts'] }, })) - expect((await host.sessions.get('p1'))?.metadata).toEqual({ + expect((await host.sessions.get({ threadId: 'p1' }))?.metadata).toEqual({ chips: ['a.ts'], theme: 'dark', }) @@ -752,7 +894,7 @@ describe('ctx.session.entry and updateEntry', () => { await session.send('hi') expect(await host.sessions.list()).toEqual({ entries: [] }) - expect(await host.sessions.get('p1')).toBeUndefined() + expect(await host.sessions.get({ threadId: 'p1' })).toBeUndefined() expect( await host.sessions.rename({ threadId: 'p1', title: 'x' }), ).toBeUndefined() diff --git a/packages/ai-harness/tests/shared-log.test.ts b/packages/ai-harness/tests/shared-log.test.ts index d723f816c4..4d25b3c0f5 100644 --- a/packages/ai-harness/tests/shared-log.test.ts +++ b/packages/ai-harness/tests/shared-log.test.ts @@ -4,7 +4,12 @@ import { memoryLogStore, memoryPersistence, } from '@tanstack/ai-persistence' -import { createHarnessHost, defineHarness, logMessageStore } from '../src' +import { + createHarnessHost, + defineHarness, + logIdOf, + logMessageStore, +} from '../src' import { gate, messageTexts, @@ -18,6 +23,8 @@ import type { ReduceOptions } from '../src' import type { Reply } from './helpers' const LOG = 'agent-1' +/** The id of the log `LOG` in the store. */ +const STORE_LOG = logIdOf({ logId: LOG }) function durablePersistence(log: LogStore = memoryLogStore()) { const { runs, metadata } = memoryPersistence().stores @@ -113,7 +120,11 @@ describe('sessions that share one log', () => { 'child answer', 'second child question', ]) - expect(await persistence.stores.log.read({ logId: 'task-1' })).toEqual([]) + expect( + await persistence.stores.log.read({ + logId: logIdOf({ logId: 'task-1' }), + }), + ).toEqual([]) await host.close().catch(() => {}) await next.close() }) @@ -134,7 +145,7 @@ describe('sessions that share one log', () => { await root.append(pair) const types = async () => - (await inner.read({ logId: LOG })).map((entry) => entry.record.type) + (await inner.read({ logId: STORE_LOG })).map((entry) => entry.record.type) expect(await types()).toEqual(['app.child_created', 'app.child_linked']) // The writer tries twice, so both tries fail. @@ -233,9 +244,9 @@ describe('sessions that share one log', () => { await root.append([{ type: 'app.signal' }]) // Another writer takes the next position of the log. - const written = await persistence.stores.log.read({ logId: LOG }) + const written = await persistence.stores.log.read({ logId: STORE_LOG }) await persistence.stores.log.append({ - logId: LOG, + logId: STORE_LOG, seq: written.length + 1, epoch: UNOWNED_EPOCH, records: [{ type: 'app.other' }], diff --git a/packages/ai-harness/tests/status-feed.test.ts b/packages/ai-harness/tests/status-feed.test.ts index 7a45735ddf..74ce9cf64c 100644 --- a/packages/ai-harness/tests/status-feed.test.ts +++ b/packages/ai-harness/tests/status-feed.test.ts @@ -123,7 +123,7 @@ describe('host.events', () => { const feed = watch((signal) => host.events({ signal })) await host.sessions.rename({ threadId: 'b', title: 'Greetings' }) - await host.sessions.delete('b') + await host.sessions.delete({ threadId: 'b' }) await vi.waitFor(() => expect(feed.seen).toHaveLength(3)) expect(feed.seen).toMatchObject([ diff --git a/packages/ai-harness/tests/subagent-children.test.ts b/packages/ai-harness/tests/subagent-children.test.ts index be46407a85..28fb85c249 100644 --- a/packages/ai-harness/tests/subagent-children.test.ts +++ b/packages/ai-harness/tests/subagent-children.test.ts @@ -83,11 +83,12 @@ describe('subagent children of a harness session', () => { await session.send('Notes, please') const threadId = `subagent:${childIdIn(lead.calls[1], 'call-1')}` - expect(await host.sessions.get(threadId)).toEqual({ + expect(await host.sessions.get({ threadId })).toEqual({ threadId, parentThreadId: 't', parentToolCallId: 'call-1', harness: 'test/subagent-children', + members: [], createdAt: expect.any(Number), updatedAt: expect.any(Number), }) @@ -143,11 +144,12 @@ describe('subagent children of a harness session', () => { expect(messageTexts(child.calls[0]).at(-1)).toBe('Look around') expect(messageTexts(lead.calls[2]).at(-1)).toBe(WAKE) const threadId = `subagent:${childIdIn(lead.calls[1], 'call-1')}` - expect(await host.sessions.get(threadId)).toEqual({ + expect(await host.sessions.get({ threadId })).toEqual({ threadId, parentThreadId: 't', parentToolCallId: 'call-1', harness: 'test/subagent-children', + members: [], createdAt: expect.any(Number), updatedAt: expect.any(Number), }) diff --git a/packages/ai-harness/tests/tenant-keys.test.ts b/packages/ai-harness/tests/tenant-keys.test.ts new file mode 100644 index 0000000000..a90fe0c6f9 --- /dev/null +++ b/packages/ai-harness/tests/tenant-keys.test.ts @@ -0,0 +1,284 @@ +import { describe, expect, it, vi } from 'vitest' +import { z } from 'zod' +import { toolDefinition } from '@tanstack/ai' +import { + UNOWNED_EPOCH, + memoryLogStore, + memoryPersistence, +} from '@tanstack/ai-persistence' +import { + configOption, + createHarnessHost, + defineCommand, + defineHarness, + definePlugin, + logIdOf, +} from '../src' +import { tenantStores } from '../src/tenant-stores' +import { mockAdapter, text, toolCall } from './helpers' +import type { ModelMessage } from '@tanstack/ai' +import type { Reply } from './helpers' + +const THREAD = 't1' +const NAME = 'test/tenant-keys' +const acme = { id: 'a', tenantId: 'acme' } +const globex = { id: 'b', tenantId: 'globex' } +/** The store key of `THREAD` in the tenant of `acme`. */ +const ACME_THREAD = logIdOf({ tenantId: acme.tenantId, logId: THREAD }) +const PROMPT = { op: 'prompt', message: 'go', busy: 'queue' } + +/** A plugin with a setting, and a counter in its state. */ +const counter = definePlugin({ + name: 'test/counter', + setup: (ctx) => { + const state = ctx.state({ count: 0 }) + return { + config: { tone: configOption.text({ default: 'plain' }) }, + commands: { + bump: defineCommand({ + description: 'Add one', + run: async () => + (await state.update((value) => ({ count: value.count + 1 }))).count, + }), + }, + } + }, +}) + +const deploy = toolDefinition({ + name: 'deploy', + description: 'Deploy. Needs approval.', + needsApproval: true, + inputSchema: z.object({ env: z.string() }), +}).server(async () => ({ ok: true })) + +function harnessWith(replies: Array = []) { + const model = mockAdapter(replies) + const harness = defineHarness({ + name: NAME, + adapter: model.adapter, + tools: [deploy], + plugins: () => [counter], + }) + return { harness, calls: model.calls } +} + +/** The role and content of each message. */ +const contents = (messages: ReadonlyArray) => + messages.map(({ role, content }) => ({ role, content })) + +/** A host with `stores.log`. */ +function durableHost() { + const { runs, interrupts, metadata, workClaims } = memoryPersistence().stores + const log = memoryLogStore() + const persistence = { + stores: { log, runs, interrupts, metadata, workClaims }, + } + return { + persistence, + /** A prompt that a host logged for `key` just before it stopped. */ + seedInput: (key: string) => + log.append({ + logId: key, + seq: 1, + epoch: UNOWNED_EPOCH, + records: [ + { type: 'harness.input', inputId: 'in-1', input: PROMPT, at: 1 }, + ], + }), + } +} + +/** A host that keeps the transcript in `stores.messages`. */ +function messageStoreHost() { + const persistence = memoryPersistence() + return { + persistence, + /** A prompt that a host left in the inbox of `key` before it stopped. */ + seedInput: (key: string) => + persistence.stores.inbox.append({ + inputId: 'in-1', + threadId: key, + input: PROMPT, + createdAt: Date.now() - 1_000, + }), + } +} + +describe.each([ + { shape: 'a durable host', setup: durableHost }, + { shape: 'a message-store host', setup: messageStoreHost }, +])('one thread id in two tenants on $shape', ({ setup }) => { + it('keeps the settings and the config of one tenant from the other', async () => { + const { persistence } = setup() + const { harness } = harnessWith() + const host = createHarnessHost({ persistence }) + const a = await host.open(harness, { threadId: THREAD, principal: acme }) + await a.configure({ settings: { instructions: 'Answer as acme.' } }) + await a.setConfig({ key: 'tone', value: 'warm' }) + + const b = await host.open(harness, { threadId: THREAD, principal: globex }) + expect(b.settings()).toEqual({}) + expect(b.config().tone?.value).toBe('plain') + await host.close() + + // A new host reads the values of acme back. + const again = createHarnessHost({ persistence }) + const reopened = await again.open(harness, { + threadId: THREAD, + principal: acme, + }) + expect(reopened.settings()).toEqual({ instructions: 'Answer as acme.' }) + expect(reopened.config().tone?.value).toBe('warm') + await again.close() + }) + + it('keeps the plugin state of one tenant from the other', async () => { + const { persistence } = setup() + const { harness } = harnessWith() + const host = createHarnessHost({ persistence }) + const a = await host.open(harness, { threadId: THREAD, principal: acme }) + expect(await a.command({ name: 'bump' })).toBe(1) + + const b = await host.open(harness, { threadId: THREAD, principal: globex }) + expect(b.snapshot().plugins).toEqual({ 'test/counter': { count: 0 } }) + expect(await b.command({ name: 'bump' })).toBe(1) + expect(await a.command({ name: 'bump' })).toBe(2) + await host.close() + }) + + it('keeps the transcript of one tenant from the other', async () => { + const { persistence } = setup() + const { harness } = harnessWith([() => text('Hello, acme.')]) + const host = createHarnessHost({ persistence }) + const a = await host.open(harness, { threadId: THREAD, principal: acme }) + expect((await a.send('Hi.')).text).toBe('Hello, acme.') + + const b = await host.open(harness, { threadId: THREAD, principal: globex }) + expect(await b.transcript()).toEqual([]) + expect(contents(await a.transcript())).toEqual([ + { role: 'user', content: 'Hi.' }, + { role: 'assistant', content: 'Hello, acme.' }, + ]) + await host.close() + }) + + it('keeps a pending interrupt of one tenant from the other', async () => { + const { persistence } = setup() + const { harness } = harnessWith([() => toolCall('deploy', { env: 'prod' })]) + const host = createHarnessHost({ persistence }) + const a = await host.open(harness, { threadId: THREAD, principal: acme }) + expect((await a.send('Deploy.')).interrupts).toHaveLength(1) + + const b = await host.open(harness, { threadId: THREAD, principal: globex }) + expect(b.snapshot().pendingInterrupts).toEqual([]) + await host.close() + + // A new host finds the interrupt of acme again. + const again = createHarnessHost({ persistence }) + const reopened = await again.open(harness, { + threadId: THREAD, + principal: acme, + }) + expect(reopened.snapshot().pendingInterrupts).toHaveLength(1) + await again.close() + }) + + it('refuses a client run id that the other tenant used', async () => { + const { persistence } = setup() + const { harness } = harnessWith([() => text('Hello, acme.')]) + const host = createHarnessHost({ persistence }) + const a = await host.open(harness, { threadId: THREAD, principal: acme }) + expect((await a.send('Hi.', { runId: 'run-1' })).text).toBe('Hello, acme.') + + const b = await host.open(harness, { threadId: THREAD, principal: globex }) + expect(await b.send('Hi.', { runId: 'run-1' }).receipt).toMatchObject({ + status: 'rejected', + reason: 'conflict', + }) + await host.close() + }) + + it('resumes a claimed thread in the tenant of its claim', async () => { + const { persistence, seedInput } = setup() + await seedInput(ACME_THREAD) + await persistence.stores.workClaims.claim({ + threadId: ACME_THREAD, + harness: NAME, + ownerId: 'host-gone', + until: Date.now() - 1_000, + }) + const { harness, calls } = harnessWith([() => text('Ran.')]) + const host = createHarnessHost({ persistence }) + + expect(await host.resumePending({ harnesses: [harness] })).toEqual([ + { threadId: THREAD, harness: NAME }, + ]) + await vi.waitFor(() => expect(calls).toHaveLength(1)) + await vi.waitFor(async () => { + const a = await host.open(harness, { threadId: THREAD, principal: acme }) + expect(contents(await a.transcript())).toEqual([ + { role: 'user', content: 'go' }, + { role: 'assistant', content: 'Ran.' }, + ]) + }) + const other = await host.open(harness, { threadId: THREAD }) + expect(await other.transcript()).toEqual([]) + await host.close() + }) +}) + +describe('tenantStores', () => { + /** The run and interrupt stores of one memory persistence, seen by two tenants. */ + function twoTenants() { + const { runs, interrupts } = memoryPersistence().stores + return { + runs, + acme: tenantStores({ runs, interrupts }, acme.tenantId), + globex: tenantStores({ runs, interrupts }, globex.tenantId), + } + } + + it('keeps each record under its store key, and hides it from the other tenant', async () => { + const { runs, acme: mine, globex: theirs } = twoTenants() + await mine.runs?.createOrResume({ + runId: 'run-1', + threadId: THREAD, + startedAt: 1, + }) + await mine.interrupts?.create({ + interruptId: 'i-1', + runId: 'run-1', + threadId: THREAD, + requestedAt: 1, + payload: {}, + }) + + expect(await runs.get('run-1')).toMatchObject({ threadId: ACME_THREAD }) + expect(await mine.runs?.get('run-1')).toMatchObject({ threadId: THREAD }) + expect(await mine.interrupts?.listPending(THREAD)).toMatchObject([ + { interruptId: 'i-1', threadId: THREAD }, + ]) + expect(await theirs.runs?.get('run-1')).toBeNull() + expect(await theirs.runs?.listByThread?.(THREAD)).toEqual([]) + expect(await theirs.interrupts?.get('i-1')).toBeNull() + expect(await theirs.interrupts?.listByRun('run-1')).toEqual([]) + }) + + it('refuses to go on with a run id of the other tenant', async () => { + const { acme: mine, globex: theirs } = twoTenants() + await mine.runs?.createOrResume({ + runId: 'run-1', + threadId: THREAD, + startedAt: 1, + }) + + await expect( + theirs.runs?.createOrResume({ + runId: 'run-1', + threadId: THREAD, + startedAt: 2, + }), + ).rejects.toThrow('Run run-1 is in another tenant.') + }) +}) diff --git a/packages/ai-harness/tests/title-usage.test.ts b/packages/ai-harness/tests/title-usage.test.ts index 4b1ccf0d26..5db22711b2 100644 --- a/packages/ai-harness/tests/title-usage.test.ts +++ b/packages/ai-harness/tests/title-usage.test.ts @@ -49,7 +49,9 @@ describe('title()', () => { expect(turn.text).toBe('answer') await vi.waitFor(async () => - expect((await host.sessions.get('t'))?.title).toBe('Trip to Rome'), + expect((await host.sessions.get({ threadId: 't' }))?.title).toBe( + 'Trip to Rome', + ), ) expect(messageTexts(titler.calls[0])[0]).toContain( 'Plan a trip to Rome for me', @@ -65,13 +67,17 @@ describe('title()', () => { const { host, session } = await open([title({ adapter: titler.adapter })]) await session.send('Plan a trip to Rome for me') await vi.waitFor(async () => - expect((await host.sessions.get('t'))?.title).toBe('Trip to Rome'), + expect((await host.sessions.get({ threadId: 't' }))?.title).toBe( + 'Trip to Rome', + ), ) await session.send('Add a day in Florence') expect(titler.calls).toHaveLength(1) - expect((await host.sessions.get('t'))?.title).toBe('Trip to Rome') + expect((await host.sessions.get({ threadId: 't' }))?.title).toBe( + 'Trip to Rome', + ) await host.close() }) @@ -97,7 +103,7 @@ describe('title()', () => { await vi.waitFor(() => expect(failures).toEqual([{ message: 'title model down' }]), ) - expect((await host.sessions.get('t'))?.title).toBeUndefined() + expect((await host.sessions.get({ threadId: 't' }))?.title).toBeUndefined() await host.close() }) @@ -181,7 +187,9 @@ describe('usage()', () => { cacheWriteTokens: 500_000, cost: 0.5, }) - expect((await host.sessions.get('t'))?.usage?.cost).toBeCloseTo(3.75, 10) + expect( + (await host.sessions.get({ threadId: 't' }))?.usage?.cost, + ).toBeCloseTo(3.75, 10) await host.close() }) @@ -195,7 +203,7 @@ describe('usage()', () => { expect(await session.command({ name: 'usage' })).toBe( `${CALL_LINE} Cost: $0.5000.`, ) - expect((await host.sessions.get('t'))?.usage?.cost).toBe(0.5) + expect((await host.sessions.get({ threadId: 't' }))?.usage?.cost).toBe(0.5) await host.close() }) @@ -208,7 +216,9 @@ describe('usage()', () => { await session.send('two') // Two calls at $0.50. The host writes no cost of its own. - expect((await host.sessions.get('t'))?.usage?.cost).toBeCloseTo(1, 10) + expect( + (await host.sessions.get({ threadId: 't' }))?.usage?.cost, + ).toBeCloseTo(1, 10) await host.close() }) @@ -223,7 +233,7 @@ describe('usage()', () => { await session.send('two') // Two lead calls and one worker call. - expect((await host.sessions.get('t'))?.usage).toEqual({ + expect((await host.sessions.get({ threadId: 't' }))?.usage).toEqual({ turns: 3, promptTokens: 3_000_000, completionTokens: 300_000, @@ -279,7 +289,7 @@ describe('usage()', () => { await session.send('one') await session.send('two') - expect((await host.sessions.get('t'))?.usage).toEqual({ + expect((await host.sessions.get({ threadId: 't' }))?.usage).toEqual({ turns: 2, promptTokens: 2_000_000, completionTokens: 200_000, @@ -301,7 +311,7 @@ describe('usage()', () => { expect(await session.command({ name: 'usage' })).toBe( `${CALL_LINE} Cost: $0.0000 (cost unknown for 1 call).`, ) - expect((await host.sessions.get('t'))?.usage).toEqual({ + expect((await host.sessions.get({ threadId: 't' }))?.usage).toEqual({ turns: 1, promptTokens: 1_000_000, completionTokens: 100_000, diff --git a/packages/ai-harness/tests/transport-edges.test.ts b/packages/ai-harness/tests/transport-edges.test.ts index 56bec4d2b7..33475fd337 100644 --- a/packages/ai-harness/tests/transport-edges.test.ts +++ b/packages/ai-harness/tests/transport-edges.test.ts @@ -12,9 +12,14 @@ import { parseHarnessInput, } from '../src/server' import { createHarnessClient } from '../src/client' -import { gate, mockAdapter, text, toolCall, untilAborted } from './helpers' -import type { WebSocketLike } from '@tanstack/ai' -import type { HostFrame } from '../src/server' +import { + fakeSocket, + gate, + mockAdapter, + text, + toolCall, + untilAborted, +} from './helpers' const pricer = defineAgent({ name: 'pricer', @@ -36,6 +41,9 @@ const remove = toolDefinition({ }).server(async () => ({ removed: true })) const auth = { authorization: 'Bearer good' } +const user = { id: 'user-1' } +/** The upgrade request of a socket of `user`. */ +const upgrade = new Request('http://x/api/harness/socket', { headers: auth }) function setup(replies: Parameters[0] = []) { const { adapter, calls } = mockAdapter(replies) @@ -47,16 +55,17 @@ function setup(replies: Parameters[0] = []) { subagents: { agents: [pricer] }, expose: { agents: ['pricer'] }, }) - const host = createHarnessHost({ persistence: memoryPersistence() }) - const handler = createHarnessHandler({ - host, - harness, - authorize: (request) => - request.headers.get('authorization') === 'Bearer good' - ? { id: 'user-1' } - : null, - canAccess: (principal, threadId) => threadId.startsWith(principal.id), + const host = createHarnessHost({ + persistence: memoryPersistence(), + access: { + identify: ({ request }) => + request.headers.get('authorization') === 'Bearer good' ? user : null, + // A user may use the threads whose id starts with the user id. + authorize: ({ principal, session }) => + session.threadId.startsWith(principal.id), + }, }) + const handler = createHarnessHandler({ host, harness }) const call = (path: string, init: RequestInit = {}) => handler( new Request(`http://x/api/harness/${path}`, { @@ -85,7 +94,11 @@ const runBody = (threadId: string, extra: Record = {}) => describe('HTTP handler edges', () => { it('answers bad requests with 400, 403, and 404', async () => { - const { host, call } = setup() + const { host, harness, call } = setup() + await host.open(harness, { + threadId: 'someone-else', + principal: { id: 'someone' }, + }) const statusOf = async (path: string, init?: RequestInit) => (await call(path, init)).status @@ -227,44 +240,15 @@ describe('HTTP handler edges', () => { }) describe('WebSocket edges', () => { - function fakeSocket(options: { failSend?: boolean } = {}) { - const handlers: Record void>> = {} - const sent: Array = [] - const closed: Array = [] - const socket: WebSocketLike = { - send: (data) => { - if (options.failSend) throw new Error('socket gone') - sent.push(JSON.parse(data)) - }, - close: (code?: number) => { - closed.push(code) - handlers.close?.forEach((handler) => handler()) - }, - addEventListener: (type: string, handler: (event?: any) => void) => { - ;(handlers[type] ??= []).push(handler) - }, - } - const receive = (frame: unknown) => - handlers.message?.forEach((handler) => - handler({ - data: typeof frame === 'string' ? frame : JSON.stringify(frame), - }), - ) - const fail = () => handlers.error?.forEach((handler) => handler()) - return { socket, sent, closed, receive, fail } - } - it('refuses a thread the principal cannot access', async () => { const { host, harness } = setup() - const { socket, sent, closed, receive } = fakeSocket() - handleHarnessSocket({ - host, - harness, - socket, - principal: { id: 'user-1' }, - canAccess: () => false, + await host.open(harness, { + threadId: 'user-2-x', + principal: { id: 'user-2' }, }) - receive({ type: 'harness.subscribe', threadId: 'user-1-x' }) + const { socket, sent, closed, receive } = fakeSocket() + handleHarnessSocket({ host, harness, socket, request: upgrade }) + receive({ type: 'harness.subscribe', threadId: 'user-2-x' }) await vi.waitFor(() => expect(closed).toEqual([4403])) expect(sent).toEqual([{ type: 'harness.error', message: 'forbidden' }]) await host.close() @@ -277,8 +261,7 @@ describe('WebSocket edges', () => { host, harness, socket: first.socket, - principal: { id: 'user-1' }, - canAccess: () => true, + request: upgrade, }) first.receive({ type: 'harness.subscribe', threadId: 'user-1-ws2' }) await vi.waitFor(() => @@ -323,8 +306,7 @@ describe('WebSocket edges', () => { host, harness, socket: second.socket, - principal: { id: 'user-1' }, - canAccess: () => true, + request: upgrade, }) second.receive({ type: 'harness.subscribe', @@ -351,8 +333,7 @@ describe('WebSocket edges', () => { host, harness, socket, - principal: { id: 'user-1' }, - canAccess: () => true, + request: upgrade, }) receive({ type: 'harness.subscribe', threadId: 'user-1-ws3' }) receive({ @@ -377,8 +358,7 @@ describe('WebSocket edges', () => { host, harness, socket, - principal: { id: 'user-1' }, - canAccess: () => true, + request: upgrade, }) receive({ type: 'harness.subscribe', threadId: 'user-1-dead' }) // Nothing to assert on the socket: the send failure must not throw. @@ -433,7 +413,10 @@ describe('protocol parsing and inputs', () => { () => text('queued answer'), () => text('follow-up answer'), ]) - const session = await host.open(harness, { threadId: 'user-1-apply' }) + const session = await host.open(harness, { + threadId: 'user-1-apply', + principal: user, + }) const first = await applyInput({ harness, @@ -640,7 +623,10 @@ describe('createHarnessClient edges', () => { describe('an operation that waits for cancel', () => { it('cancels through a control input', async () => { const { host, harness } = setup([untilAborted()]) - const session = await host.open(harness, { threadId: 'user-1-cancel' }) + const session = await host.open(harness, { + threadId: 'user-1-cancel', + principal: user, + }) const operation = session.send('wait') await vi.waitFor(() => expect(session.snapshot().status).toBe('running')) await applyInput({ diff --git a/packages/ai-harness/tests/turn-control.test.ts b/packages/ai-harness/tests/turn-control.test.ts index 1fe0535903..0ecd76bbd1 100644 --- a/packages/ai-harness/tests/turn-control.test.ts +++ b/packages/ai-harness/tests/turn-control.test.ts @@ -2,7 +2,12 @@ import { describe, expect, it, vi } from 'vitest' import { z } from 'zod' import { CONTEXT_OVERFLOW_CODE, EventType, toolDefinition } from '@tanstack/ai' import { memoryLogStore, memoryPersistence } from '@tanstack/ai-persistence' -import { createHarnessHost, defineHarness, retryTransientErrors } from '../src' +import { + createHarnessHost, + defineHarness, + logIdOf, + retryTransientErrors, +} from '../src' import { TRUNCATED_TOOL_RESULT } from '../src/resume' import { gate, @@ -17,6 +22,8 @@ import type { HarnessTurnOptions, ProjectOptions } from '../src' import type { Reply } from './helpers' const THREAD = 't1' +/** The id of the log of `THREAD` in the store. */ +const LOG = logIdOf({ logId: THREAD }) /** The start of an assistant message that streamed `delta`. */ const partialText = (delta: string): Array => [ @@ -903,7 +910,7 @@ describe('turn.beforeFinish', () => { expect(await session.send('go')).toEqual({ text: 'done', media: [] }) expect(calls).toHaveLength(1) expect( - (await log.read({ logId: THREAD })).map((entry) => entry.record.type), + (await log.read({ logId: LOG })).map((entry) => entry.record.type), ).toContain('app.state_write') await host.close() }) @@ -982,7 +989,7 @@ describe('turn.beforeFinish', () => { REMINDER, ]) expect(JSON.stringify(await session.transcript())).not.toContain(REMINDER) - expect(JSON.stringify(await log.read({ logId: THREAD }))).not.toContain( + expect(JSON.stringify(await log.read({ logId: LOG }))).not.toContain( REMINDER, ) await host.close() diff --git a/packages/ai-harness/tests/turn-overrides.test.ts b/packages/ai-harness/tests/turn-overrides.test.ts index 349482a6f7..eb2541c0a5 100644 --- a/packages/ai-harness/tests/turn-overrides.test.ts +++ b/packages/ai-harness/tests/turn-overrides.test.ts @@ -12,6 +12,7 @@ import { defineHarness, definePlugin, durableTool, + logIdOf, } from '../src' import { after, @@ -33,6 +34,8 @@ import type { HarnessConfig, PluginContributions } from '../src' import type { Reply } from './helpers' const THREAD = 't1' +/** The id of the log of `THREAD` in the store. */ +const LOG = logIdOf({ logId: THREAD }) const HIGH: ReasoningRequest = { level: 'high', summary: true } const LOW: ReasoningRequest = { level: 'low', summary: false } @@ -108,7 +111,7 @@ const toolResult = ( /** The records of the log, of one type. */ async function recordsOf(persistence: Durable, type: string) { - return (await persistence.stores.log.read({ logId: THREAD })) + return (await persistence.stores.log.read({ logId: LOG })) .map((entry) => entry.record) .filter((record) => record.type === type) } @@ -509,7 +512,13 @@ describe('turn overrides: plugin adapter picker', () => { expect(await turn).toEqual({ text: 'turn', media: [] }) expect(seen).toEqual([ - { operationId: turn.id, inputId: 'in-7', message: 'go', overrides }, + { + operationId: turn.id, + inputId: 'in-7', + message: 'go', + overrides, + run: expect.objectContaining({ runId: turn.id }), + }, ]) expect(picked.calls).toHaveLength(0) expect(calls).toHaveLength(0) @@ -537,7 +546,7 @@ describe('ephemeral input', () => { expect(messageTexts(calls[0])).toEqual(['go', NOTE]) expect(JSON.stringify(await session.transcript())).not.toContain(NOTE) expect( - JSON.stringify(await persistence.stores.log.read({ logId: THREAD })), + JSON.stringify(await persistence.stores.log.read({ logId: LOG })), ).not.toContain(NOTE) await host.close() }) @@ -545,7 +554,7 @@ describe('ephemeral input', () => { it('gives the ephemeral messages of continue to the model, and no store keeps them', async () => { const persistence = durablePersistence() await persistence.stores.log.append({ - logId: THREAD, + logId: LOG, seq: 1, epoch: UNOWNED_EPOCH, records: [ @@ -570,7 +579,7 @@ describe('ephemeral input', () => { expect(messageTexts(calls[0])).toEqual(['go', NOTE]) expect(JSON.stringify(await session.transcript())).not.toContain(NOTE) expect( - JSON.stringify(await persistence.stores.log.read({ logId: THREAD })), + JSON.stringify(await persistence.stores.log.read({ logId: LOG })), ).not.toContain(NOTE) await host.close() }) diff --git a/packages/ai-harness/tests/types.test-d.ts b/packages/ai-harness/tests/types.test-d.ts index 00dacb6c98..0d33227389 100644 --- a/packages/ai-harness/tests/types.test-d.ts +++ b/packages/ai-harness/tests/types.test-d.ts @@ -5,6 +5,7 @@ import { createHarnessHost, defineHarness, definePlugin, + membersOnly, oauthConnector, } from '../src' import { createHarnessHandler } from '../src/server' @@ -14,10 +15,12 @@ import type { SessionIndexEntry } from '@tanstack/ai-persistence' import type { AgentRun, HarnessConfig, + HarnessHost, HarnessInput, HarnessSession, Operation, PluginLifetime, + Principal, ProjectOptions, Receipt, RecoverHook, @@ -182,15 +185,54 @@ it('refuses a plugin whose metadata the harness does not have', () => { }) }) -it('needs an explicit canAccess and an explicit login mode', () => { +it('types the principal of the host from what access.identify returns', () => { + type User = { id: string; tenantId: string; role: 'admin' | 'user' } + // The app's own login check, as `verifySession(request)` in an app. + const userOf = async (request: Request): Promise => + request.headers.has('cookie') + ? { id: 'u1', tenantId: 'acme', role: 'user' } + : null + const host = createHarnessHost({ + access: { + identify: ({ request }) => userOf(request), + authorize: ({ principal }) => { + expectTypeOf(principal).toEqualTypeOf() + return principal.role === 'admin' + }, + }, + }) + expectTypeOf(host).toEqualTypeOf>() + expectTypeOf[1]['principal']>().toEqualTypeOf< + User | undefined + >() + const user: User = { id: 'u2', tenantId: 'acme', role: 'user' } + void host.open(harness, { threadId: 't1', principal: user }) + // @ts-expect-error The principal of this host has a role. + void host.open(harness, { threadId: 't1', principal: { id: 'u3' } }) + // membersOnly() fits any principal type. + expectTypeOf( + createHarnessHost({ + access: { identify: () => user, authorize: membersOnly() }, + }), + ).toEqualTypeOf>() + // Without access, the principal is the plain one. + expectTypeOf(createHarnessHost()).toEqualTypeOf>() +}) + +it('types the openAccess principal by the host, and needs an explicit login mode', () => { const host = createHarnessHost() - // @ts-expect-error `canAccess` is required: allow-all must be explicit. - createHarnessHandler({ host, harness, authorize: () => ({ id: 'u' }) }) + createHarnessHandler({ host, harness, openAccess: true }) + createHarnessHandler({ + host, + harness, + openAccess: { principal: { id: 'u' } }, + }) createHarnessHandler({ host, harness, + openAccess: true, + // @ts-expect-error The host access decides who calls: no authorize here. authorize: () => ({ id: 'u' }), - canAccess: () => true, }) const oauth = { authorizationUrl: 'a', tokenUrl: 't', clientId: 'c' } // @ts-expect-error `login` is required: no implicit loopback listener. diff --git a/packages/ai-harness/tests/usage.test.ts b/packages/ai-harness/tests/usage.test.ts index 9b45fdf9d3..76cd72fb68 100644 --- a/packages/ai-harness/tests/usage.test.ts +++ b/packages/ai-harness/tests/usage.test.ts @@ -2,7 +2,7 @@ import { describe, expect, it } from 'vitest' import { z } from 'zod' import { EventType, defineAgent, toolDefinition } from '@tanstack/ai' import { memoryLogStore, memoryPersistence } from '@tanstack/ai-persistence' -import { createHarnessHost, defineHarness } from '../src' +import { createHarnessHost, defineHarness, logIdOf } from '../src' import { addUsage, emptyUsage } from '../src/usage' import { mockAdapter, text, toolCall } from './helpers' import type { StreamChunk } from '@tanstack/ai' @@ -10,6 +10,8 @@ import type { SessionEvent } from '../src' import type { Reply } from './helpers' const THREAD = 't1' +/** The id of the log of `THREAD` in the store. */ +const LOG = logIdOf({ logId: THREAD }) const alice = { id: 'alice' } const bob = { id: 'bob' } @@ -233,7 +235,7 @@ describe('usage totals after a restart', () => { const before = session.usage() await host.close() - const records = (await log.read({ logId: THREAD })).filter( + const records = (await log.read({ logId: LOG })).filter( (entry) => entry.record.type === 'harness.usage', ) expect(records).toHaveLength(2) diff --git a/packages/ai-harness/tests/view-client-tools.test.ts b/packages/ai-harness/tests/view-client-tools.test.ts index 68a7925549..ffa0e56524 100644 --- a/packages/ai-harness/tests/view-client-tools.test.ts +++ b/packages/ai-harness/tests/view-client-tools.test.ts @@ -50,8 +50,7 @@ describe('session view client tools', () => { const handler = createHarnessHandler({ host, harness, - canAccess: () => true, - authorize: () => ({ id: 'u' }), + openAccess: { principal: { id: 'u' } }, }) const controls: Array = [] const client = createHarnessClient({ diff --git a/packages/ai-harness/tests/view.test.ts b/packages/ai-harness/tests/view.test.ts index e2d74b0ae5..983ac399ad 100644 --- a/packages/ai-harness/tests/view.test.ts +++ b/packages/ai-harness/tests/view.test.ts @@ -310,8 +310,7 @@ describe('createSessionView', () => { const handler = createHarnessHandler({ host, harness, - canAccess: () => true, - authorize: () => ({ id: 'u' }), + openAccess: { principal: { id: 'u' } }, }) const client = createHarnessClient({ url: 'http://local/api/harness', @@ -347,8 +346,7 @@ describe('createSessionView', () => { const handler = createHarnessHandler({ host, harness, - canAccess: () => true, - authorize: () => ({ id: 'u' }), + openAccess: { principal: { id: 'u' } }, }) const view = createSessionView( createHarnessClient({ diff --git a/packages/ai-harness/tests/work-claims.test.ts b/packages/ai-harness/tests/work-claims.test.ts index cc14544bfa..7ba6ff105a 100644 --- a/packages/ai-harness/tests/work-claims.test.ts +++ b/packages/ai-harness/tests/work-claims.test.ts @@ -6,7 +6,7 @@ import { memoryLogStore, memoryPersistence, } from '@tanstack/ai-persistence' -import { createHarnessHost, defineHarness, definePlugin } from '../src' +import { createHarnessHost, defineHarness, definePlugin, logIdOf } from '../src' import { after, gate, @@ -20,6 +20,8 @@ import type { WorkClaimStore } from '@tanstack/ai-persistence' import type { Reply } from './helpers' const THREAD = 't1' +/** The id of the log of `THREAD` in the store, and its key in the other stores. */ +const LOG = logIdOf({ logId: THREAD }) const NAME = 'test/work-claims' const later = () => Date.now() + 60 * 60_000 @@ -43,7 +45,7 @@ function toolCallAfter( } /** Can another host claim the thread now? It takes the thread when it can. */ -const otherHostCanClaim = (claims: WorkClaimStore, threadId = THREAD) => +const otherHostCanClaim = (claims: WorkClaimStore, threadId = LOG) => claims.claim({ threadId, harness: NAME, @@ -145,7 +147,7 @@ describe('work claims while a session runs', () => { it('refuses an input while another host claim is live, without an inbox', async () => { const persistence = durablePersistence() await persistence.stores.workClaims.claim({ - threadId: THREAD, + threadId: LOG, harness: NAME, ownerId: 'other-host', until: Date.now() + 60_000, @@ -169,7 +171,7 @@ describe('work claims while a session runs', () => { const { inbox } = memoryPersistence().stores const claims = persistence.stores.workClaims await claims.claim({ - threadId: THREAD, + threadId: LOG, harness: NAME, ownerId: 'other-host', until: Date.now() + 300, @@ -234,11 +236,7 @@ describe('work claims while a session runs', () => { }) /** An expired claim: the host that held the thread stopped. */ -const expireClaim = ( - claims: WorkClaimStore, - threadId = THREAD, - harness = NAME, -) => +const expireClaim = (claims: WorkClaimStore, threadId = LOG, harness = NAME) => claims.claim({ threadId, harness, @@ -251,7 +249,7 @@ async function seedStoppedTurn( persistence: ReturnType, ) { await persistence.stores.log.append({ - logId: THREAD, + logId: LOG, seq: 1, epoch: UNOWNED_EPOCH, records: [ @@ -276,7 +274,7 @@ async function seedStoppedTurn( }) await persistence.stores.runs.createOrResume({ runId: 'op-stopped', - threadId: THREAD, + threadId: LOG, startedAt: Date.now() - 60_000, }) await persistence.stores.runs.update('op-stopped', { @@ -292,7 +290,7 @@ const settledRecord = async ( persistence: ReturnType, inputId: string, ) => - (await persistence.stores.log.read({ logId: THREAD })) + (await persistence.stores.log.read({ logId: LOG })) .map((entry) => entry.record) .find( (record) => @@ -304,7 +302,7 @@ async function seedStoppedAgent( persistence: ReturnType, ) { await persistence.stores.log.append({ - logId: THREAD, + logId: LOG, seq: 1, epoch: UNOWNED_EPOCH, records: [ @@ -324,7 +322,7 @@ async function seedStoppedAgent( }) await persistence.stores.runs.createOrResume({ runId: 'op-agent', - threadId: THREAD, + threadId: LOG, startedAt: Date.now() - 60_000, kind: 'agent', agent: 'waiter', @@ -375,7 +373,7 @@ describe('host.resumePending', () => { it('runs an input that was stored just before the host stopped', async () => { const persistence = durablePersistence() await persistence.stores.log.append({ - logId: THREAD, + logId: LOG, seq: 1, epoch: UNOWNED_EPOCH, records: [ @@ -407,7 +405,7 @@ describe('host.resumePending', () => { const persistence = memoryPersistence() await persistence.stores.inbox.append({ inputId: 'in-1', - threadId: THREAD, + threadId: LOG, input: { op: 'prompt', message: 'go', busy: 'queue' }, createdAt: Date.now() - 1_000, }) @@ -423,7 +421,7 @@ describe('host.resumePending', () => { expect( ( await persistence.stores.inbox.get({ - threadId: THREAD, + threadId: LOG, inputId: 'in-1', }) )?.status, @@ -473,7 +471,7 @@ describe('host.resumePending', () => { it('continues a resumable background agent whose host stopped', async () => { const persistence = durablePersistence() await persistence.stores.log.append({ - logId: THREAD, + logId: LOG, seq: 1, epoch: UNOWNED_EPOCH, records: [ @@ -493,7 +491,7 @@ describe('host.resumePending', () => { }) await persistence.stores.runs.createOrResume({ runId: 'op-agent', - threadId: THREAD, + threadId: LOG, startedAt: Date.now() - 60_000, kind: 'agent', agent: 'worker', @@ -601,11 +599,8 @@ describe('host.resumePending', () => { it('skips a thread of a harness it was not given, and leaves it listed', async () => { const persistence = durablePersistence() - await expireClaim( - persistence.stores.workClaims, - 'other-thread', - 'test/other', - ) + const other = logIdOf({ logId: 'other-thread' }) + await expireClaim(persistence.stores.workClaims, other, 'test/other') const { adapter } = mockAdapter([]) const host = createHarnessHost({ persistence }) @@ -616,7 +611,7 @@ describe('host.resumePending', () => { ).toEqual([]) expect( await persistence.stores.workClaims.listExpired({ now: Date.now() }), - ).toEqual([{ threadId: 'other-thread', harness: 'test/other' }]) + ).toEqual([{ threadId: other, harness: 'test/other' }]) await host.close() }) @@ -626,7 +621,7 @@ describe('host.resumePending', () => { // Older than this thread's claim, so the store lists them first. for (let index = 0; index < 3; index += 1) { await persistence.stores.workClaims.claim({ - threadId: `other-${index}`, + threadId: logIdOf({ logId: `other-${index}` }), harness: 'test/other', ownerId: 'host-gone', until: Date.now() - 60_000, @@ -692,7 +687,7 @@ describe('host.resumePending', () => { const persistence = durablePersistence() // The failing thread has the older claim, so it comes first. await persistence.stores.workClaims.claim({ - threadId: 't-bad', + threadId: logIdOf({ logId: 't-bad' }), harness: NAME, ownerId: 'host-gone', until: Date.now() - 3_000, @@ -739,7 +734,7 @@ describe('host.resumePending', () => { // Three threads whose host stopped, the oldest claim first. for (const [index, threadId] of ['t-a', 't-b', 't-c'].entries()) { await claims.claim({ - threadId, + threadId: logIdOf({ logId: threadId }), harness: NAME, ownerId: 'host-gone', until: now - 3_000 + index, diff --git a/packages/ai-mcp/src/connector.ts b/packages/ai-mcp/src/connector.ts index 4a80a62d2f..728925fc0e 100644 --- a/packages/ai-mcp/src/connector.ts +++ b/packages/ai-mcp/src/connector.ts @@ -15,7 +15,11 @@ import type { StoredOAuthClientInformation, } from '@modelcontextprotocol/client' import type { AnyTool } from '@tanstack/ai' -import type { CredentialsAccess, Principal } from '@tanstack/ai-harness' +import type { + CredentialsAccess, + RunContext, + TurnInfo, +} from '@tanstack/ai-harness' import type { MCPClient } from './client' import type { MCPClientOptions } from './types' @@ -216,10 +220,11 @@ export function mcpConnector(options: McpConnectorOptions) { return definePlugin({ name: `connector/${id}`, setup: async (ctx) => { - // One sign-in per sender: in a shared thread, one person's sign-in, - // client, and tools never serve another person's turn. + // One sign-in per user of a run, by tenant and id: in a shared thread, + // one person's sign-in, client, and tools never serve another + // person's turn. const connections = new Map() - const keyOf = (principal: Principal | undefined) => + const keyOf = (principal: RunContext['principal']) => principal ? JSON.stringify([principal.tenantId ?? null, principal.id]) : '' @@ -228,7 +233,7 @@ export function mcpConnector(options: McpConnectorOptions) { } /** Set the sign-in of `principal`, and drop its old client and tools. */ const signedIn = async ( - principal: Principal | undefined, + principal: RunContext['principal'], connected: boolean, ) => { const key = keyOf(principal) @@ -245,20 +250,22 @@ export function mcpConnector(options: McpConnectorOptions) { prompts: [ { id: `connector/${id}:status`, - // Read after `discoverTools`, which finds the sender's state. - text: () => - connections.get(keyOf(ctx.session.principal))?.connected + // Read after `discoverTools`, which finds the state of the + // turn's user. + text: (turn: TurnInfo) => + connections.get(keyOf(turn.run.principal))?.connected ? '' : `${label} is not connected. If the user asks for ${label}, tell them to run /connect ${id}.`, }, ], - // Runs in a turn, so `ctx.session.principal` and `ctx.credentials` - // are the turn's sender. - discoverTools: async () => { - const key = keyOf(ctx.session.principal) + // Runs for each turn, with the run of that turn: its user and that + // user's credentials. + discoverTools: async (turn: TurnInfo) => { + const { principal, credentials } = turn.run + const key = keyOf(principal) let connection = connections.get(key) if (!connection) { - connection = { connected: (await ctx.credentials.get(id)) !== null } + connection = { connected: (await credentials.get(id)) !== null } connections.set(key, connection) } if (!connection.connected) return [] @@ -266,7 +273,7 @@ export function mcpConnector(options: McpConnectorOptions) { // The SDK only refreshes tokens for a client with a redirect URL. // Without one it treats the client as machine-to-machine. A new // browser sign-in still happens only through /connect. - const saved = await ctx.credentials.get(id) + const saved = await credentials.get(id) const redirectUri = (saved?.type === 'oauth' ? saved.client?.redirectUri @@ -277,7 +284,7 @@ export function mcpConnector(options: McpConnectorOptions) { transport: { type: 'http', url, - authProvider: credentialProvider(id, ctx.credentials, { + authProvider: credentialProvider(id, credentials, { redirectUri, clientName, scopes: options.scopes, @@ -327,11 +334,11 @@ export function mcpConnector(options: McpConnectorOptions) { [`connect:${id}`]: defineCommand({ description: `Sign in to ${label}`, // The sign-in belongs to the user who runs the command. - run: async (_input, { credentials, principal }) => { + run: async (_input, { run }) => { const receiver = await startLoopbackReceiver() const state = randomBytes(16).toString('base64url') try { - const provider = credentialProvider(id, credentials, { + const provider = credentialProvider(id, run.credentials, { redirectUri: receiver.redirectUri, clientName, scopes: options.scopes, @@ -365,15 +372,15 @@ export function mcpConnector(options: McpConnectorOptions) { } finally { receiver.close() } - await signedIn(principal, true) + await signedIn(run.principal, true) return `Connected to ${label}.` }, }), [`disconnect:${id}`]: defineCommand({ description: `Sign out of ${label}`, - run: async (_input, { credentials, principal }) => { - await credentials.delete(id) - await signedIn(principal, false) + run: async (_input, { run }) => { + await run.credentials.delete(id) + await signedIn(run.principal, false) return `Disconnected from ${label}.` }, }), diff --git a/packages/ai-mcp/src/harness-plugin.ts b/packages/ai-mcp/src/harness-plugin.ts index 545300c145..a22a739f2a 100644 --- a/packages/ai-mcp/src/harness-plugin.ts +++ b/packages/ai-mcp/src/harness-plugin.ts @@ -6,6 +6,7 @@ import type { AnyCommand, PluginPrompt, PluginSessionApi, + TurnInfo, } from '@tanstack/ai-harness' import type { OnElicit } from './client' import type { AskInput } from './tools' @@ -162,7 +163,7 @@ export function mcp(options: { const isOAuth = server.type === 'http' && server.oauth === true if (!isOAuth) { const connecting = connect(name, server) - return () => connecting + return (_turn: TurnInfo) => connecting } // Loaded here: the connector needs node:crypto and node:http. const { mcpConnector } = await import('./connector') @@ -177,11 +178,12 @@ export function mcp(options: { prompts.push(...connector.prompts) // The sign-in belongs to the sender of the turn, so the status is // the one of the last turn. - return async () => { + return async (turn: TurnInfo) => { try { - const tools = await connector.discoverTools() + const tools = await connector.discoverTools(turn) const signedIn = - tools.length > 0 || (await ctx.credentials.get(name)) !== null + tools.length > 0 || + (await turn.run.credentials.get(name)) !== null await setStatus( signedIn ? { name, status: 'connected', toolCount: tools.length } @@ -210,8 +212,8 @@ export function mcp(options: { return { commands, prompts, - discoverTools: async () => { - const found = await Promise.all(lists.map((list) => list())) + discoverTools: async (turn) => { + const found = await Promise.all(lists.map((list) => list(turn))) return found.flat() }, } diff --git a/packages/ai-mcp/src/harness.ts b/packages/ai-mcp/src/harness.ts index cc00099497..9331ee69a6 100644 --- a/packages/ai-mcp/src/harness.ts +++ b/packages/ai-mcp/src/harness.ts @@ -4,7 +4,7 @@ import { readUnopenedInterruptBinding, toolDefinition, } from '@tanstack/ai' -import { HARNESS_EVENTS } from '@tanstack/ai-harness' +import { AccessDeniedError, HARNESS_EVENTS } from '@tanstack/ai-harness' import { isMediaRecord, kindOf, @@ -12,11 +12,22 @@ import { mimeTypeOf, } from '@tanstack/ai-harness/media' import { applyInput, parseHarnessInput } from '@tanstack/ai-harness/server' +import { + OAuthError, + OAuthErrorCode, + bearerAuthChallengeResponse, +} from '@modelcontextprotocol/server' // The inferred return type names `MCPHandleOptions`. Without an import of // `./server/index`, the .d.ts emit writes `./server.js`, which does not resolve. import { createMCPServer, resourceDefinition } from './server/index' +import { isServedOverStdio } from './server/registry' import { toCallToolResult } from './server/tasks' -import type { ContentPart, Interrupt, JSONSchema } from '@tanstack/ai' +import type { + ContentPart, + Interrupt, + JSONSchema, + ToolExecutionContext, +} from '@tanstack/ai' import type { AnyHarness, Cursor, @@ -24,19 +35,34 @@ import type { HarnessSession, MediaRecord, Operation, + Principal, } from '@tanstack/ai-harness' +import type { OpenAccess } from '@tanstack/ai-harness/server' import type { ContentBlock } from '@modelcontextprotocol/server' import type { MCPToolContext } from './server/context' +import type { MCPHandleOptions } from './server/index' export { mcp } from './harness-plugin' export type { McpServerConfig, McpServerStatus } from './harness-plugin' /** Options for {@link createHarnessMcpServer}. */ -export interface HarnessMcpServerOptions { - /** The host that runs the sessions, from `createHarnessHost`. */ - host: HarnessHost +export interface HarnessMcpServerOptions

{ + /** + * The host that runs the sessions, from `createHarnessHost`. With + * `access`, `access.identify` reads the user of each HTTP request, and the + * host checks each tool call of that user. + */ + host: HarnessHost

/** The harness to serve. */ harness: AnyHarness + /** + * For a host without `access` only: who the tool calls run as. `true`: + * with no principal (one user). `{ principal }`: as that principal, for + * example the one local user of a stdio server. Without host access and + * without this option, the server throws at the start, so it never serves + * every thread to everyone by mistake. + */ + openAccess?: OpenAccess

/** The conversation of a tool call that names no `threadId`. Default `'main'`. */ threadId?: string /** @@ -84,6 +110,12 @@ const MEDIA_URI_PREFIX = 'harness-media://' /** The biggest image or audio file that a result carries inline. */ const INLINE_MEDIA_MAX_BYTES = 5 * 1024 * 1024 +/** The `ctx.context` key of the token that names the user of a request. */ +const CALLER_KEY = 'harnessCaller' + +/** What a tool call of this server gets on `ctx.context`. */ +type CallContext = MCPToolContext> + const threadIdSchema: JSONSchema = { type: 'string', description: @@ -155,7 +187,22 @@ const attachmentsSchema: JSONSchema = { * `harness.expose.commands`, and `tool_` for each tool in * `harness.expose.tools`. * Every tool takes an optional `threadId`. Sessions open with - * `host.open(harness, { threadId })`. + * `host.open(harness, { threadId, principal })`, so the host checks each + * call. + * + * Who makes the calls: + * + * - A host with `access`: the user that `access.identify` finds for each + * HTTP request. `null` gets a 401. A refused action is a tool error. Each + * spec 2025 session belongs to its user, and `handle` replaces the + * `authInfo` of its options for that. Such a server needs HTTP: over stdio + * it throws. The server lists the agents of the harness definition, and + * no commands or agents of plugins, as it opens no session at the start. + * - A host without `access`: the `openAccess` principal, or no principal + * with `openAccess: true`. + * + * The server throws at the start for a host without `access` and without + * `openAccess`, and for a host with both. * * Each interrupt in a result has a `kind`: `approval`, `client-tool`, or * `generic`. `approve` and `reject` answer approvals only. `resolve` answers @@ -171,29 +218,84 @@ const attachmentsSchema: JSONSchema = { * The result is the server from `createMCPServer`. Mount `server.fetch` on * an HTTP route, or pass the server to `serveMCPStdio`. * - * @param options - The host, the harness, the default thread, the approval mode, and the folders for `path` attachments + * @param options - The host, the harness, who the calls run as, the default thread, the approval mode, and the folders for `path` attachments * * @example * ```ts - * const server = await createHarnessMcpServer({ host, harness: assistant }) + * const server = await createHarnessMcpServer({ + * host, + * harness: assistant, + * openAccess: { principal: { id: 'local-user' } }, + * }) * serveMCPStdio(server) * ``` */ -export async function createHarnessMcpServer(options: HarnessMcpServerOptions) { - const { host, harness } = options +export async function createHarnessMcpServer

( + options: HarnessMcpServerOptions

, +) { + const { host, harness, openAccess } = options + const access = host.hasAccess + if (access && openAccess !== undefined) { + throw new Error( + 'createHarnessMcpServer: the host has access, so access.identify reads the user of each request. Remove openAccess.', + ) + } + if (!access && openAccess === undefined) { + throw new Error( + 'createHarnessMcpServer needs host access ({ identify, authorize }) or openAccess: true or { principal }.', + ) + } + const fixed = openAccess === true ? undefined : openAccess?.principal const defaultThread = options.threadId ?? 'main' const approvals = options.approvals ?? 'ask' const filePaths = options.filePaths ?? [] // Work that stopped for a question, by question id. `answer` continues it. const waiting = new Map>() + // The user of each request, by the token in its context. A token that + // this server did not make is in no map, so a context cannot name a user. + const callers = new WeakMap() - const open = (args: unknown) => - host.open(harness, { - threadId: - isRecord(args) && typeof args.threadId === 'string' - ? args.threadId - : defaultThread, - }) + /** + * Who makes a tool call or a resource read: on a host with access, the + * user that `identify` found for its request. Else the `openAccess` one. + */ + const principalOf = (context: Record) => { + if (!access) return fixed + const token = context[CALLER_KEY] + const principal = isRecord(token) ? callers.get(token) : undefined + if (principal === undefined) { + throw new Error( + 'This MCP server reads the user of each HTTP request. Call it with server.fetch or server.handle.', + ) + } + return principal + } + + const open = (threadId: string, context: Record) => + host.open(harness, { threadId, principal: principalOf(context) }) + + /** + * The execute of a tool that works on the session of its thread. The + * session opens for the user of the call, so the host checks each action. + * A refused action becomes a tool error that says so. + */ + const onSession = + ( + work: ( + session: HarnessSession, + args: unknown, + context: CallContext, + ) => Promise, + ) => + async (args: unknown, ctx: ToolExecutionContext) => { + try { + const session = await open(threadOf(args, defaultThread), ctx.context) + return await work(session, args, ctx.context) + } catch (error) { + if (error instanceof AccessDeniedError) throw deniedError(error) + throw error + } + } // Answers the approvals with the approval mode. `undefined` means that the // client must answer. One resolve answers every interrupt, and only the @@ -272,11 +374,10 @@ export async function createHarnessMcpServer(options: HarnessMcpServerOptions) { } async function decideAll( - args: unknown, + session: HarnessSession, approved: boolean, context: MCPToolContext, ) { - const session = await open(args) const interrupts = session.snapshot().pendingInterrupts const other = interrupts.find((interrupt) => !isApproval(interrupt)) if (other !== undefined) { @@ -302,13 +403,14 @@ export async function createHarnessMcpServer(options: HarnessMcpServerOptions) { }, required: ['message'], }), - }).server(async (args, ctx) => { - const session = await open(args) - const message = await userInputOf(session, args, filePaths) - const from = session.snapshot().cursor - const turn = session.send(message, { busy: 'queue' }) - return finish(session, turn, from, ctx.context) - }) + }).server( + onSession(async (session, args, context) => { + const message = await userInputOf(session, args, filePaths) + const from = session.snapshot().cursor + const turn = session.send(message, { busy: 'queue' }) + return finish(session, turn, from, context) + }), + ) const steer = toolDefinition({ name: 'steer', @@ -321,25 +423,28 @@ export async function createHarnessMcpServer(options: HarnessMcpServerOptions) { }, required: ['message'], }), - }).server( - async (args) => - (await open(args)).send(textArg(args, 'message'), { mode: 'steer' }) - .receipt, + }).server( + onSession( + async (session, args) => + session.send(textArg(args, 'message'), { mode: 'steer' }).receipt, + ), ) const cancel = toolDefinition({ name: 'cancel', description: 'Cancel the running turn.', inputSchema: withThreadId(undefined), - }).server(async (args) => (await open(args)).cancel()) + }).server(onSession(async (session) => session.cancel())) const approve = toolDefinition({ name: 'approve', description: 'Approve every tool call that waits for approval, then wait for the turn to continue. When an interrupt of another kind also waits, use resolve.', inputSchema: withThreadId(undefined), - }).server(async (args, ctx) => - decideAll(args, true, ctx.context), + }).server( + onSession(async (session, _args, context) => + decideAll(session, true, context), + ), ) const reject = toolDefinition({ @@ -347,8 +452,10 @@ export async function createHarnessMcpServer(options: HarnessMcpServerOptions) { description: 'Reject every tool call that waits for approval, then wait for the turn to continue. When an interrupt of another kind also waits, use resolve.', inputSchema: withThreadId(undefined), - }).server(async (args, ctx) => - decideAll(args, false, ctx.context), + }).server( + onSession(async (session, _args, context) => + decideAll(session, false, context), + ), ) const resolve = toolDefinition({ @@ -383,12 +490,13 @@ export async function createHarnessMcpServer(options: HarnessMcpServerOptions) { }, required: ['decisions'], }), - }).server(async (args, ctx) => { - const session = await open(args) - const interrupts = session.snapshot().pendingInterrupts - const answers = answersOf(decisionsOf(args), interrupts) - return resume(session, answers, ctx.context) - }) + }).server( + onSession(async (session, args, context) => { + const interrupts = session.snapshot().pendingInterrupts + const answers = answersOf(decisionsOf(args), interrupts) + return resume(session, answers, context) + }), + ) const answer = toolDefinition({ name: 'answer', @@ -402,52 +510,67 @@ export async function createHarnessMcpServer(options: HarnessMcpServerOptions) { }, required: ['questionId', 'value'], }), - }).server(async (args, ctx) => { - const session = await open(args) - const questionId = textArg(args, 'questionId') - const from = session.snapshot().cursor - const receipt = await session.answer({ - questionId, - value: isRecord(args) ? args.value : undefined, - }) - if (receipt.status === 'rejected') { - throw new Error( - `The harness refused the answer: ${receipt.reason ?? 'no reason'}.`, - ) - } - const operation = waiting.get(questionId) - waiting.delete(questionId) - // A question that no tool call waits on: nothing to continue. - if (operation === undefined) { - return { status: 'answered', ...pending(session) } - } - return finish(session, operation, from, ctx.context) - }) + }).server( + onSession(async (session, args, context) => { + const questionId = textArg(args, 'questionId') + const from = session.snapshot().cursor + const receipt = await session.answer({ + questionId, + value: isRecord(args) ? args.value : undefined, + }) + if (receipt.status === 'rejected') { + throw new Error( + `The harness refused the answer: ${receipt.reason ?? 'no reason'}.`, + ) + } + const operation = waiting.get(questionId) + waiting.delete(questionId) + // A question that no tool call waits on: nothing to continue. + if (operation === undefined) { + return { status: 'answered', ...pending(session) } + } + return finish(session, operation, from, context) + }), + ) const status = toolDefinition({ name: 'status', description: 'Show what the harness does now: the status, the interrupts and questions that wait, the background agents, and the queued turns.', inputSchema: withThreadId(undefined), - }).server(async (args) => { - const session = await open(args) - const snapshot = session.snapshot() - return { - status: snapshot.status, - ...pending(session), - agents: snapshot.activeOperations - .filter((operation) => operation.kind === 'agent') - .map(({ id, agent }) => ({ id, agent })), - queuedTurns: snapshot.queuedTurns, - } - }) + }).server( + onSession(async (session) => { + const snapshot = session.snapshot() + return { + status: snapshot.status, + ...pending(session), + agents: snapshot.activeOperations + .filter((operation) => operation.kind === 'agent') + .map(({ id, agent }) => ({ id, agent })), + queuedTurns: snapshot.queuedTurns, + } + }), + ) - // Agents and commands come from a session. Every thread has the same ones. - const session = await host.open(harness, { threadId: defaultThread }) + // Agents and commands come from a session: every thread has the same + // ones. A host with access opens a session only for a user, so the server + // opens none at the start. It then has the agents of the definition, and + // no plugin commands. + const session = access + ? undefined + : await host.open(harness, { threadId: defaultThread, principal: fixed }) + const defined = [ + ...(harness.agents ?? []), + ...(harness.subagents?.agents ?? []), + ] + const agentNamed = (name: string) => + session + ? session.registry.get(name) + : defined.find((agent) => agent.name === name) const exposed: ReadonlyArray = harness.expose?.agents ?? [] const agents = exposed.flatMap((name) => { - const agent = session.registry.get(name) + const agent = agentNamed(name) return agent === undefined ? [] : [{ name, description: agent.description, agent }] @@ -458,21 +581,22 @@ export async function createHarnessMcpServer(options: HarnessMcpServerOptions) { name: toolName, description, inputSchema: withThreadId(convertSchemaToJsonSchema(agent.inputSchema)), - }).server(async (args, ctx) => { - const target = await open(args) - const handle = target.agent(name) - if (handle === undefined) throw new Error(`Unknown agent: ${name}`) - const from = target.snapshot().cursor - return finish(target, handle.start(inputOf(args)), from, ctx.context) - }), + }).server( + onSession(async (target, args, context) => { + const handle = target.agent(name) + if (handle === undefined) throw new Error(`Unknown agent: ${name}`) + const from = target.snapshot().cursor + return finish(target, handle.start(inputOf(args)), from, context) + }), + ), ) // A command can change the session, for example `/mode bypass`, so an MCP // client runs only the commands in `expose.commands`, like a protocol client. const exposedCommands: ReadonlyArray = harness.expose?.commands ?? [] - const commands = session - .describe() - .commands.filter(({ name }) => exposedCommands.includes(name)) + const commands = (session?.describe().commands ?? []).filter(({ name }) => + exposedCommands.includes(name), + ) const commandTools = toolNames('command', commands).map( ({ item: command, toolName, description }) => toolDefinition({ @@ -481,15 +605,16 @@ export async function createHarnessMcpServer(options: HarnessMcpServerOptions) { inputSchema: withThreadId( isJsonSchema(command.input) ? command.input : undefined, ), - }).server(async (args, ctx) => { - const target = await open(args) - const from = target.snapshot().cursor - const operation = target.command({ - name: command.name, - input: inputOf(args), - }) - return finish(target, operation, from, ctx.context) - }), + }).server( + onSession(async (target, args, context) => { + const from = target.snapshot().cursor + const operation = target.command({ + name: command.name, + input: inputOf(args), + }) + return finish(target, operation, from, context) + }), + ), ) // A tool runs with no model, so an MCP client runs only `expose.tools`. @@ -503,12 +628,16 @@ export async function createHarnessMcpServer(options: HarnessMcpServerOptions) { name: toolName, description, inputSchema: withThreadId(convertSchemaToJsonSchema(tool.inputSchema)), - }).server(async (args, ctx) => { - const target = await open(args) - const from = target.snapshot().cursor - const operation = target.tool({ name: tool.name, args: inputOf(args) }) - return finish(target, operation, from, ctx.context) - }), + }).server( + onSession(async (target, args, context) => { + const from = target.snapshot().cursor + const operation = target.tool({ + name: tool.name, + args: inputOf(args), + }) + return finish(target, operation, from, context) + }), + ), ) // A client adds only the record types in `expose.records`, and only to @@ -538,21 +667,22 @@ export async function createHarnessMcpServer(options: HarnessMcpServerOptions) { }, required: ['records'], }), - }).server(async (args) => { - const target = await open(args) - const records = isRecord(args) ? args.records : undefined - const receipt = await applyInput({ - harness, - session: target, - input: parseHarnessInput({ op: 'append', records }), - }) - if (receipt.status === 'rejected') { - throw new Error( - `The harness refused the records: ${receipt.reason}.`, - ) - } - return { appended: Array.isArray(records) ? records.length : 0 } - }), + }).server( + onSession(async (target, args) => { + const records = isRecord(args) ? args.records : undefined + const receipt = await applyInput({ + harness, + session: target, + input: parseHarnessInput({ op: 'append', records }), + }) + if (receipt.status === 'rejected') { + throw new Error( + `The harness refused the records: ${receipt.reason}.`, + ) + } + return { appended: Array.isArray(records) ? records.length : 0 } + }), + ), ] // The bytes of a `resource_link` from a result. The link names its thread, @@ -562,17 +692,22 @@ export async function createHarnessMcpServer(options: HarnessMcpServerOptions) { name: 'media', mimeType: 'application/octet-stream', argsSchema: { parse: mediaAddress }, - }).read(async (_uri, { threadId, id }) => { - const target = await host.open(harness, { threadId }) - const record = await target.getMedia(id) - if (record === null) { - throw new Error(`Media ${id} was not found in thread ${threadId}.`) + }).read(async (_uri, { threadId, id }, ctx) => { + try { + const target = await open(threadId, ctx.context) + const record = await target.getMedia(id) + if (record === null) { + throw new Error(`Media ${id} was not found in thread ${threadId}.`) + } + const blob = toBase64(await target.loadMedia({ id })) + return { blob, mimeType: record.mimeType } + } catch (error) { + if (error instanceof AccessDeniedError) throw deniedError(error) + throw error } - const blob = toBase64(await target.loadMedia({ id })) - return { blob, mimeType: record.mimeType } }) - return createMCPServer({ + const server = createMCPServer({ name: options.name ?? harness.name, version: options.version ?? '1.0.0', // `approvals: 'ask'` asks by elicitation, and elicitation needs a spec 2025 @@ -595,6 +730,63 @@ export async function createHarnessMcpServer(options: HarnessMcpServerOptions) { ...appendTools, ], }) + + /** + * Serves one MCP HTTP request. On a host with access, the calls of the + * request run as the user that `identify` finds for it, and `null` gets a + * 401 with a bearer challenge. + */ + const handle = async (request: Request, handleOptions?: MCPHandleOptions) => { + if (!access) return server.handle(request, handleOptions) + // stdio has no request with a login, so `identify` cannot find a user. + if (isServedOverStdio(served)) { + throw new Error( + 'createHarnessMcpServer: a host with access reads the user of each HTTP request, so it cannot serve over stdio. For one local user, use a host without access and openAccess: { principal }.', + ) + } + const principal = await host.identify({ request }) + if (principal === null) { + return bearerAuthChallengeResponse( + new OAuthError(OAuthErrorCode.InvalidToken, 'The request has no user.'), + ) + } + const token = {} + callers.set(token, principal) + return server.handle(request, { + ...handleOptions, + authInfo: authInfoOf(principal), + context: { ...handleOptions?.context, [CALLER_KEY]: token }, + }) + } + const served = { + ...server, + fetch: (request: Request) => handle(request), + handle, + } + return served +} + +/** + * The `authInfo` of a request of `principal`. The MCP server keeps each spec + * 2025 session for the client of its token, so no other user can use it. + */ +function authInfoOf(principal: Principal) { + const clientId = JSON.stringify([principal.tenantId ?? null, principal.id]) + return { token: '', clientId, scopes: [] } +} + +/** A refused action, as an error that tells the MCP client what it was. */ +function deniedError(error: AccessDeniedError) { + return new Error( + `Access denied: the user of this call may not ${error.action} on thread ${JSON.stringify(error.threadId)}.`, + ) +} + +/** The thread that the args of a tool call name, else `fallback`. */ +function threadOf(args: unknown, fallback: string) { + return isRecord(args) && typeof args.threadId === 'string' + ? args.threadId + : fallback } // The harness needs one resolve that answers every open interrupt. An async diff --git a/packages/ai-mcp/tests/harness-media.test.ts b/packages/ai-mcp/tests/harness-media.test.ts index 4212406253..86e4b5a5e5 100644 --- a/packages/ai-mcp/tests/harness-media.test.ts +++ b/packages/ai-mcp/tests/harness-media.test.ts @@ -6,7 +6,11 @@ import { StreamableHTTPClientTransport, } from '@modelcontextprotocol/client' import { EventType, defineAgent } from '@tanstack/ai' -import { createHarnessHost, defineHarness } from '@tanstack/ai-harness' +import { + createHarnessHost, + defineHarness, + membersOnly, +} from '@tanstack/ai-harness' import { memoryPersistence } from '@tanstack/ai-persistence' import { afterEach, describe, expect, it, vi } from 'vitest' import { z } from 'zod' @@ -173,23 +177,42 @@ afterEach(async () => { type ServerOptions = Omit -/** A real MCP client (spec 2026) on a harness MCP server. */ -async function serve(harness: AnyHarness, options: ServerOptions = {}) { - const host = createHarnessHost({ persistence: memoryPersistence() }) - cleanups.push(() => host.close()) - const server = await createHarnessMcpServer({ host, harness, ...options }) +/** + * A real MCP client (spec 2026) on `server`. `user` goes in the `x-user` + * header of each request. + */ +async function connect( + server: { fetch: (request: Request) => Promise }, + user?: string, +) { const client = new Client( { name: 'tester', version: '1.0.0' }, { versionNegotiation: { mode: { pin: '2026-07-28' } } }, ) const transport = new StreamableHTTPClientTransport(serverUrl, { fetch: async (input, init) => server.fetch(new Request(input, init)), + ...(user === undefined + ? {} + : { requestInit: { headers: { 'x-user': user } } }), }) await client.connect(transport) cleanups.push(() => client.close()) return client } +/** A real MCP client (spec 2026) on a harness MCP server. */ +async function serve(harness: AnyHarness, options: ServerOptions = {}) { + const host = createHarnessHost({ persistence: memoryPersistence() }) + cleanups.push(() => host.close()) + const server = await createHarnessMcpServer({ + host, + harness, + openAccess: true, + ...options, + }) + return connect(server) +} + /** A chat harness on an MCP server. `calls` has the messages of each model call. */ async function chatServer(options: ServerOptions = {}) { const model = scripted() @@ -423,22 +446,27 @@ describe('media in results', () => { ]) }) + /** A harness whose agent `painter` paints an image over 5 MB. */ + const bigPainter = () => + defineHarness({ + name: 'test/mcp-media', + adapter: scripted().adapter, + agents: [painter(bigImage)], + expose: { agents: ['painter'] }, + }) + + /** The media link that a result of `agent_painter` has. */ + const linkOf = (reply: { content?: unknown }) => + linkShape.parse(z.array(z.unknown()).parse(reply.content)[1]).uri + /** Runs `agent_painter` with an image over 5 MB. `uri` is its media link. */ async function paintBig() { - const client = await serve( - defineHarness({ - name: 'test/mcp-media', - adapter: scripted().adapter, - agents: [painter(bigImage)], - expose: { agents: ['painter'] }, - }), - ) + const client = await serve(bigPainter()) const reply = await client.callTool({ name: 'agent_painter', arguments: {}, }) - const link = z.array(z.unknown()).parse(reply.content)[1] - return { client, reply, uri: linkShape.parse(link).uri } + return { client, reply, uri: linkOf(reply) } } it('returns an image over 5 MB as a resource link', async () => { @@ -476,4 +504,38 @@ describe('media in results', () => { `Media ${id} was not found in thread other.`, ) }) + + it('reads a media link on a host with access only for a user who may read the thread', async () => { + const olga = { id: 'olga', tenantId: 'acme' } + const cid = { id: 'cid', tenantId: 'acme' } + const users = new Map([olga, cid].map((user) => [user.id, user])) + const host = createHarnessHost({ + persistence: memoryPersistence(), + access: { + identify: ({ request }) => + users.get(request.headers.get('x-user') ?? '') ?? null, + authorize: membersOnly(), + }, + }) + cleanups.push(() => host.close()) + const server = await createHarnessMcpServer({ + host, + harness: bigPainter(), + }) + const asOlga = await connect(server, 'olga') + const asCid = await connect(server, 'cid') + const reply = await asOlga.callTool({ + name: 'agent_painter', + arguments: { threadId: 't1' }, + }) + const uri = linkOf(reply) + + const read = await asOlga.readResource({ uri }) + expect(read.contents).toEqual([ + expect.objectContaining({ uri, mimeType: 'image/png' }), + ]) + await expect(asCid.readResource({ uri })).rejects.toThrow( + 'Access denied: the user of this call may not read on thread "t1".', + ) + }) }) diff --git a/packages/ai-mcp/tests/harness-stdio.test.ts b/packages/ai-mcp/tests/harness-stdio.test.ts new file mode 100644 index 0000000000..31407dd6a7 --- /dev/null +++ b/packages/ai-mcp/tests/harness-stdio.test.ts @@ -0,0 +1,76 @@ +import { + createHarnessHost, + defineHarness, + membersOnly, +} from '@tanstack/ai-harness' +import { memoryPersistence } from '@tanstack/ai-persistence' +import { describe, expect, it, vi } from 'vitest' +import { createHarnessMcpServer } from '../src/harness' +import { serveMCPStdio } from '../src/server/stdio' + +// The stdio transport of the SDK reads the real stdin. This fake gets the +// messages of the test, and keeps what the server sends back. +const fake = vi.hoisted(() => { + class FakeTransport { + sent: Array = [] + onmessage?: (message: unknown) => void + onerror?: (error: Error) => void + constructor() { + fake.transports.push(this) + } + start() { + return Promise.resolve() + } + send(message: unknown) { + this.sent.push(message) + return Promise.resolve() + } + close() { + return Promise.resolve() + } + } + return { + FakeTransport, + transports: [] as Array>, + } +}) + +vi.mock('@modelcontextprotocol/server/stdio', () => ({ + StdioServerTransport: fake.FakeTransport, +})) + +describe('createHarnessMcpServer over stdio', () => { + it('refuses each message for a host with access, and says why on stderr', async () => { + const errors = vi.spyOn(console, 'error').mockImplementation(() => {}) + const host = createHarnessHost({ + persistence: memoryPersistence(), + access: { identify: () => null, authorize: membersOnly() }, + }) + const server = await createHarnessMcpServer({ + host, + harness: defineHarness({ name: 'test/stdio' }), + }) + const served = serveMCPStdio(server) + const transport = fake.transports.at(-1) + try { + transport?.onmessage?.({ jsonrpc: '2.0', id: 1, method: 'tools/list' }) + + await vi.waitFor(() => + expect(transport?.sent).toEqual([ + { + jsonrpc: '2.0', + id: 1, + error: { code: -32603, message: 'Internal server error' }, + }, + ]), + ) + expect(errors).toHaveBeenCalledWith( + 'createHarnessMcpServer: a host with access reads the user of each HTTP request, so it cannot serve over stdio. For one local user, use a host without access and openAccess: { principal }.', + ) + } finally { + await served.close() + await host.close() + errors.mockRestore() + } + }) +}) diff --git a/packages/ai-mcp/tests/harness.test.ts b/packages/ai-mcp/tests/harness.test.ts index 471baa67f7..d1351c6e83 100644 --- a/packages/ai-mcp/tests/harness.test.ts +++ b/packages/ai-mcp/tests/harness.test.ts @@ -15,6 +15,8 @@ import { defineCommand, defineHarness, definePlugin, + logIdOf, + membersOnly, } from '@tanstack/ai-harness' import { memoryLogStore, memoryPersistence } from '@tanstack/ai-persistence' import { afterEach, describe, expect, it } from 'vitest' @@ -379,10 +381,34 @@ type ServerOptions = Omit async function serve(harness: AnyHarness, options: ServerOptions = {}) { const host = createHarnessHost({ persistence: memoryPersistence() }) cleanups.push(() => host.close()) - const server = await createHarnessMcpServer({ host, harness, ...options }) + const server = await createHarnessMcpServer({ + host, + harness, + openAccess: true, + ...options, + }) return { server, host } } +const olga = { id: 'olga', tenantId: 'acme' } +const bea = { id: 'bea', tenantId: 'acme' } +const cid = { id: 'cid', tenantId: 'acme' } +const users = new Map([olga, bea, cid].map((user) => [user.id, user])) + +/** A membersOnly() host whose users come from the `x-user` header. */ +function accessHost() { + const host = createHarnessHost({ + persistence: memoryPersistence(), + access: { + identify: ({ request }) => + users.get(request.headers.get('x-user') ?? '') ?? null, + authorize: membersOnly(), + }, + }) + cleanups.push(() => host.close()) + return host +} + /** * A harness on an MCP server. With `reviewAt`, a middleware stops the turn at * that phase with a `review` interrupt. @@ -464,12 +490,13 @@ function accepted(value: string): ElicitReply { * A real MCP client on `server.fetch`. With `reply`, the client supports * elicitation and gives that reply to every question (spec 2025 sessions * only). Without it, the client speaks spec 2026 and has no elicitation. + * `user` goes in the `x-user` header of each request. */ async function connect( server: { fetch: (request: Request) => Promise }, - options: { reply?: ElicitReply } = {}, + options: { reply?: ElicitReply; user?: string } = {}, ) { - const reply = options.reply + const { reply, user } = options const client = new Client( { name: 'tester', version: '1.0.0' }, reply === undefined @@ -485,6 +512,9 @@ async function connect( } const transport = new StreamableHTTPClientTransport(serverUrl, { fetch: async (input, init) => server.fetch(new Request(input, init)), + ...(user === undefined + ? {} + : { requestInit: { headers: { 'x-user': user } } }), }) await client.connect(transport) cleanups.push(() => client.close()) @@ -589,7 +619,11 @@ describe('createHarnessMcpServer', () => { persistence: { stores: { log, runs, metadata } }, }) cleanups.push(() => host.close()) - const server = await createHarnessMcpServer({ host, harness }) + const server = await createHarnessMcpServer({ + host, + harness, + openAccess: true, + }) const { client } = await connect(server) const ok = await client.callTool({ @@ -597,7 +631,7 @@ describe('createHarnessMcpServer', () => { arguments: { threadId: 't-sig', records: [{ type: 'app.signal', n: 1 }] }, }) expect(ok.isError).toBeFalsy() - const stored = await log.read({ logId: 't-sig' }) + const stored = await log.read({ logId: logIdOf({ logId: 't-sig' }) }) expect(JSON.stringify(stored)).toContain('"type":"app.signal"') const refused = await client.callTool({ name: 'append', @@ -1459,3 +1493,136 @@ describe('createHarnessMcpServer', () => { }, ) }) + +describe('createHarnessMcpServer access', () => { + /** A harness whose model answers the first call with "Hello.". */ + const greeter = () => + defineHarness({ + name: 'test/access', + adapter: scripted([textTurn('Hello.')]).adapter, + }) + + it('throws at the start for a host without access and without openAccess', async () => { + const host = createHarnessHost({ persistence: memoryPersistence() }) + cleanups.push(() => host.close()) + + await expect( + createHarnessMcpServer({ host, harness: greeter() }), + ).rejects.toThrow( + 'createHarnessMcpServer needs host access ({ identify, authorize }) or openAccess: true or { principal }.', + ) + }) + + it('throws at the start for openAccess on a host with access', async () => { + const server = createHarnessMcpServer({ + host: accessHost(), + harness: greeter(), + openAccess: { principal: olga }, + }) + + await expect(server).rejects.toThrow( + 'createHarnessMcpServer: the host has access, so access.identify reads the user of each request. Remove openAccess.', + ) + }) + + it('runs every call as the openAccess principal', async () => { + const ana = { id: 'ana', tenantId: 'globex' } + const { server, host } = await serve(greeter(), { + openAccess: { principal: ana }, + }) + const { client } = await connect(server) + + await client.callTool({ + name: 'chat', + arguments: { threadId: 't-ana', message: 'Hi' }, + }) + expect( + await host.sessions.get({ threadId: 't-ana', principal: ana }), + ).toMatchObject({ tenantId: 'globex', ownerId: 'ana' }) + expect(await host.sessions.get({ threadId: 't-ana' })).toBeUndefined() + }) + + it('answers a request that identify refuses with a 401 bearer challenge', async () => { + const server = await createHarnessMcpServer({ + host: accessHost(), + harness: greeter(), + }) + + const response = await server.fetch( + new Request(serverUrl, { + method: 'POST', + headers: { + 'content-type': 'application/json', + accept: 'application/json, text/event-stream', + 'x-user': 'nobody', + }, + body: JSON.stringify({ jsonrpc: '2.0', id: 1, method: 'tools/list' }), + }), + ) + expect(response.status).toBe(401) + expect(response.headers.get('www-authenticate')).toMatch(/^Bearer /) + }) + + it('runs each request as its user, and returns a refused action as a tool error', async () => { + const host = accessHost() + const server = await createHarnessMcpServer({ host, harness: greeter() }) + const { client: asOlga } = await connect(server, { user: 'olga' }) + const { client: asBea } = await connect(server, { user: 'bea' }) + const { client: asCid } = await connect(server, { user: 'cid' }) + + const reply = await asOlga.callTool({ + name: 'chat', + arguments: { threadId: 't1', message: 'Hi' }, + }) + expect(reply.structuredContent).toMatchObject({ + status: 'completed', + text: 'Hello.', + }) + expect( + await host.sessions.get({ threadId: 't1', principal: olga }), + ).toMatchObject({ ownerId: 'olga' }) + await host.sessions.share({ + threadId: 't1', + userId: 'bea', + role: 'viewer', + principal: olga, + }) + + const stranger = await asCid.callTool({ + name: 'status', + arguments: { threadId: 't1' }, + }) + expect(stranger).toMatchObject( + toolError( + 'Access denied: the user of this call may not read on thread "t1".', + ), + ) + const viewer = await asBea.callTool({ + name: 'chat', + arguments: { threadId: 't1', message: 'Hi' }, + }) + expect(viewer).toMatchObject( + toolError( + 'Access denied: the user of this call may not send on thread "t1".', + ), + ) + }) + + it('lists the agents of the harness definition on a host with access, and no plugin commands', async () => { + const harness = defineHarness({ + name: 'test/access-tools', + adapter: scripted([]).adapter, + agents: [pricer], + expose: { agents: ['pricer'], commands: ['greet'] }, + plugins: () => [commands], + }) + const server = await createHarnessMcpServer({ host: accessHost(), harness }) + const { client } = await connect(server, { user: 'olga' }) + + const listed = await client.listTools() + expect(listed.tools.map((tool) => tool.name).sort()).toEqual([ + 'agent_pricer', + ...controlTools, + ]) + }) +}) diff --git a/packages/ai-persistence/src/index.ts b/packages/ai-persistence/src/index.ts index aa36ae843c..44ab4a5dcd 100644 --- a/packages/ai-persistence/src/index.ts +++ b/packages/ai-persistence/src/index.ts @@ -53,9 +53,11 @@ export type { TurnLeaseKey, WorkClaimStore, SessionIndexEntry, + SessionIndexKey, SessionIndexListOptions, SessionIndexPage, SessionIndexStore, + SessionMember, // Named product shapes (prefer these over a sparse bag) ChatTranscriptStores, ChatPersistenceStores, diff --git a/packages/ai-persistence/src/memory.ts b/packages/ai-persistence/src/memory.ts index 243f138d5f..17da64f5fc 100644 --- a/packages/ai-persistence/src/memory.ts +++ b/packages/ai-persistence/src/memory.ts @@ -35,8 +35,8 @@ import type { RunRecord, RunStore, SessionIndexEntry, + SessionIndexKey, SessionIndexListOptions, - SessionIndexPage, SessionIndexStore, WorkClaimStore, LeaseStore, @@ -1291,61 +1291,58 @@ function parseSessionCursor(cursor: string) { } class MemorySessionIndexStore implements SessionIndexStore { - private readonly entries = new Map< + // Nested maps, tenant then thread. A joined `${tenant}:${thread}` key would + // make `('a:b', 'c')` and `('a', 'b:c')` the same entry. A missing tenant is + // the default tenant, `''`. + private readonly tenants = new Map< string, - { entry: SessionIndexEntry; revision: string } + Map >() // 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) + get(key: SessionIndexKey) { + const stored = this.find(key) return Promise.resolve(stored && structuredClone(stored.entry)) } - getVersioned( - threadId: string, - ): Promise<{ entry: SessionIndexEntry; revision: string } | undefined> { - const stored = this.entries.get(threadId) + getVersioned(key: SessionIndexKey) { + const stored = this.find(key) return Promise.resolve(stored && structuredClone(stored)) } + upsert(entry: SessionIndexEntry) { + this.write(entry) + return Promise.resolve() + } upsertIf({ entry, - expectedRevision, + revision, }: { 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' }) + revision: string | null + }) { + const current = this.find(entry)?.revision ?? null + if (current !== revision) { + return Promise.resolve({ ok: false, reason: 'conflict' } as const) } - return Promise.resolve({ ok: true, revision: this.write(entry) }) + return Promise.resolve({ ok: true, revision: this.write(entry) } as const) } - list(options: SessionIndexListOptions = {}): Promise { - const { limit, cursor, parentThreadId, principal, search, harness } = - options + list(options: SessionIndexListOptions = {}) { + const { limit, cursor, parentThreadId, memberId, 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) + const tenant = this.tenants.get(options.tenantId ?? '') + const stored = [...(tenant?.values() ?? [])] + const matching = stored + .map(({ entry }) => 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 isMember = + memberId === undefined || + entry.members.some((member) => member.userId === memberId) const isAfterCursor = after === undefined || bySessionOrder(after, entry) < 0 const isFound = @@ -1357,7 +1354,7 @@ class MemorySessionIndexStore implements SessionIndexStore { ) return ( isChild && - isOwned && + isMember && isFound && isHarness && hasMetadata && @@ -1376,17 +1373,27 @@ class MemorySessionIndexStore implements SessionIndexStore { : {}), }) } - delete(threadId: string): Promise { - this.entries.delete(threadId) + delete(key: SessionIndexKey) { + const tenantId = key.tenantId ?? '' + const tenant = this.tenants.get(tenantId) + if (!tenant) return Promise.resolve() + tenant.delete(key.threadId) + if (tenant.size === 0) this.tenants.delete(tenantId) return Promise.resolve() } + private find(key: SessionIndexKey) { + return this.tenants.get(key.tenantId ?? '')?.get(key.threadId) + } // Copies in and out, so a caller never changes a stored entry. - private write(entry: SessionIndexEntry): string { + private write(entry: SessionIndexEntry) { + const tenantId = entry.tenantId ?? '' + let tenant = this.tenants.get(tenantId) + if (!tenant) { + tenant = new Map() + this.tenants.set(tenantId, tenant) + } const revision = String(++this.lastRevision) - this.entries.set(entry.threadId, { - entry: structuredClone(entry), - revision, - }) + tenant.set(entry.threadId, { entry: structuredClone(entry), revision }) return revision } } diff --git a/packages/ai-persistence/src/middleware.ts b/packages/ai-persistence/src/middleware.ts index 7f3dbdc723..1c8b6ecc0b 100644 --- a/packages/ai-persistence/src/middleware.ts +++ b/packages/ai-persistence/src/middleware.ts @@ -2260,6 +2260,12 @@ export interface WithPersistenceOptions { * interval for its transcript writes, with or without `snapshotStreaming`. */ snapshotIntervalMs?: number + /** + * The tenant of the thread, for `stores.sessions`. The middleware reads the + * parent entry of a subagent child from this tenant and writes the child + * entry in it. Without it, the default tenant. + */ + tenantId?: string } /** @@ -2334,6 +2340,7 @@ export function withPersistence( ? { interrupts: persistence.stores.interrupts } : {}), ...(sessions ? { sessions } : {}), + ...(options.tenantId !== undefined ? { tenantId: options.tenantId } : {}), }) function resetStreaming(state: RunStateEntry) { diff --git a/packages/ai-persistence/src/subagent-runs.ts b/packages/ai-persistence/src/subagent-runs.ts index b05b190183..be1610eeb0 100644 --- a/packages/ai-persistence/src/subagent-runs.ts +++ b/packages/ai-persistence/src/subagent-runs.ts @@ -299,6 +299,12 @@ export function createSubagentRunRecorder(stores: { interrupts?: InterruptStore /** When present, each child gets a session index entry. */ sessions?: SessionIndexStore + /** + * The tenant of the threads, for the session index. The parent entry is + * read from it and each child entry is written in it. Without it, the + * default tenant. + */ + tenantId?: string /** Minimum milliseconds between writes while a child streams. */ intervalMs?: number }) { @@ -529,41 +535,55 @@ export function createSubagentRunRecorder(stores: { 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. + // One index entry per child, in the tenant of the recorder. A write + // replaces the whole entry, so keep the fields that other writers set, like + // a title. `upsertIf` writes only over the entry that was read, and a + // conflict reads again, so a change from another host is not lost. The + // child has the owner and the `harness` name of its parent entry, so a list + // filtered by them shows it. It has no members: access to a child goes + // through its top-level thread. 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) + const tenant = + stores.tenantId === undefined ? {} : { tenantId: stores.tenantId } + const threadId = 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, + const parent = await sessions.get({ + ...tenant, + threadId: note.parentThreadId, }) + for (let attempt = 0; attempt < 5; attempt += 1) { + const current = await sessions.getVersioned({ ...tenant, threadId }) + const now = Date.now() + const written = await sessions.upsertIf({ + entry: { + members: [], + ...current?.entry, + ...tenant, + threadId, + parentThreadId: note.parentThreadId, + ...(note.parentToolCallId !== undefined && { + parentToolCallId: note.parentToolCallId, + }), + ...(parent?.ownerId !== undefined && { ownerId: parent.ownerId }), + ...(parent?.harness !== undefined && { harness: parent.harness }), + createdAt: current?.entry.createdAt ?? now, + updatedAt: now, + }, + revision: current?.revision ?? null, + }) + if (written.ok) return + } + console.warn( + `[ai-persistence] session index entry '${threadId}' changed on another host 5 times in a row`, + ) } catch (error) { console.warn( - `[ai-persistence] session index write failed for '${childThreadId}'`, + `[ai-persistence] session index write failed for '${threadId}'`, error, ) } diff --git a/packages/ai-persistence/src/testkit/conformance.ts b/packages/ai-persistence/src/testkit/conformance.ts index 9785a50b3c..3df333dab2 100644 --- a/packages/ai-persistence/src/testkit/conformance.ts +++ b/packages/ai-persistence/src/testkit/conformance.ts @@ -3060,31 +3060,51 @@ export function runPersistenceConformance( }) }) - // Each case uses its own ids, because the suite shares one persistence. + // The suite shares one persistence. Most cases write in a tenant of their + // own, so a list has only the entries of the case. The cases about the + // default tenant use new thread ids. describe('sessions', () => { const newId = () => `sessions-${crypto.randomUUID()}` const sessionEntry = ( threadId: string, updatedAt: number, extra: Partial = {}, - ) => ({ threadId, createdAt: 1, updatedAt, ...extra }) + ) => ({ threadId, members: [], createdAt: 1, updatedAt, ...extra }) + /** A new tenant, with an entry and a key factory for it. */ + const newTenant = () => { + const tenantId = newId() + return { + tenantId, + entry: ( + threadId: string, + updatedAt: number, + extra: Partial = {}, + ) => sessionEntry(threadId, updatedAt, { tenantId, ...extra }), + key: (threadId: string) => ({ tenantId, threadId }), + } + } 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 { tenantId, key } = newTenant() const written: SessionIndexEntry = { - threadId, + tenantId, + threadId: 'thread', + ownerId: 'user-1', + members: [ + { userId: 'user-1', role: 'owner' }, + { userId: 'user-2', role: 'viewer' }, + ], harness: 'coder', title: 'Fix the build', - parentThreadId: `${threadId}-parent`, + parentThreadId: 'parent', parentToolCallId: 'call-1', createdAt: 1, updatedAt: 2, pinned: true, - principal: { id: 'user-1', tenantId: 'org-1' }, usage: { turns: 2, promptTokens: 10, @@ -3099,23 +3119,80 @@ export function runPersistenceConformance( await store.upsert(written) - expect(await store.get(threadId)).toEqual(written) - expect(await store.get(newId())).toBeUndefined() + expect(await store.get(key('thread'))).toEqual(written) + expect(await store.get(key('unknown'))).toBeUndefined() }) - it('replaces the whole entry on a second upsert', async (ctx) => { + it('keeps the owner and the members in getVersioned and list', async (ctx) => { + const store = persistence.stores.sessions + if (!store) return ctx.skip('sessions store not provided') + const { tenantId, entry, key } = newTenant() + const shared = entry('shared', 1, { + ownerId: 'alice', + members: [ + { userId: 'alice', role: 'owner' }, + { userId: 'bob', role: 'member' }, + { userId: 'carol', role: 'viewer' }, + ], + }) + + await store.upsert(shared) + + expect((await store.getVersioned(key('shared')))?.entry).toEqual(shared) + expect((await store.list({ tenantId })).entries).toEqual([shared]) + }) + + it('keeps one thread id in two tenants as two entries', 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 }), + const tenantId = newId() + const inTenant = sessionEntry(threadId, 10, { tenantId, title: 'a' }) + const inDefault = sessionEntry(threadId, 20, { title: 'b' }) + + await store.upsert(inTenant) + await store.upsert(inDefault) + + expect(await store.get({ tenantId, threadId })).toEqual(inTenant) + expect(await store.get({ threadId })).toEqual(inDefault) + expect((await store.list({ tenantId })).entries).toEqual([inTenant]) + + await store.delete({ tenantId, threadId }) + + expect(await store.get({ tenantId, threadId })).toBeUndefined() + expect(await store.get({ threadId })).toEqual(inDefault) + }) + + it('lists only the default tenant without tenantId', async (ctx) => { + const store = persistence.stores.sessions + if (!store) return ctx.skip('sessions store not provided') + const id = newId() + await store.upsert(sessionEntry(`${id}-tenant`, 10, { tenantId: id })) + await store.upsert(sessionEntry(`${id}-default`, 10)) + + const page = await store.list() + + expect(threadIds(page)).toContain(`${id}-default`) + // `''` is the default tenant too. + const otherTenants = page.entries.filter( + (entry) => (entry.tenantId ?? '') !== '', ) + expect(otherTenants).toEqual([]) + }) + + 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 { tenantId, entry, key } = newTenant() + await store.upsert(entry('thread', 2, { title: 'old', pinned: true })) - await store.upsert(sessionEntry(threadId, 3, { title: 'new' })) + await store.upsert(entry('thread', 3, { title: 'new' })) // `pinned` is gone: upsert does not merge. - expect(await store.get(threadId)).toEqual({ - threadId, + expect(await store.get(key('thread'))).toEqual({ + tenantId, + threadId: 'thread', + members: [], createdAt: 1, updatedAt: 3, title: 'new', @@ -3125,215 +3202,178 @@ export function runPersistenceConformance( 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 })) + const { tenantId, entry } = newTenant() + await store.upsert(entry('c', 20)) + await store.upsert(entry('a', 30)) + await store.upsert(entry('b', 20)) // Newest first. The same `updatedAt` goes in `threadId` order. - expect(threadIds(await store.list({ principal: owner }))).toEqual([ - a, - b, - c, + expect(threadIds(await store.list({ tenantId }))).toEqual([ + 'a', + 'b', + 'c', ]) - const first = await store.list({ principal: owner, limit: 2 }) - expect(threadIds(first)).toEqual([a, b]) + const first = await store.list({ tenantId, limit: 2 }) + expect(threadIds(first)).toEqual(['a', 'b']) expect(first.truncated).toBe(true) const second = await store.list({ - principal: owner, + tenantId, limit: 2, cursor: required(first.cursor, 'the cursor of a truncated page'), }) - expect(threadIds(second)).toEqual([c]) + expect(threadIds(second)).toEqual(['c']) expect(second.truncated).toBeFalsy() expect(second.cursor).toBeUndefined() }) - it('filters by parent thread and by owner', async (ctx) => { + it('filters by parent thread', 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 })) + const { tenantId, entry } = newTenant() + await store.upsert(entry('parent', 10)) await store.upsert( - sessionEntry(child, 20, { - principal: owner, - parentThreadId: parent, + entry('child', 20, { + 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' }, - }), - ) + await store.upsert(entry('sibling', 30, { parentThreadId: 'parent' })) + await store.upsert(entry('other', 40, { parentThreadId: 'other' })) - 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([]) + threadIds(await store.list({ tenantId, parentThreadId: 'parent' })), + ).toEqual(['sibling', 'child']) }) 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 }), - ) + const { tenantId, entry } = newTenant() + await store.upsert(entry('parent', 10)) + await store.upsert(entry('child', 20, { parentThreadId: 'parent' })) expect( - threadIds( - await store.list({ principal: owner, parentThreadId: null }), - ), - ).toEqual([parent]) + threadIds(await store.list({ tenantId, parentThreadId: null })), + ).toEqual(['parent']) }) - it('filters by title search, without case', async (ctx) => { + it('filters by member id, with any role', 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` + const { tenantId, entry } = newTenant() + const bob = { userId: 'bob', role: 'owner' } as const await store.upsert( - sessionEntry(build, 10, { principal: owner, title: 'Fix the Build' }), + entry('owned', 10, { + ownerId: 'alice', + members: [{ userId: 'alice', role: 'owner' }], + }), + ) + await store.upsert( + entry('member', 20, { + ownerId: 'bob', + members: [bob, { userId: 'alice', role: 'member' }], + }), + ) + await store.upsert( + entry('viewer', 30, { + ownerId: 'bob', + members: [bob, { userId: 'alice', role: 'viewer' }], + }), ) await store.upsert( - sessionEntry(docs, 20, { principal: owner, title: 'Write docs' }), + entry('other', 40, { ownerId: 'bob', members: [bob] }), + ) + // A subagent child has the owner of its parent and no members. + await store.upsert( + entry('child', 50, { ownerId: 'alice', parentThreadId: 'owned' }), ) - 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. + threadIds(await store.list({ tenantId, memberId: 'alice' })), + ).toEqual(['viewer', 'member', 'owned']) + }) + + it('filters by title search, without case', async (ctx) => { + const store = persistence.stores.sessions + if (!store) return ctx.skip('sessions store not provided') + const { tenantId, entry } = newTenant() + await store.upsert(entry('build', 10, { title: 'Fix the Build' })) + await store.upsert(entry('docs', 20, { title: 'Write docs' })) + await store.upsert(entry('untitled', 30)) + expect( - threadIds(await store.list({ principal: owner, search: '' })), - ).toEqual([docs, build]) + threadIds(await store.list({ tenantId, search: 'BUILD' })), + ).toEqual(['build']) + // An entry with no title never matches. + expect(threadIds(await store.list({ tenantId, 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 })) + const { tenantId, entry } = newTenant() + await store.upsert(entry('coder', 10, { harness: 'coder' })) + await store.upsert(entry('writer', 20, { harness: 'writer' })) + await store.upsert(entry('none', 30)) expect( - threadIds(await store.list({ principal: owner, harness: 'coder' })), - ).toEqual([coder]) + threadIds(await store.list({ tenantId, 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` + const { tenantId, entry } = newTenant() await store.upsert( - sessionEntry(both, 10, { - principal: owner, - metadata: { project: 'p-1', branch: 'main' }, - }), + entry('both', 10, { metadata: { project: 'p-1', branch: 'main' } }), ) await store.upsert( - sessionEntry(projectOnly, 20, { - principal: owner, - metadata: { project: 'p-1' }, - }), + entry('project-only', 20, { metadata: { project: 'p-1' } }), ) await store.upsert( - sessionEntry(other, 30, { - principal: owner, - metadata: { project: 'p-2', branch: 'main' }, - }), + entry('other', 30, { metadata: { project: 'p-2', branch: 'main' } }), ) - await store.upsert(sessionEntry(none, 40, { principal: owner })) + await store.upsert(entry('none', 40)) expect( threadIds( - await store.list({ - principal: owner, - metadata: { project: 'p-1' }, - }), + await store.list({ tenantId, metadata: { project: 'p-1' } }), ), - ).toEqual([projectOnly, both]) + ).toEqual(['project-only', 'both']) expect( threadIds( await store.list({ - principal: owner, + tenantId, metadata: { project: 'p-1', branch: 'main' }, }), ), - ).toEqual([both]) + ).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 { tenantId, entry } = newTenant() const base = { - principal: owner, title: 'Fix the build', harness: 'coder', metadata: { project: 'p-1' }, } - await store.upsert(sessionEntry(match, 10, base)) + await store.upsert(entry('match', 10, base)) await store.upsert( - sessionEntry(wrongTitle, 20, { ...base, title: 'Write docs' }), + entry('wrong-title', 20, { ...base, title: 'Write docs' }), ) await store.upsert( - sessionEntry(wrongHarness, 30, { ...base, harness: 'writer' }), + entry('wrong-harness', 30, { ...base, harness: 'writer' }), ) await store.upsert( - sessionEntry(wrongProject, 40, { + entry('wrong-project', 40, { ...base, metadata: { project: 'p-2' }, }), @@ -3342,105 +3382,83 @@ export function runPersistenceConformance( expect( threadIds( await store.list({ - principal: owner, + tenantId, search: 'build', harness: 'coder', metadata: { project: 'p-1' }, }), ), - ).toEqual([match]) + ).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]) + const { tenantId, entry } = newTenant() + await store.upsert(entry('a', 50, { title: 'build a' })) + await store.upsert(entry('skip', 40, { title: 'docs' })) + await store.upsert(entry('b', 30, { title: 'build b' })) + await store.upsert(entry('c', 20, { title: 'build c' })) + + const first = await store.list({ tenantId, search: 'build', limit: 2 }) + expect(threadIds(first)).toEqual(['a', 'b']) expect(first.truncated).toBe(true) const second = await store.list({ - principal: owner, + tenantId, search: 'build', limit: 2, cursor: required(first.cursor, 'the cursor of a truncated page'), }) - expect(threadIds(second)).toEqual([c]) + expect(threadIds(second)).toEqual(['c']) expect(second.truncated).toBeFalsy() }) - it('deletes an entry, and ignores an unknown thread', async (ctx) => { + it('deletes an entry for good, 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, - ]) + const { tenantId, entry, key } = newTenant() + await store.upsert(entry('kept', 1)) + await store.upsert(entry('removed', 2, { title: 'old' })) + + await store.delete(key('removed')) + await store.delete(key('unknown')) + + expect(await store.get(key('removed'))).toBeUndefined() + expect(await store.getVersioned(key('removed'))).toBeUndefined() + expect(threadIds(await store.list({ tenantId }))).toEqual(['kept']) + // Nothing of the old entry is left, so "only when absent" writes. + const recreated = await store.upsertIf({ + entry: entry('removed', 3), + revision: null, + }) + expect(recreated.ok).toBe(true) + expect(await store.get(key('removed'))).toEqual(entry('removed', 3)) }) - // 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() + if (!store) return ctx.skip('sessions store not provided') + const { entry, key } = newTenant() + expect(await store.getVersioned(key('thread'))).toBeUndefined() // `null`: only when absent. const created = await store.upsertIf({ - entry: sessionEntry(threadId, 1), - expectedRevision: null, + entry: entry('thread', 1), + revision: null, }) if (!created.ok) throw new Error('the first upsertIf conflicted') expect( - await store.upsertIf({ - entry: sessionEntry(threadId, 9), - expectedRevision: null, - }), - ).toEqual({ - ok: false, - reason: 'conflict', - }) - expect(await store.getVersioned(threadId)).toEqual({ - entry: sessionEntry(threadId, 1), + await store.upsertIf({ entry: entry('thread', 9), revision: null }), + ).toEqual({ ok: false, reason: 'conflict' }) + expect(await store.getVersioned(key('thread'))).toEqual({ + entry: entry('thread', 1), revision: created.revision, }) const changed = await store.upsertIf({ - entry: sessionEntry(threadId, 2, { title: 'new' }), - expectedRevision: created.revision, + entry: entry('thread', 2, { title: 'new' }), + revision: created.revision, }) if (!changed.ok) throw new Error('upsertIf at the read revision conflicted') @@ -3449,51 +3467,51 @@ export function runPersistenceConformance( // A stale revision conflicts and leaves the entry as it was. expect( await store.upsertIf({ - entry: sessionEntry(threadId, 3), - expectedRevision: created.revision, + entry: entry('thread', 3), + revision: created.revision, }), ).toEqual({ ok: false, reason: 'conflict' }) - expect(await store.get(threadId)).toEqual( - sessionEntry(threadId, 2, { title: 'new' }), + expect(await store.get(key('thread'))).toEqual( + entry('thread', 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') + if (!store) return ctx.skip('sessions store not provided') + const { entry, key } = newTenant() + await store.upsert(entry('thread', 1)) + const read = required( + await store.getVersioned(key('thread')), + 'the versioned entry', + ) - await store.upsert(sessionEntry(threadId, 2)) + await store.upsert(entry('thread', 2)) expect( await store.upsertIf({ - entry: sessionEntry(threadId, 3), - expectedRevision: read.revision, + entry: entry('thread', 3), + revision: read.revision, }), ).toEqual({ ok: false, reason: 'conflict' }) - expect((await store.get(threadId))?.updatedAt).toBe(2) + expect((await store.get(key('thread')))?.updatedAt).toBe(2) // A deleted revision does not match the entry written after it. const last = required( - await store.getVersioned(threadId), + await store.getVersioned(key('thread')), 'the versioned entry', ) - await store.delete(threadId) - expect(await store.getVersioned(threadId)).toBeUndefined() + await store.delete(key('thread')) expect( await store.upsertIf({ - entry: sessionEntry(threadId, 4), - expectedRevision: last.revision, + entry: entry('thread', 4), + revision: last.revision, }), ).toEqual({ ok: false, reason: 'conflict' }) - expect(await store.get(threadId)).toBeUndefined() + expect(await store.get(key('thread'))).toBeUndefined() const recreated = await store.upsertIf({ - entry: sessionEntry(threadId, 5), - expectedRevision: null, + entry: entry('thread', 5), + revision: null, }) if (!recreated.ok) throw new Error('upsertIf after delete conflicted') expect(recreated.revision).not.toBe(last.revision) diff --git a/packages/ai-persistence/src/types.ts b/packages/ai-persistence/src/types.ts index 3dc5396d45..4f45266314 100644 --- a/packages/ai-persistence/src/types.ts +++ b/packages/ai-persistence/src/types.ts @@ -644,7 +644,7 @@ export interface BlobStore { // Opt-in stores for a host that keeps long-lived sessions: the inputs it // accepted, an append-only log, turn leases, work claims, // and a list of sessions. `withPersistence` writes only `sessions`, an entry -// for each subagent child. +// for each subagent child, in the tenant of its `tenantId` option. /** Lifecycle of one inbox entry. */ export type InboxStatus = 'pending' | 'applied' | 'rejected' @@ -1118,6 +1118,25 @@ export function defineWorkClaimStore(store: WorkClaimStore): WorkClaimStore { return store } +/** + * A user who can open a session, and the role of that user in it. The + * `owner` is the user who opened the session. The host decides what a + * `member` and a `viewer` can do. + */ +export interface SessionMember { + userId: string + role: 'owner' | 'member' | 'viewer' +} + +/** + * The key of one session index entry. Without `tenantId`, the key is in the + * default tenant. + */ +export interface SessionIndexKey { + tenantId?: string + threadId: string +} + /** * 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 @@ -1128,7 +1147,24 @@ export function defineWorkClaimStore(store: WorkClaimStore): WorkClaimStore { * sorts by it. */ export interface SessionIndexEntry> { + /** + * The tenant of the thread. Without it, the entry is in the default tenant. + * The tenant and `threadId` together are the key, so one thread id in two + * tenants is two entries. + */ + tenantId?: string threadId: string + /** + * The user id of the owner. Without it, the thread has no owner: a host + * without access rules opened it without a user. + */ + ownerId?: string + /** + * The users who can open the thread. The owner is in it with the role + * `owner`. Empty when the thread has no owner. A subagent child keeps it + * empty: access to a child goes through its top-level thread. + */ + members: Array /** The name of the setup that runs the thread. */ harness?: string title?: string @@ -1142,8 +1178,6 @@ export interface SessionIndexEntry> { 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. */ @@ -1164,6 +1198,16 @@ export interface SessionIndexEntry> { /** Options for {@link SessionIndexStore.list}. */ export interface SessionIndexListOptions { + /** + * The tenant to list. Without it, the list has only the default tenant. A + * list never has the entries of another tenant. + */ + tenantId?: string + /** + * Only the entries whose `members` have this user id, with any role. An + * entry with an `ownerId` and no `members` (a subagent child) never matches. + */ + memberId?: string /** The most entries in one page. Without it, the page has all entries. */ limit?: number /** The `cursor` of the previous page. */ @@ -1173,11 +1217,6 @@ export interface SessionIndexListOptions { * 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. @@ -1204,60 +1243,65 @@ export interface SessionIndexPage { truncated?: boolean } -/** Durable index of sessions, one entry per thread. */ +/** + * Durable index of sessions, one entry per tenant and thread. The key is + * `(tenantId ?? '', threadId)`: a key without `tenantId` is in the default + * tenant, and one thread id in two tenants is two entries. + */ export interface SessionIndexStore { + /** The entry for `key`, or `undefined`. */ + get: (key: SessionIndexKey) => Promise /** - * 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. + * The entry for `key` with its revision, or `undefined` when it is absent. + * A write that changes the entry gives it a new revision. A revision may + * name the content (a content-hash ETag, as in S3 or Netlify Blobs), so the + * same entry written again, also after a `delete`, can get the same + * revision back. */ - upsert: (entry: SessionIndexEntry) => Promise - /** The entry for `threadId`, or `undefined`. */ - get: (threadId: string) => Promise + getVersioned: ( + key: SessionIndexKey, + ) => Promise<{ entry: SessionIndexEntry; revision: string } | undefined> /** - * 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. A write that changes the entry gives it a new revision. A - * revision may name the content (a content-hash ETag, as in S3 or Netlify - * Blobs), so the same entry written again, also after a `delete`, can get - * the same revision back. + * Write the entry for the tenant and thread of `entry`. + * + * INVARIANT (full replace): a second upsert for the same key replaces the + * whole entry. A field that the new entry does not have is gone. To change + * one field, read the entry and write a changed copy with `upsertIf`. */ - getVersioned?: ( - threadId: string, - ) => Promise<{ entry: SessionIndexEntry; revision: string } | undefined> + upsert: (entry: SessionIndexEntry) => Promise /** - * 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. + * Write `entry` only when the current revision of its key is `revision` + * (`null` means "only when absent"). Returns the new revision, or + * `{ ok: false, reason: 'conflict' }` when someone else wrote first. With + * `getVersioned`, hosts on several machines do not lose each other's + * changes to one entry. A full replace, the same as `upsert`. */ - upsertIf?: (write: { + upsertIf: (input: { entry: SessionIndexEntry - expectedRevision: string | null + revision: 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. + * The entries of one tenant 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. + * is `truncated: true` with a `cursor`. Pass that `cursor` back, with the + * same filters, 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. + * Remove the entry for `key`. A hard delete: after it, `get` and + * `getVersioned` return `undefined`, and `upsertIf` with revision `null` + * writes a new entry. A no-op for an unknown key. * * 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 + delete: (key: SessionIndexKey) => Promise } /** Type a {@link SessionIndexStore} implementation inline. */ diff --git a/packages/ai-persistence/tests/persistence-types.test-d.ts b/packages/ai-persistence/tests/persistence-types.test-d.ts index e54179fd86..b1da72da92 100644 --- a/packages/ai-persistence/tests/persistence-types.test-d.ts +++ b/packages/ai-persistence/tests/persistence-types.test-d.ts @@ -151,6 +151,16 @@ expectTypeOf(memoryPersistence()).toEqualTypeOf< sessions: SessionIndexStore }> >() +// The session index is keyed by tenant and thread, and compare-and-set is part +// of the contract: a store without it is not a session index store. +expectTypeOf[0]>().toEqualTypeOf<{ + tenantId?: string + threadId: string +}>() +expectTypeOf< + Omit +>().not.toExtend() + const transcript: ChatTranscriptPersistence = messagesOnly void transcript declare const fullChat: ChatPersistence diff --git a/packages/ai-persistence/tests/subagent-persistence-cards.test.ts b/packages/ai-persistence/tests/subagent-persistence-cards.test.ts index c00f64814d..ba65f201b0 100644 --- a/packages/ai-persistence/tests/subagent-persistence-cards.test.ts +++ b/packages/ai-persistence/tests/subagent-persistence-cards.test.ts @@ -11,18 +11,26 @@ import type { } from '@tanstack/ai' import { memoryPersistence, reconstructChat, withPersistence } from '../src' import { createSubagentRunRecorder } from '../src/subagent-runs' +import type { SessionIndexStore } from '../src' const t = 1 /** A model that plays one script per call. */ -function scriptedAdapter(scripts: Array>): AnyTextAdapter { +function scriptedAdapter(scripts: Array>) { let call = 0 return { kind: 'text', name: 'scripted', model: 'test-model', - '~types': {}, - chatStream() { + '~types': { + providerOptions: undefined, + inputModalities: undefined, + messageMetadataByModality: undefined, + toolCapabilities: undefined, + toolCallMetadata: undefined, + systemPromptMetadata: undefined, + }, + chatStream: () => { const script = scripts[call++] if (!script) throw new Error('no script for this model call') return (async function* () { @@ -30,7 +38,8 @@ function scriptedAdapter(scripts: Array>): AnyTextAdapter { yield* script })() }, - } as unknown as AnyTextAdapter + structuredOutput: async () => ({ data: {}, rawText: '{}' }), + } satisfies AnyTextAdapter } /** A child that thinks, asks approval to delete a file, then reports. */ @@ -100,7 +109,7 @@ function cleaner( timestamp: t, }, ], - ] as Array>) + ]) return defineAgent({ name: 'cleaner', description: 'Deletes files', @@ -193,6 +202,32 @@ function childStarted( } satisfies StreamChunk } +/** + * `store`, where another host runs `write` between the read and the write of + * the first `times` compare-and-set writes. + */ +function racingSessions( + store: SessionIndexStore, + write: () => Promise, + times: number, +) { + let raced = 0 + return { + get: (key) => store.get(key), + getVersioned: (key) => store.getVersioned(key), + upsert: (entry) => store.upsert(entry), + upsertIf: async (input) => { + if (raced < times) { + raced += 1 + await write() + } + return store.upsertIf(input) + }, + list: (options) => store.list(options), + delete: (key) => store.delete(key), + } satisfies SessionIndexStore +} + /** Feed one parent run's chunks on thread `desk` to a recorder. */ async function feed( recorder: ReturnType, @@ -663,7 +698,9 @@ describe('subagent run recorder', () => { ...textMessage('m1', 'Inner notes', { subagentRunId: 'inner' }), ]) - expect(await stores.sessions.get('subagent:inner')).toMatchObject({ + expect( + await stores.sessions.get({ threadId: 'subagent:inner' }), + ).toMatchObject({ parentThreadId: 'subagent:outer', parentToolCallId: 'call_n', }) @@ -677,17 +714,20 @@ describe('subagent run recorder', () => { ).toBeUndefined() }) - it('gives each child the owner of its parent thread', async () => { + it('writes each child in the tenant of the recorder, with the owner of its parent thread', async () => { const { stores } = memoryPersistence() await stores.sessions.upsert({ + tenantId: 'org-1', threadId: 'desk', - principal: { id: 'alice' }, + ownerId: 'alice', + members: [{ userId: 'alice', role: 'owner' }], createdAt: t, updatedAt: t, }) const recorder = createSubagentRunRecorder({ messages: stores.messages, sessions: stores.sessions, + tenantId: 'org-1', intervalMs: 0, }) await feed(recorder, 'run-1', [ @@ -698,12 +738,89 @@ describe('subagent run recorder', () => { }), ]) - const owners = await Promise.all( - ['subagent:outer', 'subagent:inner'].map( - async (threadId) => (await stores.sessions.get(threadId))?.principal, + const children = await Promise.all( + ['subagent:outer', 'subagent:inner'].map((threadId) => + stores.sessions.get({ tenantId: 'org-1', threadId }), ), ) - expect(owners).toEqual([{ id: 'alice' }, { id: 'alice' }]) + // No members: access to a child goes through its top-level thread. + const child = { tenantId: 'org-1', ownerId: 'alice', members: [] } + expect(children).toEqual([ + expect.objectContaining(child), + expect.objectContaining(child), + ]) + // Nothing went to the default tenant. + expect(await stores.sessions.list()).toEqual({ entries: [] }) + }) + + it('keeps a change that another host makes to a child entry while it writes', async () => { + const { stores } = memoryPersistence() + const named = { + threadId: 'subagent:child', + members: [], + title: 'Named elsewhere', + createdAt: 5, + updatedAt: 5, + } + const recorder = createSubagentRunRecorder({ + messages: stores.messages, + sessions: racingSessions( + stores.sessions, + () => stores.sessions.upsert(named), + 1, + ), + intervalMs: 0, + }) + + await feed(recorder, 'run-1', [ + childStarted('child', { parentToolCallId: 'call_1' }), + ]) + + expect(await stores.sessions.get({ threadId: 'subagent:child' })).toEqual({ + ...named, + parentThreadId: 'desk', + parentToolCallId: 'call_1', + updatedAt: expect.any(Number), + }) + }) + + it('warns and goes on when another host changes a child entry on every try', async () => { + const warn = vi.spyOn(console, 'warn').mockImplementation(() => {}) + try { + const { stores } = memoryPersistence() + let writes = 0 + const elsewhere = () => + stores.sessions.upsert({ + threadId: 'subagent:child', + members: [], + createdAt: 5, + updatedAt: (writes += 1), + }) + const recorder = createSubagentRunRecorder({ + messages: stores.messages, + sessions: racingSessions(stores.sessions, elsewhere, 5), + intervalMs: 0, + }) + + await feed(recorder, 'run-1', [ + childStarted('child', { parentToolCallId: 'call_1' }), + ]) + + expect(warn).toHaveBeenCalledWith( + expect.stringContaining('changed on another host 5 times'), + ) + // The recorder wrote nothing over the other host's entry. + expect(await stores.sessions.get({ threadId: 'subagent:child' })).toEqual( + { + threadId: 'subagent:child', + members: [], + createdAt: 5, + updatedAt: 5, + }, + ) + } finally { + warn.mockRestore() + } }) it('loads a waiting continued child without the messages of the call it waits in', async () => { @@ -882,9 +999,10 @@ describe('agent tools with persistence', () => { args: { prompt: 'Take notes' }, }) const threadId = `subagent:${startedId(first)}` - const entry = await sessions.get(threadId) + const entry = await sessions.get({ threadId }) expect(entry).toEqual({ threadId, + members: [], parentThreadId: 'desk', parentToolCallId: 'call_1', createdAt: expect.any(Number), @@ -899,8 +1017,9 @@ describe('agent tools with persistence', () => { toolCallId: 'call_2', args: { sessionId: startedId(first), prompt: 'More' }, }) - expect(await sessions.get(threadId)).toEqual({ + expect(await sessions.get({ threadId })).toEqual({ threadId, + members: [], title: 'Tide notes', parentThreadId: 'desk', parentToolCallId: 'call_2', @@ -909,26 +1028,33 @@ describe('agent tools with persistence', () => { }) }) - it('gives a child entry the harness and owner of its parent entry', async () => { + it('gives a child entry the tenant, harness, and owner of its parent entry', async () => { const persistence = memoryPersistence() const sessions = persistence.stores.sessions await sessions.upsert({ + tenantId: 'org-1', threadId: 'desk', harness: 'notes/desk', - principal: { id: 'alice' }, + ownerId: 'alice', + members: [{ userId: 'alice', role: 'owner' }], createdAt: 1, updatedAt: 1, }) const { agent } = noteTaker() - const chunks = await subagentTurn(withPersistence(persistence), agent, { + const middleware = withPersistence(persistence, { tenantId: 'org-1' }) + const chunks = await subagentTurn(middleware, 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' }) + const threadId = `subagent:${startedId(chunks)}` + expect(await sessions.get({ tenantId: 'org-1', threadId })).toMatchObject({ + tenantId: 'org-1', + harness: 'notes/desk', + ownerId: 'alice', + members: [], + }) }) it('continues a stored child by sessionId', async () => { diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index e185116708..bec478d622 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -2669,9 +2669,6 @@ importers: '@tanstack/ai-acp': specifier: workspace:^ version: link:../ai-acp - '@tanstack/ai-mcp': - specifier: workspace:^ - version: link:../ai-mcp devDependencies: '@tanstack/ai': specifier: workspace:* @@ -2679,6 +2676,9 @@ importers: '@tanstack/ai-harness': specifier: workspace:* version: link:../ai-harness + '@tanstack/ai-mcp': + specifier: workspace:* + version: link:../ai-mcp '@tanstack/ai-persistence': specifier: workspace:* version: link:../ai-persistence diff --git a/testing/e2e/src/routes/api.harness-first-party.ts b/testing/e2e/src/routes/api.harness-first-party.ts index e801c5ef7a..58f53b7093 100644 --- a/testing/e2e/src/routes/api.harness-first-party.ts +++ b/testing/e2e/src/routes/api.harness-first-party.ts @@ -199,10 +199,10 @@ async function sessionTitle( ) const turn = await session.send('[first-party-title] plan a trip to Rome') const deadline = Date.now() + 5_000 - let entry = await host.sessions.get('e2e-title') + let entry = await host.sessions.get({ threadId: 'e2e-title' }) while (!entry?.title && failures.length === 0 && Date.now() < deadline) { await new Promise((resolve) => setTimeout(resolve, 50)) - entry = await host.sessions.get('e2e-title') + entry = await host.sessions.get({ threadId: 'e2e-title' }) } return { text: turn.text, title: entry?.title, failures } } diff --git a/testing/e2e/src/routes/api.harness-protocol.$.ts b/testing/e2e/src/routes/api.harness-protocol.$.ts index 7b7e1dcb9f..23b5da7439 100644 --- a/testing/e2e/src/routes/api.harness-protocol.$.ts +++ b/testing/e2e/src/routes/api.harness-protocol.$.ts @@ -2,14 +2,13 @@ import { createFileRoute } from '@tanstack/react-router' import { defineAgent } from '@tanstack/ai' import { createOpenaiChat } from '@tanstack/ai-openai' import { - type FinishContext, - type ProjectOptions, retryTransientErrors, configOption, createHarnessHost, defineCommand, defineHarness, definePlugin, + membersOnly, } from '@tanstack/ai-harness' import { createHarnessHandler } from '@tanstack/ai-harness/server' import { mcp } from '@tanstack/ai-mcp/harness' @@ -24,6 +23,11 @@ import { } from '@/lib/harness-protocol-tools' import { createImageAdapter } from '@/lib/media-providers' import { createTextAdapter } from '@/lib/providers' +import type { + FinishContext, + HarnessRunContext, + ProjectOptions, +} from '@tanstack/ai-harness' /** * The harness session protocol behind `createHarnessHandler`. The main @@ -32,15 +36,39 @@ import { createTextAdapter } from '@/lib/providers' * `painter`, which makes one image with the OpenAI image adapter. The * harness keeps that image, and the uploads, in memory. */ +// The users of the specs, by bearer token. `e2e`, `bob`, and `carol` are in +// one tenant, so a spec can share a thread between them. `mallory` is in +// another tenant. +const users = new Map([ + ['Bearer e2e-token', { id: 'e2e', tenantId: 'e2e' }], + ['Bearer e2e-token-bob', { id: 'bob', tenantId: 'e2e' }], + ['Bearer e2e-token-carol', { id: 'carol', tenantId: 'e2e' }], + ['Bearer e2e-token-mallory', { id: 'mallory', tenantId: 'other' }], +]) +const identify = ({ request }: { request: Request }) => + users.get(request.headers.get('authorization') ?? '') ?? null +// Test only: each user may do each action on each thread of their tenant. +// The host checks the tenant before it calls `authorize`. +const tenantAccess = { identify, authorize: () => true } // One host for every test. Each test uses its own thread ids, and a signed // media URL comes without the test headers. -const host = createHarnessHost({ persistence: memoryPersistence() }) +const host = createHarnessHost({ + persistence: memoryPersistence(), + access: tenantAccess, +}) // A durable host: a session log and run leases. A request picks it with the -// `x-harness-durable: 1` header. +// `x-harness-durable: 1` header. Host access keeps the owner and the members +// of each thread in `stores.sessions`. +const { runs, sessions } = memoryPersistence().stores const durableHost = createHarnessHost({ - persistence: { - stores: { log: memoryLogStore(), runs: memoryPersistence().stores.runs }, - }, + persistence: { stores: { log: memoryLogStore(), runs, sessions } }, + access: tenantAccess, +}) +// Only the members of a thread may use it. A request picks this host with +// the `x-harness-access: members` header. +const membersHost = createHarnessHost({ + persistence: memoryPersistence(), + access: { identify, authorize: membersOnly() }, }) // On the durable host, an `app.signal` record that a client appends becomes // a user message. The host without a log does not read it. @@ -51,14 +79,6 @@ const signalProject: ProjectOptions = { : [...messages], } -// Two users, so a test can share one thread between them. -const authorize = (req: Request) => { - const token = req.headers.get('authorization') - if (token === 'Bearer e2e-token') return { id: 'e2e' } - if (token === 'Bearer e2e-token-bob') return { id: 'bob' } - return null -} - // The `drafter` and `slowBuild` of each test wait until the test sends the // `release` command, so the test can act while they wait. // ponytail: one gate per test id, kept for the life of the server. @@ -228,13 +248,13 @@ function harnessFor(request: Request) { }, plugins: () => [ ...servers, - // The sender of the running turn, as plugins see it. + // The user who started the run, as its tools see it. definePlugin({ name: 'e2e/whoami', - setup: (ctx) => ({ + setup: () => ({ tools: [ - whoami.server(async () => ({ - id: ctx.session.principal?.id ?? null, + whoami.server(async (_input, toolContext) => ({ + id: toolContext.context.run.principal?.id ?? null, })), ], }), @@ -287,20 +307,20 @@ function harnessFor(request: Request) { }), ], }) + const isMembersOnly = request.headers.get('x-harness-access') === 'members' + const isDurable = request.headers.get('x-harness-durable') === '1' return { harness, - host: request.headers.get('x-harness-durable') === '1' ? durableHost : host, + host: isMembersOnly ? membersHost : isDurable ? durableHost : host, } } function handlerFor(request: Request) { const { harness, host: hostOf } = harnessFor(request) + // The `access.identify` of the host reads the user of each request. return createHarnessHandler({ host: hostOf, harness, - authorize, - // Test only: every principal may use every thread. - canAccess: () => true, // Each request makes a new handler, so a fixed secret keeps a signed URL // working on the next request. media: { secret: 'e2e-media-secret' }, @@ -314,15 +334,12 @@ function handlerFor(request: Request) { * Here the test picks it, so a test can fork onto a thread that has messages. */ async function forkThread(request: Request) { - const principal = authorize(request) + const principal = identify({ request }) if (!principal) return Response.json({ error: 'unauthorized' }, { status: 401 }) const { harness, host: hostOf } = harnessFor(request) - const body = (await request.json()) as { - threadId: string - newThreadId: string - at?: string - } + const body: { threadId: string; newThreadId: string; at?: string } = + await request.json() try { const forked = await hostOf.fork(harness, { ...body, principal }) return Response.json({ transcript: await forked.transcript() }) diff --git a/testing/e2e/tests/harness-access.spec.ts b/testing/e2e/tests/harness-access.spec.ts new file mode 100644 index 0000000000..80fead9335 --- /dev/null +++ b/testing/e2e/tests/harness-access.spec.ts @@ -0,0 +1,172 @@ +import { test, expect } from './fixtures' +import type { APIRequestContext, APIResponse } from '@playwright/test' + +/** + * Access on a `membersOnly()` host: the owner shares a thread with one + * member. A user of the same tenant who is not a member is refused (403), + * and a user of another tenant does not find the thread (404). The + * `x-harness-access: members` header picks that host in the route. + */ + +const BASE = '/api/harness-protocol' +const ANSWER = 'Hello over the harness protocol.' +const FORBIDDEN = { status: 403, body: { error: 'forbidden' } } +const NOT_FOUND = { status: 404, body: { error: 'not found' } } + +/** The text of the answer in a POST run stream. */ +const answerOf = async (response: APIResponse) => + (await response.text()) + .split('\n') + .filter((line) => line.startsWith('data: ')) + .map((line) => JSON.parse(line.slice('data: '.length))) + .filter((event) => event.type === 'TEXT_MESSAGE_CONTENT') + .map((event) => event.delta) + .join('') + +/** The status and the JSON body of a refused request. */ +const refusalOf = async (response: APIResponse) => ({ + status: response.status(), + body: await response.json(), +}) + +/** + * Open the session stream of a thread, read its first frame, and close the + * stream. A refused request answers with JSON, not with a stream. + */ +async function firstEvent(url: string, headers: Record) { + const response = await fetch(url, { headers }) + if (!response.ok || !response.body) { + return { status: response.status, body: await response.json() } + } + const reader = response.body.pipeThrough(new TextDecoderStream()).getReader() + let buffer = '' + while (!buffer.includes('\n\n')) { + const read = await reader.read() + if (read.done) break + buffer += read.value + } + await reader.cancel() + const data = buffer.split('\n').find((line) => line.startsWith('data: ')) + return { + status: response.status, + body: data ? JSON.parse(data.slice('data: '.length)) : null, + } +} + +/** One user of the `membersOnly()` host, by the token of the route. */ +function asUser( + request: APIRequestContext, + { + testId, + aimockPort, + baseURL, + }: { testId: string; aimockPort: number; baseURL: string | undefined }, + token: string, +) { + const headers = { + authorization: `Bearer ${token}`, + 'x-test-id': testId, + 'x-aimock-port': String(aimockPort), + 'x-harness-access': 'members', + } + const json = { ...headers, 'content-type': 'application/json' } + return { + /** Send the hello prompt of the protocol fixture as the run `runId`. */ + send: (threadId: string, runId: string) => + request.post(`${BASE}/run`, { + headers: json, + data: { + threadId, + runId, + messages: [ + { id: 'u1', role: 'user', content: '[harness-protocol] hello' }, + ], + tools: [], + context: [], + state: {}, + forwardedProps: {}, + }, + }), + sessions: (body: object) => + request.post(`${BASE}/sessions`, { headers: json, data: body }), + events: (threadId: string) => + firstEvent( + `${baseURL}${BASE}/events?threadId=${threadId}&from=0`, + headers, + ), + /** What a `ChatClient` with `persistence` loads for the thread. */ + hydrate: (threadId: string) => + request.get(`${BASE}/run?threadId=${threadId}`, { headers }), + } +} + +test.describe('harness access', () => { + test('a shared thread serves its two members, refuses a third user, and stays hidden from another tenant', async ({ + request, + baseURL, + testId, + aimockPort, + }) => { + const user = (token: string) => + asUser(request, { testId, aimockPort, baseURL }, token) + const owner = user('e2e-token') + const bob = user('e2e-token-bob') + const carol = user('e2e-token-carol') + const mallory = user('e2e-token-mallory') + const threadId = `access-${testId}` + + // 1. The owner starts the thread. + expect( + await answerOf(await owner.send(threadId, `access-owner-${testId}`)), + ).toBe(ANSWER) + + // 2. The owner shares the thread with bob. + const shared = await owner.sessions({ + op: 'share', + threadId, + userId: 'bob', + role: 'member', + }) + expect((await shared.json()).members).toEqual([ + { userId: 'e2e', role: 'owner' }, + { userId: 'bob', role: 'member' }, + ]) + + // 3. Bob reads the session stream and sends a prompt. + expect(await bob.events(threadId)).toMatchObject({ + status: 200, + body: { type: 'harness.event' }, + }) + expect( + await answerOf(await bob.send(threadId, `access-bob-${testId}`)), + ).toBe(ANSWER) + + // 4. Carol is in the same tenant, but she is not a member. + expect(await carol.events(threadId)).toEqual(FORBIDDEN) + expect( + await refusalOf(await carol.send(threadId, `access-carol-${testId}`)), + ).toEqual(FORBIDDEN) + + // 5. Mallory is in another tenant, where this thread does not exist. + expect(await mallory.events(threadId)).toEqual(NOT_FOUND) + expect(await (await mallory.hydrate(threadId)).json()).toEqual({ + messages: [], + activeRun: null, + interrupts: null, + }) + + // 6. The owner removes bob, so the next requests of bob are refused. + const unshared = await owner.sessions({ + op: 'unshare', + threadId, + userId: 'bob', + }) + expect((await unshared.json()).members).toEqual([ + { userId: 'e2e', role: 'owner' }, + ]) + expect(await bob.events(threadId)).toEqual(FORBIDDEN) + expect( + await refusalOf(await bob.send(threadId, `access-bob-again-${testId}`)), + ).toEqual(FORBIDDEN) + }) +}) diff --git a/testing/e2e/tests/harness-media.spec.ts b/testing/e2e/tests/harness-media.spec.ts index cf7a46e955..7578579382 100644 --- a/testing/e2e/tests/harness-media.spec.ts +++ b/testing/e2e/tests/harness-media.spec.ts @@ -1,4 +1,5 @@ import { test, expect } from './fixtures' +import { openHarnessThread } from './helpers' import type { APIRequestContext } from '@playwright/test' import type { MediaRecord } from '@tanstack/ai-harness' @@ -184,6 +185,8 @@ test.describe('harness media', () => { return JSON.stringify(await response.json()) } expect(await childOf(threadId)).toContain('Painted one image.') - expect(await childOf(`media-other-${testId}`)).toBe('[]') + const other = `media-other-${testId}` + await openHarnessThread(request, api.headers, other) + expect(await childOf(other)).toBe('[]') }) }) diff --git a/testing/e2e/tests/harness-protocol-inputs.spec.ts b/testing/e2e/tests/harness-protocol-inputs.spec.ts index a468e8b3b9..234e7a4a17 100644 --- a/testing/e2e/tests/harness-protocol-inputs.spec.ts +++ b/testing/e2e/tests/harness-protocol-inputs.spec.ts @@ -1,6 +1,7 @@ import { ChatClient, fetchServerSentEvents } from '@tanstack/ai-client' import { deploy } from '../src/lib/harness-protocol-tools' import { test, expect } from './fixtures' +import { openHarnessThread } from './helpers' import type { APIRequestContext } from '@playwright/test' type Chunk = { @@ -364,11 +365,7 @@ test.describe('harness protocol inputs', () => { const opener = asUser(request, testId, aimockPort) const bob = asUser(request, testId, aimockPort, 'e2e-token-bob') // The first user opens the session of the thread. - const snapshot = await request.get( - `/api/harness-protocol/snapshot?threadId=${threadId}`, - { headers: opener.headers }, - ) - expect(snapshot.status()).toBe(200) + await openHarnessThread(request, opener.headers, threadId) const chunks = sse( await ( @@ -438,13 +435,14 @@ test.describe('harness protocol inputs', () => { .join(''), })) const pendingQuestion = async () => { + const response = await request.get( + `/api/harness-protocol/snapshot?threadId=${threadId}`, + { headers: user.headers }, + ) + // 404 until the turn of the first tab opens the thread. + if (response.status() === 404) return undefined const snapshot: { pendingQuestions: Array<{ questionId: string }> } = - await ( - await request.get( - `/api/harness-protocol/snapshot?threadId=${threadId}`, - { headers: user.headers }, - ) - ).json() + await response.json() return snapshot.pendingQuestions[0]?.questionId } diff --git a/testing/e2e/tests/harness-protocol.spec.ts b/testing/e2e/tests/harness-protocol.spec.ts index a815485660..6d29f3fe4c 100644 --- a/testing/e2e/tests/harness-protocol.spec.ts +++ b/testing/e2e/tests/harness-protocol.spec.ts @@ -1,4 +1,5 @@ import { test, expect } from './fixtures' +import { openHarnessThread } from './helpers' test.describe('harness protocol', () => { const headers = (testId: string, aimockPort: number) => ({ @@ -400,11 +401,7 @@ test.describe('harness protocol', () => { // The stream ends with the turn. await run.text() // The second thread only opens. - const opened = await request.get( - `/api/harness-protocol/snapshot?threadId=${second}`, - { headers: auth }, - ) - expect(opened.status()).toBe(200) + await openHarnessThread(request, auth, second) const sessions = (body: object) => request.post('/api/harness-protocol/sessions', { @@ -509,11 +506,7 @@ test.describe('harness protocol', () => { [`search-c-${testId}`]: `Trip budget ${testId}`, } for (const [threadId, title] of Object.entries(titles)) { - const opened = await request.get( - `/api/harness-protocol/snapshot?threadId=${threadId}`, - { headers: auth }, - ) - expect(opened.status()).toBe(200) + await openHarnessThread(request, auth, threadId) const renamed = await request.post('/api/harness-protocol/sessions', { headers: json, data: { op: 'rename', threadId, title }, @@ -554,9 +547,7 @@ test.describe('harness protocol', () => { const auth = headers(testId, aimockPort) const threadId = `status-${testId}` // Open the session first, so the feed starts with its status. - await request.get(`/api/harness-protocol/snapshot?threadId=${threadId}`, { - headers: auth, - }) + await openHarnessThread(request, auth, threadId) const feed = await fetch(`${baseURL}/api/harness-protocol/host-events`, { headers: auth, }) diff --git a/testing/e2e/tests/helpers.ts b/testing/e2e/tests/helpers.ts index 3210214a8c..9b569912d2 100644 --- a/testing/e2e/tests/helpers.ts +++ b/testing/e2e/tests/helpers.ts @@ -1,4 +1,5 @@ -import { expect, type Page } from '@playwright/test' +import { expect } from '@playwright/test' +import type { APIRequestContext, Page } from '@playwright/test' declare global { interface Window { @@ -290,3 +291,20 @@ export async function waitForGenerationComplete(page: Page, timeout = 30_000) { timeout, }) } + +/** + * Open a thread of the `/api/harness-protocol` route and start no turn. The + * host has access, so a read of a thread with no index entry is 404 and + * opens nothing. A setting change opens the thread. + */ +export async function openHarnessThread( + request: APIRequestContext, + headers: Record, + threadId: string, +) { + const response = await request.post('/api/harness-protocol/control', { + headers: { ...headers, 'content-type': 'application/json' }, + data: { threadId, input: { op: 'config', key: 'tone', value: 'plain' } }, + }) + expect(await response.json()).toMatchObject({ status: 'accepted' }) +}