From 689da08157006727392b02ca8ac873e501a4c3ea Mon Sep 17 00:00:00 2001 From: Aditya Kumarakrishnan Date: Mon, 31 Aug 2026 13:53:37 +0530 Subject: [PATCH 1/2] feat(agents): adopt backend-issued claim write tokens and opt-in fenced session streams --- .../agents-backend-issued-write-tokens.md | 6 + packages/agents-runtime/src/process-wake.ts | 5 + .../agents-runtime/test/process-wake.test.ts | 53 +++ .../src/claim-write-token-store.ts | 54 +++- packages/agents-server/src/entity-manager.ts | 23 +- .../src/routing/internal-router.ts | 79 ++++- .../src/routing/runners-router.ts | 10 + .../src/routing/stream-append.ts | 11 + packages/agents-server/src/runtime.ts | 2 + packages/agents-server/src/server.ts | 9 + .../agents-server/src/standalone-runtime.ts | 2 + packages/agents-server/src/stream-client.ts | 34 +- .../test/claim-write-token-store.test.ts | 57 ++++ .../test/electric-agents-status.test.ts | 1 + .../agents-server/test/runners-router.test.ts | 48 +++ .../test/server-claim-write-token.test.ts | 73 +++++ .../agents-server/test/stream-append.test.ts | 109 +++++++ .../test/stream-client-fork.test.ts | 16 + .../agents-server/test/stream-client.test.ts | 43 +++ .../subscription-webhooks-routing.test.ts | 306 ++++++++++++++++++ 20 files changed, 928 insertions(+), 13 deletions(-) create mode 100644 .changeset/agents-backend-issued-write-tokens.md create mode 100644 packages/agents-server/test/stream-append.test.ts diff --git a/.changeset/agents-backend-issued-write-tokens.md b/.changeset/agents-backend-issued-write-tokens.md new file mode 100644 index 0000000000..32ff4c49aa --- /dev/null +++ b/.changeset/agents-backend-issued-write-tokens.md @@ -0,0 +1,6 @@ +--- +'@electric-ax/agents-runtime': patch +'@electric-ax/agents-server': patch +--- + +Adopt backend-issued claim write tokens: when the Durable Streams backend implements the Write Fencing extension, the Agents Server adopts the `write_token` it delivers with wake notifications, pull claims, and heartbeat acks as the claim's write token, and the runtime refreshes its token from heartbeat responses. An opt-in `fencedSessionStreams` server option (env `ELECTRIC_AGENTS_FENCED_SESSION_STREAMS`) creates entity session streams with `Write-Fence: true` and forwards the token plus the fenced-class assertion on runtime appends, so the backend itself rejects deposed or lapsed writers. Off by default, and behaviour is unchanged when the backend supplies no token. diff --git a/packages/agents-runtime/src/process-wake.ts b/packages/agents-runtime/src/process-wake.ts index 472cba39b5..cf8cf6992d 100644 --- a/packages/agents-runtime/src/process-wake.ts +++ b/packages/agents-runtime/src/process-wake.ts @@ -1248,6 +1248,7 @@ export async function processWake( ok: boolean claimToken?: string token?: string + writeToken?: string } if (!data.ok) { failBackgroundWake( @@ -1258,6 +1259,10 @@ export async function processWake( } if (data.claimToken) activeClaimToken = data.claimToken if (data.token) activeClaimToken = data.token + // The server may re-mint the write token on each heartbeat (its + // backend's token TTL tracks the claim lease); adopt the refresh + // so appends from a long-running activation keep a live token. + if (data.writeToken) writeToken = data.writeToken }) .catch((err: unknown) => { failBackgroundWake(err, `HEARTBEAT_FAILED`) diff --git a/packages/agents-runtime/test/process-wake.test.ts b/packages/agents-runtime/test/process-wake.test.ts index 49c3a5460f..0c09cbc8a2 100644 --- a/packages/agents-runtime/test/process-wake.test.ts +++ b/packages/agents-runtime/test/process-wake.test.ts @@ -1847,6 +1847,59 @@ describe(`processWake`, () => { setIntervalSpy.mockRestore() }) + it(`adopts a refreshed write token from heartbeat responses`, async () => { + defineEntity(`test-agent`, { + handler: async () => { + // Keep the wake alive long enough for several heartbeat intervals. + await new Promise((resolve) => setTimeout(resolve, 150)) + }, + }) + + fetchMock.mockImplementation(async (url, opts) => { + if (String(url).includes(`/_electric/wakes/wake-abc`)) { + const body = JSON.parse(String(opts?.body ?? `{}`)) as Record< + string, + unknown + > + const isClaim = body.wakeId !== undefined + const isHeartbeat = !isClaim && body.done === undefined + return new Response( + JSON.stringify({ + ok: true, + ...(isClaim ? { writeToken: `wt-initial` } : {}), + ...(isHeartbeat ? { writeToken: `wt-refreshed` } : {}), + }), + { status: 200, headers: { 'content-type': `application/json` } } + ) + } + return new Response(JSON.stringify({ ok: true }), { + status: 200, + headers: { 'content-type': `application/json` }, + }) + }) + + await processWake(makeNotification(), { + ...BASE_CONFIG, + heartbeatInterval: 20, + }) + + // The producer reads the write token per request, so appends issued + // after a heartbeat refresh must carry the refreshed token. + const producer = mockConstructedProducers.find( + (constructed) => + constructed.producerId === + `entity-http://localhost:3000/test-agent/agent-1` + ) + const producerFetch = producer!.opts!.fetch as typeof fetch + await producerFetch(`http://localhost:3000/streams/entity:agent-1`, { + method: `POST`, + }) + const [, init] = fetchMock.mock.calls.at(-1)! + expect(new Headers(init?.headers).get(`authorization`)).toBe( + `Bearer wt-refreshed` + ) + }) + it(`flushes producer on completion`, async () => { defineEntity(`test-agent`, { handler: () => {}, diff --git a/packages/agents-server/src/claim-write-token-store.ts b/packages/agents-server/src/claim-write-token-store.ts index 1975e78b46..1e350fccea 100644 --- a/packages/agents-server/src/claim-write-token-store.ts +++ b/packages/agents-server/src/claim-write-token-store.ts @@ -2,14 +2,27 @@ import { randomUUID } from 'node:crypto' interface ActiveClaimWriteToken { token: string + /** + * The token this consumer's previous mint for the same stream issued. + * Kept valid so an append already in flight with the pre-refresh token + * is not rejected when a heartbeat re-mints. A different consumer's mint + * evicts the claim outright, previous token included. + */ + previousToken?: string consumerId: string } export class ClaimWriteTokenStore { private readonly claimsByStream = new Map() private readonly streamKeysByConsumer = new Map>() + private readonly deliveredTokensByConsumer = new Map() - mint(service: string, streamPath: string, consumerId: string): string { + mint( + service: string, + streamPath: string, + consumerId: string, + token: string = randomUUID() + ): string { const streamKey = this.streamKey(service, streamPath) const consumerKey = this.consumerKey(service, consumerId) const previousClaimForStream = this.claimsByStream.get(streamKey) @@ -20,16 +33,46 @@ export class ClaimWriteTokenStore { ) } - const token = randomUUID() - this.claimsByStream.set(streamKey, { token, consumerId }) + this.claimsByStream.set(streamKey, { + token, + consumerId, + ...(previousClaimForStream?.consumerId === consumerId + ? { previousToken: previousClaimForStream.token } + : {}), + }) this.addConsumerStream(consumerKey, streamKey) return token } + /** + * Remembers a write token the Durable Streams backend issued with a wake + * delivery (webhook notification or pull-wake claim), keyed by the wake's + * consumer id, until that consumer's claim callback mints it as the active + * claim write token. + */ + recordDelivered(service: string, consumerId: string, token: string): void { + this.deliveredTokensByConsumer.set( + this.consumerKey(service, consumerId), + token + ) + } + + takeDelivered(service: string, consumerId: string): string | undefined { + const consumerKey = this.consumerKey(service, consumerId) + const token = this.deliveredTokensByConsumer.get(consumerKey) + if (token !== undefined) { + this.deliveredTokensByConsumer.delete(consumerKey) + } + return token + } + isValid(service: string, streamPath: string, token: string): boolean { + const activeClaim = this.claimsByStream.get( + this.streamKey(service, streamPath) + ) return ( - this.claimsByStream.get(this.streamKey(service, streamPath))?.token === - token + activeClaim !== undefined && + (activeClaim.token === token || activeClaim.previousToken === token) ) } @@ -54,6 +97,7 @@ export class ClaimWriteTokenStore { clearConsumer(service: string, consumerId: string): void { const consumerKey = this.consumerKey(service, consumerId) + this.deliveredTokensByConsumer.delete(consumerKey) const streamKeys = this.streamKeysByConsumer.get(consumerKey) if (!streamKeys) return diff --git a/packages/agents-server/src/entity-manager.ts b/packages/agents-server/src/entity-manager.ts index 200f518661..4a764373b1 100644 --- a/packages/agents-server/src/entity-manager.ts +++ b/packages/agents-server/src/entity-manager.ts @@ -386,6 +386,15 @@ export class EntityManager { SpawnPersistResult > private readonly stopWakeRegistryOnShutdown: boolean + /** + * When enabled, entity session streams are created (and forked) with + * `Write-Fence: true` and runtime appends forward the claim write token + * plus the fenced-class assertion to the Durable Streams backend, so a + * backend implementing the Write Fencing extension can reject deposed or + * lapsed writers itself. Off by default; a backend without the extension + * ignores the headers either way. + */ + readonly fencedSessionStreams: boolean constructor(opts: { registry: PostgresRegistry @@ -397,6 +406,7 @@ export class EntityManager { writeTokenValidator?: WriteTokenValidator spawnConcurrency?: number stopWakeRegistryOnShutdown?: boolean + fencedSessionStreams?: boolean }) { this.registry = opts.registry this.tenantId = opts.registry.tenantId ?? DEFAULT_TENANT_ID @@ -407,6 +417,9 @@ export class EntityManager { this.entityBridgeManager = opts.entityBridgeManager ?? null this.writeTokenValidator = opts.writeTokenValidator ?? null this.stopWakeRegistryOnShutdown = opts.stopWakeRegistryOnShutdown ?? true + this.fencedSessionStreams = + opts.fencedSessionStreams ?? + process.env.ELECTRIC_AGENTS_FENCED_SESSION_STREAMS === `true` const spawnConcurrency = opts.spawnConcurrency ?? @@ -853,6 +866,7 @@ export class EntityManager { this.streamClient.create(mainPath, { contentType, body: initialBody, + writeFence: this.fencedSessionStreams, }), ]) @@ -1194,9 +1208,12 @@ export class EntityManager { await this.streamClient.fork( plan.fork.streams.main, plan.source.streams.main, - isRoot && effectiveForkPointer - ? { forkPointer: effectiveForkPointer } - : undefined + { + ...(isRoot && effectiveForkPointer + ? { forkPointer: effectiveForkPointer } + : {}), + writeFence: this.fencedSessionStreams, + } ) createdStreams.push(plan.fork.streams.main) } diff --git a/packages/agents-server/src/routing/internal-router.ts b/packages/agents-server/src/routing/internal-router.ts index 87b78ee00d..d03403002d 100644 --- a/packages/agents-server/src/routing/internal-router.ts +++ b/packages/agents-server/src/routing/internal-router.ts @@ -84,6 +84,7 @@ const subscriptionWebhookBodySchema = Type.Object( streams: Type.Optional(Type.Array(Type.Record(Type.String(), Type.Any()))), callback_url: Type.Optional(Type.String()), callback_token: Type.Optional(Type.String()), + write_token: Type.Optional(Type.String()), primary_stream: Type.Optional(Type.String()), primaryStream: Type.Optional(Type.String()), streamPath: Type.Optional(Type.String()), @@ -315,6 +316,7 @@ function newWebhookPayload(body: SubscriptionWebhookBody | undefined): { tailOffset: string callbackUrl: string callbackToken: string + writeToken?: string } | null { if ( !body || @@ -358,6 +360,7 @@ function newWebhookPayload(body: SubscriptionWebhookBody | undefined): { tailOffset: selectedStream.tail_offset, callbackUrl: body.callback_url, callbackToken: body.callback_token, + ...(body.write_token ? { writeToken: body.write_token } : {}), } } @@ -585,6 +588,22 @@ async function subscriptionWebhook( ) } + if (newWebhook?.writeToken) { + // The backend minted this wake's write token before delivery (Write + // Fencing extension); hold it for the runtime's claim callback to + // adopt. Recorded only after the stopped/paused auto-ack and + // fork-lock rejections above, so a wake that never reaches the + // runtime — and therefore never claims — leaves no store entry + // behind. It must be recorded before the forward below, because the + // runtime's claim callback can arrive while the forward is still in + // flight. + ctx.runtime.claimWriteTokens.recordDelivered( + ctx.service, + newWebhook.wakeId, + newWebhook.writeToken + ) + } + if (entity) { rootSpan?.setAttribute(ATTR.ENTITY_URL, entity.url) await tracer.startActiveSpan( @@ -789,6 +808,29 @@ async function wakeCallback( responseBytes = new TextEncoder().encode(JSON.stringify(responseBody)) } } + } else if (!isDoneRequest && upstream.ok && target.primaryStream) { + // Heartbeat: a backend implementing the Write Fencing extension re-mints + // the claim's write token on every ack. Adopt the refresh so the active + // token tracks the backend's, and surface it to the runtime as + // `writeToken` alongside the passed-through ack body. + const responseBody = decodeJsonObject(responseBytes) + const refreshedToken = responseBody?.write_token + if ( + responseBody?.ok === true && + typeof refreshedToken === `string` && + refreshedToken !== `` + ) { + const writeToken = await mintClaimWriteToken( + ctx, + target.primaryStream, + consumerId, + refreshedToken + ) + if (writeToken) { + responseBody.writeToken = writeToken + responseBytes = new TextEncoder().encode(JSON.stringify(responseBody)) + } + } } try { @@ -904,12 +946,45 @@ async function wakeCallback( async function mintClaimWriteToken( ctx: TenantContext, streamPath: string, - consumerId: string + consumerId: string, + backendToken?: string ): Promise { const entity = await ctx.entityManager.registry.getEntityByStream(streamPath) if (!entity) return undefined - return ctx.runtime.claimWriteTokens.mint(ctx.service, streamPath, consumerId) + // When the Durable Streams backend issued a write token for this claim + // (Write Fencing extension), adopt it as the claim's write token so the + // store stays the single validation authority while the backend is the + // mint. When the backend supplied none, the store mints its own token and + // behaviour is byte-for-byte what it is today. + // + // Version-skew matrix — adoption is data-driven (it follows `write_token` + // fields wherever the backend sends them), while stream fencing is a + // separate opt-in (`fencedSessionStreams`): + // - Base backend (no Write Fencing), any server/runtime: no `write_token` + // ever appears, so the store mints and nothing changes. With + // `fencedSessionStreams` on, the fencing headers are sent but a base + // backend ignores them (additive headers, base spec §11), so + // enforcement remains this store only. + // - Token-minting backend, older server (no adoption): the optional + // fields are ignored and the old server mints its own tokens; it never + // creates streams fenced, so the backend enforces nothing. + // - Token-minting backend, this server, older runtime (one that ignores + // the heartbeat response's `writeToken`): adoption still happens here, + // so a backend that rotates the token on ack refresh ages the runtime's + // original token out of the store's one-refresh grace, and the + // runtime's appends start failing 401 mid-activation. Upgrade runtimes + // before pointing this server at a token-minting backend (and before + // enabling `fencedSessionStreams`). + const token = + backendToken ?? + ctx.runtime.claimWriteTokens.takeDelivered(ctx.service, consumerId) + return ctx.runtime.claimWriteTokens.mint( + ctx.service, + streamPath, + consumerId, + token + ) } function encodeWakeCallbackBody( diff --git a/packages/agents-server/src/routing/runners-router.ts b/packages/agents-server/src/routing/runners-router.ts index 8ffcaf45b0..cab02950e7 100644 --- a/packages/agents-server/src/routing/runners-router.ts +++ b/packages/agents-server/src/routing/runners-router.ts @@ -612,6 +612,16 @@ async function notificationFromClaim( }, }) + if (input.claim.write_token) { + // The backend minted a write token with this claim (Write Fencing + // extension); hold it for the runner's claim callback to adopt. + ctx.runtime.claimWriteTokens.recordDelivered( + ctx.service, + input.claim.wake_id, + input.claim.write_token + ) + } + await ctx.entityManager.registry.materializeActiveClaim({ consumerId: input.claim.wake_id, epoch: input.claim.generation, diff --git a/packages/agents-server/src/routing/stream-append.ts b/packages/agents-server/src/routing/stream-append.ts index e6f6d31784..86e794749d 100644 --- a/packages/agents-server/src/routing/stream-append.ts +++ b/packages/agents-server/src/routing/stream-append.ts @@ -10,6 +10,7 @@ import { ErrCodeUnauthorized, } from '../electric-agents-types.js' import { serverLog } from '../utils/log.js' +import { WRITE_FENCE_HEADER, WRITE_TOKEN_HEADER } from '../stream-client.js' import type { EntityManager } from '../entity-manager.js' import type { IRequest, RouterType } from 'itty-router' @@ -106,6 +107,16 @@ async function handleStreamAppend( if (!manager.isValidWriteToken(entity, token)) { return apiError(401, ErrCodeUnauthorized, `Invalid write token`) } + if (manager.fencedSessionStreams) { + // Forward the runtime's claim capability to the Durable Streams + // backend and assert the fenced write class, so a stale or lost token + // is a loud 401 downstream instead of a silent open-class write under + // this server's forwarded identity. The runtime's own bearer is + // overwritten with the server's when the request is forwarded, so + // `Write-Token` is the only carrier that survives. + request.headers.set(WRITE_TOKEN_HEADER, token) + request.headers.set(WRITE_FENCE_HEADER, `true`) + } if (manager.isForkWriteLockedEntity(entity.url)) { return apiError( 409, diff --git a/packages/agents-server/src/runtime.ts b/packages/agents-server/src/runtime.ts index c545fa714c..164ac965e6 100644 --- a/packages/agents-server/src/runtime.ts +++ b/packages/agents-server/src/runtime.ts @@ -49,6 +49,7 @@ export interface ElectricAgentsTenantRuntimeOptions { pgSync?: PgSyncBridgeManagerOptions claimWriteTokens?: ClaimWriteTokenStore stopWakeRegistryOnShutdown?: boolean + fencedSessionStreams?: boolean } export class ElectricAgentsTenantRuntime { @@ -99,6 +100,7 @@ export class ElectricAgentsTenantRuntime { token ), stopWakeRegistryOnShutdown: options.stopWakeRegistryOnShutdown ?? false, + fencedSessionStreams: options.fencedSessionStreams, }) this.pgSyncBridgeManager = options.pgSyncBridgeManager ?? diff --git a/packages/agents-server/src/server.ts b/packages/agents-server/src/server.ts index 8e265dca4d..249d794e43 100644 --- a/packages/agents-server/src/server.ts +++ b/packages/agents-server/src/server.ts @@ -84,6 +84,14 @@ export interface ElectricAgentsServerOptions { * Defaults to dispatchRecoveryIntervalMs when periodic recovery is enabled. */ staleOutstandingWakeAfterMs?: number + /** + * Create entity session streams with `Write-Fence: true` and forward the + * claim write token on runtime appends, letting a Durable Streams backend + * that implements the Write Fencing extension fence deposed writers + * itself. Off by default; falls back to the + * ELECTRIC_AGENTS_FENCED_SESSION_STREAMS env var. + */ + fencedSessionStreams?: boolean } interface MockAgentBootstrap { @@ -245,6 +253,7 @@ export class ElectricAgentsServer { electricUrl: this.options.electricUrl, electricSecret: this.options.electricSecret, pgSync: this.options.pgSync, + fencedSessionStreams: this.options.fencedSessionStreams, }) this.electricAgentsManager = this.standaloneRuntime.manager this.entityBridgeManager = this.standaloneRuntime.entityBridgeManager diff --git a/packages/agents-server/src/standalone-runtime.ts b/packages/agents-server/src/standalone-runtime.ts index bddf870b43..5b2e6073ba 100644 --- a/packages/agents-server/src/standalone-runtime.ts +++ b/packages/agents-server/src/standalone-runtime.ts @@ -39,6 +39,7 @@ export interface StandaloneAgentsRuntimeOptions { entityBridgeManager?: EntityBridgeCoordinator pgSyncBridgeManager?: PgSyncBridgeCoordinator pgSync?: PgSyncBridgeManagerOptions + fencedSessionStreams?: boolean } export interface StartedStandaloneAgentsRuntime { @@ -115,6 +116,7 @@ export async function startStandaloneAgentsRuntime( pgSyncBridgeManager: options.pgSyncBridgeManager, pgSync: options.pgSync, stopWakeRegistryOnShutdown: options.wakeRegistry ? false : true, + fencedSessionStreams: options.fencedSessionStreams, }) const startWakeRegistry = options.startWakeRegistry ?? true diff --git a/packages/agents-server/src/stream-client.ts b/packages/agents-server/src/stream-client.ts index 96e92de279..31b45a0bfe 100644 --- a/packages/agents-server/src/stream-client.ts +++ b/packages/agents-server/src/stream-client.ts @@ -11,6 +11,18 @@ import type { HeadersRecord, MaybePromise } from '@durable-streams/client' export type DurableStreamsBearerProvider = string | (() => MaybePromise) +/** + * Headers of the Durable Streams Write Fencing extension. `Write-Fence: true` + * on a create opts the stream into fenced appends; on an append it asserts + * the fenced write class, which the backend only honours together with a + * claim-scoped write token carried in `Write-Token`. A backend without the + * extension ignores both headers. + * + * https://github.com/adityavkk/chronicle/blob/main/docs/spec/WRITE-FENCING.md + */ +export const WRITE_FENCE_HEADER = `Write-Fence` +export const WRITE_TOKEN_HEADER = `Write-Token` + export interface StreamClientOptions { bearer?: DurableStreamsBearerProvider } @@ -72,6 +84,11 @@ export interface SubscriptionClaimResponse { wake_id: string generation: number token: string + /** + * Claim-scoped write token minted by a backend that implements the Write + * Fencing extension; absent on backends without it. + */ + write_token?: string streams: Array lease_ttl_ms?: number } @@ -226,7 +243,12 @@ export class StreamClient { async create( path: string, - opts: { contentType: string; body?: Uint8Array | string; closed?: boolean } + opts: { + contentType: string + body?: Uint8Array | string + closed?: boolean + writeFence?: boolean + } ): Promise { return await withSpan(`stream.create`, async (span) => { span.setAttributes({ @@ -235,7 +257,9 @@ export class StreamClient { }) await DurableStream.create({ url: this.streamUrl(path), - headers: this.streamHeaders(), + headers: opts.writeFence + ? { ...this.streamHeaders(), [WRITE_FENCE_HEADER]: `true` } + : this.streamHeaders(), contentType: opts.contentType, body: opts.body, closed: opts.closed, @@ -246,7 +270,7 @@ export class StreamClient { async fork( path: string, sourcePath: string, - opts?: { forkPointer?: EventPointer } + opts?: { forkPointer?: EventPointer; writeFence?: boolean } ): Promise { return await withSpan(`stream.fork`, async (span) => { span.setAttributes({ @@ -257,6 +281,10 @@ export class StreamClient { 'content-type': `application/json`, 'Stream-Forked-From': new URL(this.streamUrl(sourcePath)).pathname, } + if (opts?.writeFence) { + // Forks never inherit the fence; a fenced fork opts in explicitly. + headers[WRITE_FENCE_HEADER] = `true` + } if (opts?.forkPointer) { // The durable-streams server returns 400 if Stream-Fork-Sub-Offset // > 0 without an accompanying Stream-Fork-Offset. When our diff --git a/packages/agents-server/test/claim-write-token-store.test.ts b/packages/agents-server/test/claim-write-token-store.test.ts index f7ebed3a2a..d11cce5ec9 100644 --- a/packages/agents-server/test/claim-write-token-store.test.ts +++ b/packages/agents-server/test/claim-write-token-store.test.ts @@ -51,4 +51,61 @@ describe(`ClaimWriteTokenStore`, () => { expect(store.isValid(`tenant-a`, `/two/main`, secondToken)).toBe(false) expect(store.isValid(`tenant-a`, `/three/main`, otherToken)).toBe(true) }) + + it(`adopts a backend-issued token instead of minting one`, () => { + const store = new ClaimWriteTokenStore() + + const token = store.mint(`tenant-a`, `/one/main`, `wake-1`, `backend-token`) + + expect(token).toBe(`backend-token`) + expect(store.isValid(`tenant-a`, `/one/main`, `backend-token`)).toBe(true) + }) + + it(`keeps the previous token valid across a same-consumer refresh`, () => { + const store = new ClaimWriteTokenStore() + + const first = store.mint(`tenant-a`, `/one/main`, `wake-1`) + const second = store.mint(`tenant-a`, `/one/main`, `wake-1`) + + expect(store.isValid(`tenant-a`, `/one/main`, first)).toBe(true) + expect(store.isValid(`tenant-a`, `/one/main`, second)).toBe(true) + + const third = store.mint(`tenant-a`, `/one/main`, `wake-1`) + + expect(store.isValid(`tenant-a`, `/one/main`, first)).toBe(false) + expect(store.isValid(`tenant-a`, `/one/main`, second)).toBe(true) + expect(store.isValid(`tenant-a`, `/one/main`, third)).toBe(true) + }) + + it(`evicts the previous token when a different consumer claims the stream`, () => { + const store = new ClaimWriteTokenStore() + + const first = store.mint(`tenant-a`, `/one/main`, `wake-1`) + const refreshed = store.mint(`tenant-a`, `/one/main`, `wake-1`) + const successor = store.mint(`tenant-a`, `/one/main`, `wake-2`) + + expect(store.isValid(`tenant-a`, `/one/main`, first)).toBe(false) + expect(store.isValid(`tenant-a`, `/one/main`, refreshed)).toBe(false) + expect(store.isValid(`tenant-a`, `/one/main`, successor)).toBe(true) + }) + + it(`hands a delivered token to exactly one claim`, () => { + const store = new ClaimWriteTokenStore() + + store.recordDelivered(`tenant-a`, `wake-1`, `backend-token`) + + expect(store.takeDelivered(`tenant-a`, `wake-1`)).toBe(`backend-token`) + expect(store.takeDelivered(`tenant-a`, `wake-1`)).toBeUndefined() + }) + + it(`scopes delivered tokens by tenant and drops them with the consumer`, () => { + const store = new ClaimWriteTokenStore() + + store.recordDelivered(`tenant-a`, `wake-1`, `token-a`) + store.recordDelivered(`tenant-b`, `wake-1`, `token-b`) + store.clearConsumer(`tenant-a`, `wake-1`) + + expect(store.takeDelivered(`tenant-a`, `wake-1`)).toBeUndefined() + expect(store.takeDelivered(`tenant-b`, `wake-1`)).toBe(`token-b`) + }) }) diff --git a/packages/agents-server/test/electric-agents-status.test.ts b/packages/agents-server/test/electric-agents-status.test.ts index 277b83119b..794f56a8cb 100644 --- a/packages/agents-server/test/electric-agents-status.test.ts +++ b/packages/agents-server/test/electric-agents-status.test.ts @@ -811,6 +811,7 @@ describe(`ElectricAgentsManager.forkSubtree`, () => { expect(rootForkCall).toBeDefined() expect(rootForkCall?.opts).toEqual({ forkPointer: { offset: `C`, subOffset: 1 }, + writeFence: false, }) }) diff --git a/packages/agents-server/test/runners-router.test.ts b/packages/agents-server/test/runners-router.test.ts index c59d5ca8c5..b493e2a8af 100644 --- a/packages/agents-server/test/runners-router.test.ts +++ b/packages/agents-server/test/runners-router.test.ts @@ -1,4 +1,5 @@ import { describe, expect, it, vi } from 'vitest' +import { ClaimWriteTokenStore } from '../src/claim-write-token-store' import { globalRouter } from '../src/routing/global-router' import { DurableStreamsSubscriptionError } from '../src/stream-client' import type { TenantContext } from '../src/routing/context' @@ -521,6 +522,53 @@ describe(`runner routes`, () => { ) }) + it(`holds a backend-issued write token from a runner claim for the claim callback`, async () => { + const claimWriteTokens = new ClaimWriteTokenStore() + const ctx = buildContext({ + runtime: { claimWriteTokens } as any, + }) + vi.mocked(ctx.streamClient.claimSubscription).mockResolvedValue({ + wake_id: `wake-1`, + generation: 7, + token: `claim-token`, + write_token: `backend-token-1`, + streams: [{ path: `chat/one/main`, tail_offset: `12` }], + lease_ttl_ms: 30_000, + }) + vi.mocked(ctx.entityManager.registry.getEntityByStream).mockResolvedValue({ + url: `/chat/one`, + type: `chat`, + status: `idle`, + streams: { main: `/chat/one/main` }, + subscription_id: `runner:runner-1`, + write_token: `entity-token`, + tags: {}, + created_at: 1, + updated_at: 1, + }) + + const response = await globalRouter.fetch( + request(`POST`, `/_electric/runners/runner-1/claim`, { + subscription_id: `runner:runner-1`, + stream: `chat/one/main`, + generation: 7, + ts: 123, + }), + ctx + ) + + expect(response.status).toBe(200) + const body = (await response.json()) as Record + expect(body.claimToken).toBe(`claim-token`) + // The backend's write token is not part of the notification; the runner + // receives it from its claim callback, which adopts the held token. + expect(body.write_token).toBeUndefined() + expect(body.writeToken).toBeUndefined() + expect(claimWriteTokens.takeDelivered(`tenant-test`, `wake-1`)).toBe( + `backend-token-1` + ) + }) + it(`releases paused entity claims without dispatching pending work`, async () => { const ctx = buildContext() vi.mocked(ctx.streamClient.claimSubscription).mockResolvedValue({ diff --git a/packages/agents-server/test/server-claim-write-token.test.ts b/packages/agents-server/test/server-claim-write-token.test.ts index cab6275338..662522e589 100644 --- a/packages/agents-server/test/server-claim-write-token.test.ts +++ b/packages/agents-server/test/server-claim-write-token.test.ts @@ -399,6 +399,79 @@ describe(`Claim-scoped write tokens`, () => { expect(currentOwnerTokenRes.status).toBe(204) }, 20_000) + it(`heartbeat acks carrying a backend write token refresh the active claim token`, async () => { + const typeName = `claim-heartbeat-refresh-${Date.now()}` + const entity = await createEntity(typeName, `owner`) + const pgDb = (electricAgentsServer as any).pgDb + + // A callback receiver that answers heartbeats the way a Durable Streams + // backend with the Write Fencing extension does: ok plus a re-minted + // write token. + const refreshingReceiver = createServer((_req, res) => { + res.writeHead(200, { 'content-type': `application/json` }) + res.end(JSON.stringify({ ok: true, write_token: `backend-refresh-1` })) + }) + await new Promise((resolve) => + refreshingReceiver.listen(0, `127.0.0.1`, () => resolve()) + ) + const refreshingReceiverUrl = `http://127.0.0.1:${ + (refreshingReceiver.address() as { port: number }).port + }` + + try { + await pgDb.insert(consumerCallbacks).values({ + consumerId: `consumer-refresh`, + callbackUrl: refreshingReceiverUrl, + primaryStream: entity.streams.main, + }) + + const claim = await claimConsumer({ + consumerId: `consumer-refresh`, + epoch: 4, + wakeId: `wake-refresh`, + }) + expect(claim.ok).toBe(true) + expect(claim.writeToken).toBeTruthy() + + const heartbeatRes = await fetch( + `${baseUrl}/_electric/wake-callbacks/consumer-refresh`, + { + method: `POST`, + headers: { 'content-type': `application/json` }, + body: JSON.stringify({ epoch: 4, acks: [] }), + } + ) + expect(heartbeatRes.status).toBe(200) + const heartbeatBody = (await heartbeatRes.json()) as { + ok: boolean + writeToken?: string + } + expect(heartbeatBody.ok).toBe(true) + expect(heartbeatBody.writeToken).toBe(`backend-refresh-1`) + + // The refreshed token writes, and the pre-refresh token stays valid + // so an append already in flight with it is not rejected. + const refreshedTokenRes = await appendEntityEvent({ + streamPath: entity.streams.main, + writeToken: `backend-refresh-1`, + key: `manifest-refreshed-token`, + }) + expect(refreshedTokenRes.status).toBe(204) + + const previousTokenRes = await appendEntityEvent({ + streamPath: entity.streams.main, + writeToken: claim.writeToken!, + key: `manifest-previous-token`, + }) + expect(previousTokenRes.status).toBe(204) + } finally { + refreshingReceiver.closeAllConnections() + await new Promise((resolve) => + refreshingReceiver.close(() => resolve()) + ) + } + }, 20_000) + it(`kill clears the active claim token for the entity stream`, async () => { const typeName = `claim-kill-cleanup-${Date.now()}` const entity = await createEntity(typeName, `owner`) diff --git a/packages/agents-server/test/stream-append.test.ts b/packages/agents-server/test/stream-append.test.ts new file mode 100644 index 0000000000..841e148f74 --- /dev/null +++ b/packages/agents-server/test/stream-append.test.ts @@ -0,0 +1,109 @@ +import { describe, expect, it, vi } from 'vitest' +import { + createStreamAppendRouteRequest, + electricAgentsStreamAppendRouter, +} from '../src/routing/stream-append' +import type { ElectricAgentsStreamAppendRuntime } from '../src/routing/stream-append' + +const entity = { + url: `/horton/demo`, + type: `horton`, + status: `idle`, + streams: { main: `/horton/demo/main` }, + write_token: `claim-token-1`, +} + +function buildRuntime(opts: { + fencedSessionStreams: boolean + hasEntity?: boolean +}): ElectricAgentsStreamAppendRuntime { + return { + manager: { + registry: { + getEntityByStream: vi + .fn() + .mockResolvedValue(opts.hasEntity === false ? null : entity), + }, + fencedSessionStreams: opts.fencedSessionStreams, + isAttachmentStreamPath: vi.fn(() => false), + isValidWriteToken: vi.fn(() => true), + isForkWriteLockedEntity: vi.fn(() => false), + isForkWriteLockedStream: vi.fn(() => false), + validateWriteEvent: vi.fn().mockResolvedValue(null), + } as any, + evaluateWakePayload: vi.fn().mockResolvedValue(undefined), + checkRunFinished: vi.fn(), + syncManifestWakes: vi.fn().mockResolvedValue(undefined), + syncManifestEntitySources: vi.fn().mockResolvedValue(undefined), + syncManifestSchedules: vi.fn().mockResolvedValue(undefined), + } +} + +function appendRequest(path: string) { + return createStreamAppendRouteRequest( + new Request(`http://agents.local${path}`, { + method: `POST`, + headers: { + 'content-type': `application/json`, + authorization: `Bearer claim-token-1`, + }, + body: JSON.stringify({ type: `default`, key: `k1`, value: {} }), + }) + ) +} + +describe(`entity stream appends`, () => { + it(`forwards the write token and fenced-class assertion when fenced session streams are enabled`, async () => { + let forwardedHeaders: Headers | null = null + const forward = vi.fn(async (req: { headers: Headers }) => { + forwardedHeaders = req.headers + return new Response(null, { status: 204 }) + }) + + const response = await electricAgentsStreamAppendRouter.fetch( + appendRequest(`/horton/demo/main`), + buildRuntime({ fencedSessionStreams: true }), + forward as any + ) + + expect(response?.status).toBe(204) + expect(forwardedHeaders!.get(`Write-Token`)).toBe(`claim-token-1`) + expect(forwardedHeaders!.get(`Write-Fence`)).toBe(`true`) + }) + + it(`forwards appends without fencing headers by default`, async () => { + let forwardedHeaders: Headers | null = null + const forward = vi.fn(async (req: { headers: Headers }) => { + forwardedHeaders = req.headers + return new Response(null, { status: 204 }) + }) + + const response = await electricAgentsStreamAppendRouter.fetch( + appendRequest(`/horton/demo/main`), + buildRuntime({ fencedSessionStreams: false }), + forward as any + ) + + expect(response?.status).toBe(204) + expect(forwardedHeaders!.has(`Write-Token`)).toBe(false) + expect(forwardedHeaders!.has(`Write-Fence`)).toBe(false) + }) + + it(`never asserts the fenced class on shared-state appends`, async () => { + let forwardedHeaders: Headers | null = null + const forward = vi.fn(async (req: { headers: Headers }) => { + forwardedHeaders = req.headers + return new Response(null, { status: 204 }) + }) + + const response = await electricAgentsStreamAppendRouter.fetch( + appendRequest(`/_electric/shared-state/board-1`), + buildRuntime({ fencedSessionStreams: true, hasEntity: false }), + forward as any + ) + + expect(response?.status).toBe(204) + expect(forwardedHeaders!.has(`Write-Token`)).toBe(false) + expect(forwardedHeaders!.has(`Write-Fence`)).toBe(false) + }) +}) diff --git a/packages/agents-server/test/stream-client-fork.test.ts b/packages/agents-server/test/stream-client-fork.test.ts index 40a440def2..9540bbba42 100644 --- a/packages/agents-server/test/stream-client-fork.test.ts +++ b/packages/agents-server/test/stream-client-fork.test.ts @@ -35,6 +35,22 @@ describe(`StreamClient.fork`, () => { ).resolves.toEqual({ offset: expect.any(String) }) }) + it(`create and fork with writeFence succeed against a backend without the extension`, async () => { + // A base Durable Streams server ignores the Write-Fence header, so the + // fencedSessionStreams knob is safe to enable against any backend. + await client.create(`/source-fenced`, { + contentType: `application/json`, + body: `[]`, + writeFence: true, + }) + + await client.fork(`/fork-fenced`, `/source-fenced`, { writeFence: true }) + + await expect( + client.append(`/fork-fenced`, JSON.stringify({ type: `reconcile` })) + ).resolves.toEqual({ offset: expect.any(String) }) + }) + it(`preserves source history when reading the fork`, async () => { const sourceEvent = { type: `inbox`, diff --git a/packages/agents-server/test/stream-client.test.ts b/packages/agents-server/test/stream-client.test.ts index 6daa4aa173..c0579bbe71 100644 --- a/packages/agents-server/test/stream-client.test.ts +++ b/packages/agents-server/test/stream-client.test.ts @@ -1,4 +1,5 @@ import { beforeEach, describe, expect, it, vi } from 'vitest' +import { DurableStream } from '@durable-streams/client' import { StreamClient } from '../src/stream-client' @@ -222,6 +223,48 @@ describe(`StreamClient`, () => { } }) + it(`opts a stream into write fencing only when asked`, async () => { + const createMock = vi.mocked(DurableStream.create).mockClear() + createMock.mockResolvedValue(undefined as never) + const client = new StreamClient(`http://127.0.0.1:4545`) + + await client.create(`/horton/demo/main`, { + contentType: `application/json`, + writeFence: true, + }) + await client.create(`/horton/other/main`, { + contentType: `application/json`, + }) + + expect(createMock.mock.calls[0]![0]).toMatchObject({ + headers: { 'Write-Fence': `true` }, + }) + expect(createMock.mock.calls[1]![0]!.headers ?? {}).not.toHaveProperty( + `Write-Fence` + ) + }) + + it(`opts a fork into write fencing only when asked`, async () => { + const fetchMock = vi.spyOn(globalThis, `fetch`).mockResolvedValue( + new Response(null, { status: 200 }) + ) + const client = new StreamClient(`http://127.0.0.1:4545`) + + try { + await client.fork(`/fork/main`, `/source/main`, { writeFence: true }) + await client.fork(`/fork/other`, `/source/main`) + + expect( + new Headers(fetchMock.mock.calls[0]?.[1]?.headers).get(`Write-Fence`) + ).toBe(`true`) + expect( + new Headers(fetchMock.mock.calls[1]?.[1]?.headers).has(`Write-Fence`) + ).toBe(false) + } finally { + fetchMock.mockRestore() + } + }) + it(`preserves claim token authorization on subscription ack`, async () => { const fetchMock = vi.spyOn(globalThis, `fetch`).mockResolvedValueOnce( new Response(JSON.stringify({ ok: true }), { diff --git a/packages/agents-server/test/subscription-webhooks-routing.test.ts b/packages/agents-server/test/subscription-webhooks-routing.test.ts index 9c2bf2cef7..541264ae6d 100644 --- a/packages/agents-server/test/subscription-webhooks-routing.test.ts +++ b/packages/agents-server/test/subscription-webhooks-routing.test.ts @@ -633,6 +633,192 @@ describe(`subscription webhooks for Durable Streams subscriptions`, () => { } }) + it(`adopts the write token a backend delivers with a webhook wake`, async () => { + const claimWriteTokens = new ClaimWriteTokenStore() + const webhookSelect = selectDb([ + { webhookUrl: `http://runtime.local/_electric/builtin-agent-handler` }, + ]) + const insert = insertDb() + const fetchSpy = vi.spyOn(globalThis, `fetch`).mockResolvedValue( + new Response(JSON.stringify({ ok: true }), { + headers: { 'content-type': `application/json` }, + }) + ) + + try { + const webhookResponse = await globalRouter.fetch( + request(`POST`, `/_electric/subscription-webhooks/horton-handler`, { + subscription_id: `horton-handler`, + wake_id: `wake-fenced`, + generation: 9, + streams: [ + { + path: `horton/demo/main`, + acked_offset: `0`, + tail_offset: `1`, + has_pending: true, + }, + ], + callback_url: `http://durable.local/v1/stream/tenant-a/__ds/subscriptions/horton-handler/callback`, + callback_token: `callback-token`, + write_token: `backend-token-1`, + }), + buildContext({ + pgDb: { + select: webhookSelect.select, + insert: insert.insert, + } as any, + runtime: { claimWriteTokens } as any, + }) + ) + expect(webhookResponse.status).toBe(200) + + const callbackSelect = selectDb([ + { + callbackUrl: `http://durable.local/v1/stream/tenant-a/__ds/subscriptions/horton-handler/callback`, + primaryStream: `/horton/demo/main`, + }, + ]) + const claimResponse = await globalRouter.fetch( + request(`POST`, `/_electric/wake-callbacks/wake-fenced`, { + wakeId: `wake-fenced`, + epoch: 9, + }), + buildContext({ + pgDb: { select: callbackSelect.select } as any, + runtime: { claimWriteTokens } as any, + }) + ) + + expect(claimResponse.status).toBe(200) + const body = await responseJson(claimResponse) + expect(body.ok).toBe(true) + expect(body.writeToken).toBe(`backend-token-1`) + expect( + claimWriteTokens.isValid( + `tenant-a`, + `/horton/demo/main`, + `backend-token-1` + ) + ).toBe(true) + } finally { + fetchSpy.mockRestore() + } + }) + + it(`refreshes the claim write token from heartbeat ack responses`, async () => { + const select = selectDb([ + { + callbackUrl: `http://durable.local/v1/stream/tenant-a/__ds/subscriptions/horton-handler/callback`, + primaryStream: `/horton/demo/main`, + }, + ]) + const fetchSpy = vi.spyOn(globalThis, `fetch`).mockResolvedValue( + new Response(JSON.stringify({ ok: true, write_token: `ds-refresh-1` }), { + headers: { 'content-type': `application/json` }, + }) + ) + const ctx = buildContext({ + pgDb: { select: select.select } as any, + }) + const activeToken = ctx.runtime.claimWriteTokens.mint( + `tenant-a`, + `/horton/demo/main`, + `wake-1` + ) + + try { + const response = await globalRouter.fetch( + new Request(`http://agents.local/_electric/wake-callbacks/wake-1`, { + method: `POST`, + headers: { + 'content-type': `application/json`, + authorization: `Bearer callback-token`, + }, + body: JSON.stringify({ + epoch: 7, + acks: [{ path: `/horton/demo/main`, offset: `1` }], + }), + }), + ctx + ) + + expect(response.status).toBe(200) + const body = await responseJson(response) + expect(body.ok).toBe(true) + expect(body.writeToken).toBe(`ds-refresh-1`) + // The refreshed token is active; the pre-refresh token stays valid so + // an append already in flight with it is not rejected. + expect( + ctx.runtime.claimWriteTokens.isValid( + `tenant-a`, + `/horton/demo/main`, + `ds-refresh-1` + ) + ).toBe(true) + expect( + ctx.runtime.claimWriteTokens.isValid( + `tenant-a`, + `/horton/demo/main`, + activeToken + ) + ).toBe(true) + } finally { + fetchSpy.mockRestore() + } + }) + + it(`keeps heartbeat responses unchanged when the backend supplies no write token`, async () => { + const select = selectDb([ + { + callbackUrl: `http://durable.local/v1/stream/tenant-a/__ds/subscriptions/horton-handler/callback`, + primaryStream: `/horton/demo/main`, + }, + ]) + const fetchSpy = vi.spyOn(globalThis, `fetch`).mockResolvedValue( + new Response(JSON.stringify({ ok: true }), { + headers: { 'content-type': `application/json` }, + }) + ) + const ctx = buildContext({ + pgDb: { select: select.select } as any, + }) + const activeToken = ctx.runtime.claimWriteTokens.mint( + `tenant-a`, + `/horton/demo/main`, + `wake-1` + ) + + try { + const response = await globalRouter.fetch( + new Request(`http://agents.local/_electric/wake-callbacks/wake-1`, { + method: `POST`, + headers: { + 'content-type': `application/json`, + authorization: `Bearer callback-token`, + }, + body: JSON.stringify({ + epoch: 7, + acks: [{ path: `/horton/demo/main`, offset: `1` }], + }), + }), + ctx + ) + + expect(response.status).toBe(200) + await expect(responseJson(response)).resolves.toEqual({ ok: true }) + expect( + ctx.runtime.claimWriteTokens.isValid( + `tenant-a`, + `/horton/demo/main`, + activeToken + ) + ).toBe(true) + } finally { + fetchSpy.mockRestore() + } + }) + it(`auto-acks webhook wakes for stopped entities`, async () => { const select = selectDb([ { webhookUrl: `http://runtime.local/_electric/builtin-agent-handler` }, @@ -686,6 +872,126 @@ describe(`subscription webhooks for Durable Streams subscriptions`, () => { } }) + it(`holds no delivered write token for a wake auto-acked for a stopped entity`, async () => { + const claimWriteTokens = new ClaimWriteTokenStore() + const select = selectDb([ + { webhookUrl: `http://runtime.local/_electric/builtin-agent-handler` }, + ]) + const insert = insertDb() + const stoppedEntity = { + url: `/horton/demo`, + type: `horton`, + status: `stopped`, + streams: { + main: `/horton/demo/main`, + }, + } + const fetchSpy = vi.spyOn(globalThis, `fetch`) + + try { + const response = await globalRouter.fetch( + request(`POST`, `/_electric/subscription-webhooks/horton-handler`, { + subscription_id: `horton-handler`, + wake_id: `wake-stopped`, + generation: 8, + streams: [ + { + path: `horton/demo/main`, + acked_offset: `1`, + tail_offset: `2`, + has_pending: true, + }, + ], + callback_url: `http://durable.local/v1/stream/tenant-a/__ds/subscriptions/horton-handler/callback`, + callback_token: `callback-token`, + write_token: `backend-token-unclaimed`, + }), + buildContext({ + pgDb: { select: select.select, insert: insert.insert } as any, + runtime: { claimWriteTokens } as any, + entityManager: { + registry: { + getEntityByStream: vi.fn().mockResolvedValue(stoppedEntity), + updateStatus: vi.fn().mockResolvedValue(undefined), + }, + enrichPayload: vi.fn(async (payload) => payload), + isForkWorkLockedEntity: vi.fn(() => false), + } as any, + }) + ) + + expect(response.status).toBe(200) + await expect(response.json()).resolves.toEqual({ done: true }) + // The wake was auto-acked before any claim could happen, so the + // delivered token must not be left behind in the store. + expect( + claimWriteTokens.takeDelivered(`tenant-a`, `wake-stopped`) + ).toBeUndefined() + expect(fetchSpy).not.toHaveBeenCalled() + } finally { + fetchSpy.mockRestore() + } + }) + + it(`holds no delivered write token for a wake rejected while the entity is fork-locked`, async () => { + const claimWriteTokens = new ClaimWriteTokenStore() + const select = selectDb([ + { webhookUrl: `http://runtime.local/_electric/builtin-agent-handler` }, + ]) + const insert = insertDb() + const fetchSpy = vi.spyOn(globalThis, `fetch`) + + try { + const response = await globalRouter.fetch( + request(`POST`, `/_electric/subscription-webhooks/horton-handler`, { + subscription_id: `horton-handler`, + wake_id: `wake-locked`, + generation: 8, + streams: [ + { + path: `horton/demo/main`, + acked_offset: `1`, + tail_offset: `2`, + has_pending: true, + }, + ], + callback_url: `http://durable.local/v1/stream/tenant-a/__ds/subscriptions/horton-handler/callback`, + callback_token: `callback-token`, + write_token: `backend-token-unclaimed`, + }), + buildContext({ + pgDb: { select: select.select, insert: insert.insert } as any, + runtime: { claimWriteTokens } as any, + entityManager: { + registry: { + getEntityByStream: vi.fn().mockResolvedValue({ + url: `/horton/demo`, + type: `horton`, + status: `idle`, + streams: { + main: `/horton/demo/main`, + }, + }), + updateStatus: vi.fn().mockResolvedValue(undefined), + }, + enrichPayload: vi.fn(async (payload) => payload), + isForkWorkLockedEntity: vi.fn(() => true), + } as any, + }) + ) + + expect(response.status).toBe(409) + // The wake was rejected before any claim could happen, so the + // delivered token must not be left behind in the store. + expect( + claimWriteTokens.takeDelivered(`tenant-a`, `wake-locked`) + ).toBeUndefined() + expect(fetchSpy).not.toHaveBeenCalled() + } finally { + fetchSpy.mockRestore() + } + }) + it(`translates runtime done callbacks to the new Durable Streams callback shape`, async () => { const select = selectDb([ { From 9e9f593d6d6212a5a211098f98097fcb938f0185 Mon Sep 17 00:00:00 2001 From: Aditya Kumarakrishnan Date: Wed, 2 Sep 2026 05:15:20 +0530 Subject: [PATCH 2/2] feat(agents-server): the backend is the write authority for fenced session streams MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit With fencedSessionStreams on, the append path forwarded the claim's write token to the backend but still required the in-memory ClaimWriteTokenStore to know it first. That store is process memory: a claim adopted by another server instance, or by this one before a restart, was refused 401 while the backend still considered it live — the process-local write authority #4286 describes, one layer up. In fenced mode the backend minted the token and verifies it atomically with the append, so the server now defers to it and does not consult the store on the append path; with the flag off, validation is unchanged. The token stays opaque (spec §4): nothing inspects its shape. Because a backend without the extension ignores Write-Fence and creates the stream unfenced, a client that requires fencing must check for the `Write-Fence: true` echo before treating the stream as fenced [WF-02]. Fenced creates now HEAD for the echo (the client library does not surface the PUT response) and fenced forks check the PUT response, and both fail instead of yielding a stream the fence never protects. The pin against DurableStreamTestServer flips from "silently unfenced" to "refused". --- .../agents-backend-issued-write-tokens.md | 2 +- .../src/routing/stream-append.ts | 21 ++--- packages/agents-server/src/stream-client.ts | 33 +++++++- .../agents-server/test/stream-append.test.ts | 41 +++++++++- .../test/stream-client-fork.test.ts | 25 +++--- .../agents-server/test/stream-client.test.ts | 80 +++++++++++++++---- 6 files changed, 163 insertions(+), 39 deletions(-) diff --git a/.changeset/agents-backend-issued-write-tokens.md b/.changeset/agents-backend-issued-write-tokens.md index 32ff4c49aa..b8c4f5278c 100644 --- a/.changeset/agents-backend-issued-write-tokens.md +++ b/.changeset/agents-backend-issued-write-tokens.md @@ -3,4 +3,4 @@ '@electric-ax/agents-server': patch --- -Adopt backend-issued claim write tokens: when the Durable Streams backend implements the Write Fencing extension, the Agents Server adopts the `write_token` it delivers with wake notifications, pull claims, and heartbeat acks as the claim's write token, and the runtime refreshes its token from heartbeat responses. An opt-in `fencedSessionStreams` server option (env `ELECTRIC_AGENTS_FENCED_SESSION_STREAMS`) creates entity session streams with `Write-Fence: true` and forwards the token plus the fenced-class assertion on runtime appends, so the backend itself rejects deposed or lapsed writers. Off by default, and behaviour is unchanged when the backend supplies no token. +Adopt backend-issued claim write tokens: when the Durable Streams backend implements the Write Fencing extension, the Agents Server adopts the `write_token` it delivers with wake notifications, pull claims, and heartbeat acks as the claim's write token, and the runtime refreshes its token from heartbeat responses. An opt-in `fencedSessionStreams` server option (env `ELECTRIC_AGENTS_FENCED_SESSION_STREAMS`) creates entity session streams with `Write-Fence: true` and forwards the token plus the fenced-class assertion on runtime appends, so the backend itself is the write authority and rejects deposed or lapsed writers — including across a server restart or from another server instance, where the in-memory token store would not know the claim. In this mode a session-stream create or fork fails unless the backend echoes `Write-Fence: true`, so the option never yields silently unfenced streams. Off by default, and behaviour is unchanged when the backend supplies no token. diff --git a/packages/agents-server/src/routing/stream-append.ts b/packages/agents-server/src/routing/stream-append.ts index 86e794749d..7daa67413d 100644 --- a/packages/agents-server/src/routing/stream-append.ts +++ b/packages/agents-server/src/routing/stream-append.ts @@ -104,18 +104,21 @@ async function handleStreamAppend( if (entity) { const token = writeTokenFromHeaders(request.headers) - if (!manager.isValidWriteToken(entity, token)) { - return apiError(401, ErrCodeUnauthorized, `Invalid write token`) - } if (manager.fencedSessionStreams) { - // Forward the runtime's claim capability to the Durable Streams - // backend and assert the fenced write class, so a stale or lost token - // is a loud 401 downstream instead of a silent open-class write under - // this server's forwarded identity. The runtime's own bearer is - // overwritten with the server's when the request is forwarded, so - // `Write-Token` is the only carrier that survives. + // The backend minted this token for the claim and verifies it + // atomically with the append, so write authority is the backend's and + // not this process's memory: a claim adopted by another instance, or + // before a restart, is honoured here exactly as the backend honours + // it. Forward the token and assert the fenced write class, so a stale + // or lost token is a loud 401 downstream instead of a silent + // open-class write under this server's forwarded identity. The + // runtime's own bearer is overwritten with the server's when the + // request is forwarded, so `Write-Token` is the only carrier that + // survives. request.headers.set(WRITE_TOKEN_HEADER, token) request.headers.set(WRITE_FENCE_HEADER, `true`) + } else if (!manager.isValidWriteToken(entity, token)) { + return apiError(401, ErrCodeUnauthorized, `Invalid write token`) } if (manager.isForkWriteLockedEntity(entity.url)) { return apiError( diff --git a/packages/agents-server/src/stream-client.ts b/packages/agents-server/src/stream-client.ts index 31b45a0bfe..275a57f6c5 100644 --- a/packages/agents-server/src/stream-client.ts +++ b/packages/agents-server/src/stream-client.ts @@ -16,7 +16,8 @@ export type DurableStreamsBearerProvider = string | (() => MaybePromise) * on a create opts the stream into fenced appends; on an append it asserts * the fenced write class, which the backend only honours together with a * claim-scoped write token carried in `Write-Token`. A backend without the - * extension ignores both headers. + * extension ignores both headers, so a fenced create is only trusted once the + * backend echoes `Write-Fence: true` back [WF-02]. * * https://github.com/adityavkk/chronicle/blob/main/docs/spec/WRITE-FENCING.md */ @@ -264,6 +265,15 @@ export class StreamClient { body: opts.body, closed: opts.closed, }) + if (opts.writeFence) { + // The client library does not surface the PUT response; a fenced + // stream echoes on HEAD as well [WF-02]. + const head = await fetch(this.streamUrl(path), { + method: `HEAD`, + headers: await this.requestHeaders(), + }) + this.assertWriteFenced(path, head) + } }) } @@ -303,7 +313,10 @@ export class StreamClient { headers: await this.requestHeaders(headers), }) - if (response.ok) return + if (response.ok) { + if (opts?.writeFence) this.assertWriteFenced(path, response) + return + } throw new Error( `Stream fork failed: ${response.status} ${await response.text()}` @@ -311,6 +324,22 @@ export class StreamClient { }) } + /** + * A backend without the Write Fencing extension ignores `Write-Fence` and + * creates the stream unfenced, so a client that requires fencing checks for + * the echo before treating the stream as fenced [WF-02]. Throwing keeps + * `fencedSessionStreams` honest: a session stream is fenced, or its + * creation fails — never silently unfenced. + */ + private assertWriteFenced(path: string, response: Response): void { + if (response.headers.get(WRITE_FENCE_HEADER)?.toLowerCase() === `true`) { + return + } + throw new Error( + `Stream ${path} is not write-fenced: the Durable Streams backend did not echo ${WRITE_FENCE_HEADER} (fencedSessionStreams requires a backend implementing the Write Fencing extension)` + ) + } + async append( path: string, data: Uint8Array | string, diff --git a/packages/agents-server/test/stream-append.test.ts b/packages/agents-server/test/stream-append.test.ts index 841e148f74..99fd8659ea 100644 --- a/packages/agents-server/test/stream-append.test.ts +++ b/packages/agents-server/test/stream-append.test.ts @@ -16,6 +16,7 @@ const entity = { function buildRuntime(opts: { fencedSessionStreams: boolean hasEntity?: boolean + knownToken?: boolean }): ElectricAgentsStreamAppendRuntime { return { manager: { @@ -26,7 +27,7 @@ function buildRuntime(opts: { }, fencedSessionStreams: opts.fencedSessionStreams, isAttachmentStreamPath: vi.fn(() => false), - isValidWriteToken: vi.fn(() => true), + isValidWriteToken: vi.fn(() => opts.knownToken ?? true), isForkWriteLockedEntity: vi.fn(() => false), isForkWriteLockedStream: vi.fn(() => false), validateWriteEvent: vi.fn().mockResolvedValue(null), @@ -71,6 +72,44 @@ describe(`entity stream appends`, () => { expect(forwardedHeaders!.get(`Write-Fence`)).toBe(`true`) }) + it(`lets the backend judge a token the local store does not know when fenced session streams are enabled`, async () => { + // The claim may have been adopted by another instance, or by this one + // before a restart: in fenced mode the backend is the write authority. + let forwardedHeaders: Headers | null = null + const forward = vi.fn(async (req: { headers: Headers }) => { + forwardedHeaders = req.headers + return new Response(null, { status: 204 }) + }) + const runtime = buildRuntime({ + fencedSessionStreams: true, + knownToken: false, + }) + + const response = await electricAgentsStreamAppendRouter.fetch( + appendRequest(`/horton/demo/main`), + runtime, + forward as any + ) + + expect(response?.status).toBe(204) + expect(forwardedHeaders!.get(`Write-Token`)).toBe(`claim-token-1`) + expect(forwardedHeaders!.get(`Write-Fence`)).toBe(`true`) + expect(runtime.manager.isValidWriteToken).not.toHaveBeenCalled() + }) + + it(`rejects a token the local store does not know when fencing is off`, async () => { + const forward = vi.fn(async () => new Response(null, { status: 204 })) + + const response = await electricAgentsStreamAppendRouter.fetch( + appendRequest(`/horton/demo/main`), + buildRuntime({ fencedSessionStreams: false, knownToken: false }), + forward as any + ) + + expect(response?.status).toBe(401) + expect(forward).not.toHaveBeenCalled() + }) + it(`forwards appends without fencing headers by default`, async () => { let forwardedHeaders: Headers | null = null const forward = vi.fn(async (req: { headers: Headers }) => { diff --git a/packages/agents-server/test/stream-client-fork.test.ts b/packages/agents-server/test/stream-client-fork.test.ts index 9540bbba42..94e365762c 100644 --- a/packages/agents-server/test/stream-client-fork.test.ts +++ b/packages/agents-server/test/stream-client-fork.test.ts @@ -35,20 +35,25 @@ describe(`StreamClient.fork`, () => { ).resolves.toEqual({ offset: expect.any(String) }) }) - it(`create and fork with writeFence succeed against a backend without the extension`, async () => { - // A base Durable Streams server ignores the Write-Fence header, so the - // fencedSessionStreams knob is safe to enable against any backend. - await client.create(`/source-fenced`, { + it(`refuses writeFence against a backend without the extension`, async () => { + // A base Durable Streams server ignores the Write-Fence header and + // creates the stream unfenced, so a fenced create or fork must fail + // rather than hand back a stream the fence never protects [WF-02]. + await expect( + client.create(`/source-fenced`, { + contentType: `application/json`, + body: `[]`, + writeFence: true, + }) + ).rejects.toThrow(/did not echo Write-Fence/) + + await client.create(`/source-unfenced`, { contentType: `application/json`, body: `[]`, - writeFence: true, }) - - await client.fork(`/fork-fenced`, `/source-fenced`, { writeFence: true }) - await expect( - client.append(`/fork-fenced`, JSON.stringify({ type: `reconcile` })) - ).resolves.toEqual({ offset: expect.any(String) }) + client.fork(`/fork-fenced`, `/source-unfenced`, { writeFence: true }) + ).rejects.toThrow(/did not echo Write-Fence/) }) it(`preserves source history when reading the fork`, async () => { diff --git a/packages/agents-server/test/stream-client.test.ts b/packages/agents-server/test/stream-client.test.ts index c0579bbe71..1eb9c19303 100644 --- a/packages/agents-server/test/stream-client.test.ts +++ b/packages/agents-server/test/stream-client.test.ts @@ -226,28 +226,61 @@ describe(`StreamClient`, () => { it(`opts a stream into write fencing only when asked`, async () => { const createMock = vi.mocked(DurableStream.create).mockClear() createMock.mockResolvedValue(undefined as never) + const fetchMock = vi + .spyOn(globalThis, `fetch`) + .mockResolvedValue( + new Response(null, { status: 200, headers: { 'Write-Fence': `true` } }) + ) const client = new StreamClient(`http://127.0.0.1:4545`) - await client.create(`/horton/demo/main`, { - contentType: `application/json`, - writeFence: true, - }) - await client.create(`/horton/other/main`, { - contentType: `application/json`, - }) + try { + await client.create(`/horton/demo/main`, { + contentType: `application/json`, + writeFence: true, + }) + await client.create(`/horton/other/main`, { + contentType: `application/json`, + }) - expect(createMock.mock.calls[0]![0]).toMatchObject({ - headers: { 'Write-Fence': `true` }, - }) - expect(createMock.mock.calls[1]![0]!.headers ?? {}).not.toHaveProperty( - `Write-Fence` - ) + expect(createMock.mock.calls[0]![0]).toMatchObject({ + headers: { 'Write-Fence': `true` }, + }) + expect(createMock.mock.calls[1]![0]!.headers ?? {}).not.toHaveProperty( + `Write-Fence` + ) + // Only the fenced create checks for the echo [WF-02]. + expect(fetchMock).toHaveBeenCalledTimes(1) + expect(fetchMock.mock.calls[0]?.[1]?.method).toBe(`HEAD`) + } finally { + fetchMock.mockRestore() + } + }) + + it(`refuses a fenced create the backend did not echo`, async () => { + vi.mocked(DurableStream.create).mockResolvedValue(undefined as never) + const fetchMock = vi + .spyOn(globalThis, `fetch`) + .mockResolvedValue(new Response(null, { status: 200 })) + const client = new StreamClient(`http://127.0.0.1:4545`) + + try { + await expect( + client.create(`/horton/demo/main`, { + contentType: `application/json`, + writeFence: true, + }) + ).rejects.toThrow(/did not echo Write-Fence/) + } finally { + fetchMock.mockRestore() + } }) it(`opts a fork into write fencing only when asked`, async () => { - const fetchMock = vi.spyOn(globalThis, `fetch`).mockResolvedValue( - new Response(null, { status: 200 }) - ) + const fetchMock = vi + .spyOn(globalThis, `fetch`) + .mockResolvedValue( + new Response(null, { status: 200, headers: { 'Write-Fence': `true` } }) + ) const client = new StreamClient(`http://127.0.0.1:4545`) try { @@ -265,6 +298,21 @@ describe(`StreamClient`, () => { } }) + it(`refuses a fenced fork the backend did not echo`, async () => { + const fetchMock = vi + .spyOn(globalThis, `fetch`) + .mockResolvedValue(new Response(null, { status: 200 })) + const client = new StreamClient(`http://127.0.0.1:4545`) + + try { + await expect( + client.fork(`/fork/main`, `/source/main`, { writeFence: true }) + ).rejects.toThrow(/did not echo Write-Fence/) + } finally { + fetchMock.mockRestore() + } + }) + it(`preserves claim token authorization on subscription ack`, async () => { const fetchMock = vi.spyOn(globalThis, `fetch`).mockResolvedValueOnce( new Response(JSON.stringify({ ok: true }), {