From 689da08157006727392b02ca8ac873e501a4c3ea Mon Sep 17 00:00:00 2001 From: Aditya Kumarakrishnan Date: Mon, 31 Aug 2026 13:53:37 +0530 Subject: [PATCH 1/6] 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/6] 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 }), { From af5dd05419cdd9ff10568bb67a0461fb577db311 Mon Sep 17 00:00:00 2001 From: Aditya Kumarakrishnan Date: Wed, 2 Sep 2026 13:11:01 +0530 Subject: [PATCH 3/6] fix(agents-server): expire claim write tokens the store never sees released MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The ClaimWriteTokenStore only ever shrank on a done or a kill handled by the same process. A server instance that received a wake's delivery but not its claim kept the delivered token forever, and one that minted or refreshed a claim whose done reached a sibling instance kept the claim entry — and kept answering isValid() for it — forever. Entries now carry an expiry of the default claim lease plus the grace a backend keeps a write token alive past it (PROTOCOL §7 lease_ttl_ms, 30s + 5s), refreshed by every mint and by a new renew() the heartbeat path calls when the backend answers without re-minting, so a claim this process is the sole authority for lives as long as it heartbeats. A lazy sweep on the paths that add entries drops what expired. The public methods keep their signatures; the constructor takes an optional { ttlMs, now } for tests. --- .../src/claim-write-token-store.ts | 106 +++++++++++++++--- .../test/claim-write-token-store.test.ts | 70 ++++++++++++ 2 files changed, 161 insertions(+), 15 deletions(-) diff --git a/packages/agents-server/src/claim-write-token-store.ts b/packages/agents-server/src/claim-write-token-store.ts index 1e350fccea..88cd3ac3dc 100644 --- a/packages/agents-server/src/claim-write-token-store.ts +++ b/packages/agents-server/src/claim-write-token-store.ts @@ -10,12 +10,42 @@ interface ActiveClaimWriteToken { */ previousToken?: string consumerId: string + expiresAt: number +} + +interface DeliveredWriteToken { + token: string + expiresAt: number +} + +/** + * How long an entry stays valid without a mint or renewal: the Durable + * Streams default claim lease (30s, PROTOCOL §7) plus the grace a backend + * keeps a write token alive past the lease. A heartbeat renews the entry, so + * only a claim that stopped heartbeating — or one whose heartbeats and done + * all reached other server instances — ages out. + */ +const DEFAULT_TTL_MS = 30_000 + 5_000 + +export interface ClaimWriteTokenStoreOptions { + ttlMs?: number + now?: () => number } export class ClaimWriteTokenStore { private readonly claimsByStream = new Map() private readonly streamKeysByConsumer = new Map>() - private readonly deliveredTokensByConsumer = new Map() + private readonly deliveredTokensByConsumer = new Map< + string, + DeliveredWriteToken + >() + private readonly ttlMs: number + private readonly now: () => number + + constructor(options: ClaimWriteTokenStoreOptions = {}) { + this.ttlMs = options.ttlMs ?? DEFAULT_TTL_MS + this.now = options.now ?? Date.now + } mint( service: string, @@ -23,6 +53,7 @@ export class ClaimWriteTokenStore { consumerId: string, token: string = randomUUID() ): string { + this.sweep() const streamKey = this.streamKey(service, streamPath) const consumerKey = this.consumerKey(service, consumerId) const previousClaimForStream = this.claimsByStream.get(streamKey) @@ -36,6 +67,7 @@ export class ClaimWriteTokenStore { this.claimsByStream.set(streamKey, { token, consumerId, + expiresAt: this.now() + this.ttlMs, ...(previousClaimForStream?.consumerId === consumerId ? { previousToken: previousClaimForStream.token } : {}), @@ -44,6 +76,23 @@ export class ClaimWriteTokenStore { return token } + /** + * Extends the expiry of every entry a consumer holds. Called on a + * heartbeat the backend answered without re-minting, so an entry this + * process is the sole authority for (no Write Fencing extension) lives as + * long as the claim does. + */ + renew(service: string, consumerId: string): void { + const consumerKey = this.consumerKey(service, consumerId) + const expiresAt = this.now() + this.ttlMs + const delivered = this.deliveredTokensByConsumer.get(consumerKey) + if (delivered) delivered.expiresAt = expiresAt + for (const streamKey of this.streamKeysByConsumer.get(consumerKey) ?? []) { + const claim = this.claimsByStream.get(streamKey) + if (claim) claim.expiresAt = expiresAt + } + } + /** * Remembers a write token the Durable Streams backend issued with a wake * delivery (webhook notification or pull-wake claim), keyed by the wake's @@ -51,25 +100,23 @@ export class ClaimWriteTokenStore { * claim write token. */ recordDelivered(service: string, consumerId: string, token: string): void { - this.deliveredTokensByConsumer.set( - this.consumerKey(service, consumerId), - token - ) + this.sweep() + this.deliveredTokensByConsumer.set(this.consumerKey(service, consumerId), { + token, + expiresAt: this.now() + this.ttlMs, + }) } 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 + const delivered = this.deliveredTokensByConsumer.get(consumerKey) + if (delivered === undefined) return undefined + this.deliveredTokensByConsumer.delete(consumerKey) + return delivered.expiresAt > this.now() ? delivered.token : undefined } isValid(service: string, streamPath: string, token: string): boolean { - const activeClaim = this.claimsByStream.get( - this.streamKey(service, streamPath) - ) + const activeClaim = this.liveClaim(this.streamKey(service, streamPath)) return ( activeClaim !== undefined && (activeClaim.token === token || activeClaim.previousToken === token) @@ -78,8 +125,8 @@ export class ClaimWriteTokenStore { owns(service: string, streamPath: string, consumerId: string): boolean { return ( - this.claimsByStream.get(this.streamKey(service, streamPath)) - ?.consumerId === consumerId + this.liveClaim(this.streamKey(service, streamPath))?.consumerId === + consumerId ) } @@ -107,6 +154,35 @@ export class ClaimWriteTokenStore { } } + private liveClaim(streamKey: string): ActiveClaimWriteToken | undefined { + const claim = this.claimsByStream.get(streamKey) + return claim !== undefined && claim.expiresAt > this.now() + ? claim + : undefined + } + + /** + * Drops expired entries. Lazy — run from the paths that add entries — so a + * server instance that saw a wake's delivery or claim but none of its + * later callbacks does not accumulate them forever. + */ + private sweep(): void { + const now = this.now() + for (const [consumerKey, streamKeys] of this.streamKeysByConsumer) { + for (const streamKey of streamKeys) { + const claim = this.claimsByStream.get(streamKey) + if (claim !== undefined && claim.expiresAt > now) continue + this.claimsByStream.delete(streamKey) + this.removeConsumerStream(consumerKey, streamKey) + } + } + for (const [consumerKey, delivered] of this.deliveredTokensByConsumer) { + if (delivered.expiresAt <= now) { + this.deliveredTokensByConsumer.delete(consumerKey) + } + } + } + private addConsumerStream(consumerKey: string, streamKey: string): void { let streamKeys = this.streamKeysByConsumer.get(consumerKey) if (!streamKeys) { 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 d11cce5ec9..086a8834cf 100644 --- a/packages/agents-server/test/claim-write-token-store.test.ts +++ b/packages/agents-server/test/claim-write-token-store.test.ts @@ -108,4 +108,74 @@ describe(`ClaimWriteTokenStore`, () => { expect(store.takeDelivered(`tenant-a`, `wake-1`)).toBeUndefined() expect(store.takeDelivered(`tenant-b`, `wake-1`)).toBe(`token-b`) }) + + describe(`expiry`, () => { + function clockedStore(ttlMs = 1_000) { + const clock = { now: 0 } + return { + clock, + store: new ClaimWriteTokenStore({ ttlMs, now: () => clock.now }), + } + } + + it(`expires a claim that is neither re-minted nor renewed within the TTL`, () => { + const { clock, store } = clockedStore() + + const token = store.mint(`tenant-a`, `/one/main`, `wake-1`) + clock.now = 999 + expect(store.isValid(`tenant-a`, `/one/main`, token)).toBe(true) + expect(store.owns(`tenant-a`, `/one/main`, `wake-1`)).toBe(true) + + clock.now = 1_000 + expect(store.isValid(`tenant-a`, `/one/main`, token)).toBe(false) + expect(store.owns(`tenant-a`, `/one/main`, `wake-1`)).toBe(false) + }) + + it(`re-minting and renewing both extend the claim`, () => { + const { clock, store } = clockedStore() + + const first = store.mint(`tenant-a`, `/one/main`, `wake-1`) + clock.now = 800 + const second = store.mint(`tenant-a`, `/one/main`, `wake-1`) + clock.now = 1_500 + expect(store.isValid(`tenant-a`, `/one/main`, first)).toBe(true) + expect(store.isValid(`tenant-a`, `/one/main`, second)).toBe(true) + + store.renew(`tenant-a`, `wake-1`) + clock.now = 2_400 + expect(store.isValid(`tenant-a`, `/one/main`, second)).toBe(true) + expect(store.owns(`tenant-a`, `/one/main`, `wake-1`)).toBe(true) + }) + + it(`expires a delivered token nobody claimed`, () => { + const { clock, store } = clockedStore() + + store.recordDelivered(`tenant-a`, `wake-1`, `backend-token`) + store.recordDelivered(`tenant-a`, `wake-2`, `backend-token-2`) + clock.now = 500 + store.renew(`tenant-a`, `wake-2`) + clock.now = 1_000 + + expect(store.takeDelivered(`tenant-a`, `wake-1`)).toBeUndefined() + expect(store.takeDelivered(`tenant-a`, `wake-2`)).toBe(`backend-token-2`) + }) + + it(`sweeps expired entries so a later mint by a different consumer starts clean`, () => { + const { clock, store } = clockedStore() + + const stale = store.mint(`tenant-a`, `/one/main`, `wake-1`) + store.recordDelivered(`tenant-a`, `wake-1`, `delivered-1`) + clock.now = 1_000 + // A mint for another stream triggers the sweep; the stale consumer's + // entries are gone, not merely hidden, so clearing it later is a no-op + // that cannot touch a successor's claim. + store.mint(`tenant-a`, `/two/main`, `wake-2`) + const successor = store.mint(`tenant-a`, `/one/main`, `wake-3`) + store.clearConsumer(`tenant-a`, `wake-1`) + + expect(store.isValid(`tenant-a`, `/one/main`, stale)).toBe(false) + expect(store.isValid(`tenant-a`, `/one/main`, successor)).toBe(true) + expect(store.takeDelivered(`tenant-a`, `wake-1`)).toBeUndefined() + }) + }) }) From 01825eb5d9a4304e798c15aa506f4a303303a7a8 Mon Sep 17 00:00:00 2001 From: Aditya Kumarakrishnan Date: Wed, 2 Sep 2026 15:11:43 +0530 Subject: [PATCH 4/6] feat(agents-server): make wake claims and dones safe across server instances MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit With more than one server instance behind one address, the runtime's claim callback for a wake need not reach the instance that received the delivery, and its done need not reach the one that handled the claim. Three defects share that root, all in the wake-callback route: - The claim was answered locally and never reached the backend. On the instance that had not seen the delivery, takeDelivered() found nothing and the store minted a random token; with fencedSessionStreams on, the runtime's first append was refused 401, the wake failed, and its done — sent at the delivery's tail offset — acked the message away with nothing on the stream. Silent loss, not a retry. - Webhook wakes never materialised a consumer_claims row, so a done on an instance that had not minted the token found neither entityCleared nor an owning store entry and left the entity `running` (or `stopping`, from which nothing else ever leaves). - The status write was a read of entity.status followed by updateStatus(), which a done racing the next wake's delivery could clobber. A claim is now forwarded like a heartbeat: a non-done callback that renews the lease and, on a backend with the Write Fencing extension, returns the wake's write token (PROTOCOL §7.1), which the claim adopts over anything in local memory. A done carries a wake id too, so the claim path excludes it explicitly — routed as a claim, a done would be refused for a write token no backend mints for it and would never release the claim. In fenced mode a claim the backend issues no token for is refused (502 WRITE_TOKEN_UNAVAILABLE) rather than answered with a token the backend would reject; the runtime drops the wake unclaimed and the lease lapse re-wakes it. Without fencing the store is still the only write authority, so nothing is gained by holding the wake on the backend's answer: a claim it refuses, or never answers, falls back to a local mint and the wake runs exactly as it did before claims were forwarded. The delivered token and the store's own mint remain the fallback when fencing is off. Once the backend has accepted the claim it is materialised in consumer_claims (webhook wakes only — pull-wake claims already were), so materializeReleasedClaim's compare-and-clear yields entityCleared on any instance. The same insert fences a redelivered wake's second claim: an active row for the same (consumerId, epoch) answers 409 WAKE_ALREADY_CLAIMED before any token is minted, so the duplicate never claims, runs, or acks. The backend's own (generation, wake_id) fence rejects lapsed and re-armed wakes before this point, so an active row can only mean a wake another runtime is still processing. The status transition is one conditional statement — running → idle, stopping → stopped, nothing else touched — behind the unchanged `entityCleared || stillOwnsClaim` gate. It is a statement of its own rather than part of the release: the next wake can be delivered in between and set `running` again, which this write would report `idle` under until that wake's own done settles it. The done clears everything the consumer holds in the store (clearConsumer, consumer-scoped so a newer wake's entry survives), which also drops the delivered token an instance kept for a claim it never handled. A runtime refusing a forwarded wake reverts the entity to idle, as a failed forward already did, but only while it is still `running` — a rolling runtime restart refuses wakes in bulk, and an unconditional write would report a stopping entity, or one a sibling has already woken again, wrong. --- .changeset/agents-replica-safe-claims.md | 5 + .../src/claim-write-token-store.ts | 24 +- .../src/electric-agents-types.ts | 2 + packages/agents-server/src/entity-registry.ts | 76 ++- .../src/routing/internal-router.ts | 243 +++++-- .../test/consumer-claim-registry.test.ts | 133 ++++ .../test/server-claim-write-token.test.ts | 9 +- .../subscription-webhooks-routing.test.ts | 639 ++++++++++++++++-- 8 files changed, 1018 insertions(+), 113 deletions(-) create mode 100644 .changeset/agents-replica-safe-claims.md diff --git a/.changeset/agents-replica-safe-claims.md b/.changeset/agents-replica-safe-claims.md new file mode 100644 index 0000000000..fe79cf9dc3 --- /dev/null +++ b/.changeset/agents-replica-safe-claims.md @@ -0,0 +1,5 @@ +--- +'@electric-ax/agents-server': patch +--- + +Make wake claims and dones safe across Agents Server instances. A runtime's claim callback is now forwarded to the Durable Streams backend as a lease-renewing callback, like a heartbeat, so the claim adopts the write token the backend re-mints for it rather than one held in the memory of whichever instance happened to receive the delivery; with `fencedSessionStreams` on, a claim the backend issues no token for is refused (`502 WRITE_TOKEN_UNAVAILABLE`) instead of being answered with a token the backend would reject. The claim is materialised in `consumer_claims` once the backend has accepted it, which lets a done that lands on a different instance release it and settle the entity's status (`running` → `idle`, `stopping` → `stopped`) in one conditional update that cannot clobber a newer wake; a redelivered wake's second claim is refused with `409 WAKE_ALREADY_CLAIMED`. Without `fencedSessionStreams` this store remains the sole write authority, so a claim the backend refuses or never answers still falls back to a local mint and the wake runs as before. In-memory claim write tokens expire a grace period past the claim lease unless a heartbeat refreshes them, and a runtime that refuses a forwarded wake reverts the entity to `idle` only if it is still `running`, as a failed forward now does too. diff --git a/packages/agents-server/src/claim-write-token-store.ts b/packages/agents-server/src/claim-write-token-store.ts index 88cd3ac3dc..7424b13e3e 100644 --- a/packages/agents-server/src/claim-write-token-store.ts +++ b/packages/agents-server/src/claim-write-token-store.ts @@ -19,13 +19,25 @@ interface DeliveredWriteToken { } /** - * How long an entry stays valid without a mint or renewal: the Durable - * Streams default claim lease (30s, PROTOCOL §7) plus the grace a backend - * keeps a write token alive past the lease. A heartbeat renews the entry, so - * only a claim that stopped heartbeating — or one whose heartbeats and done - * all reached other server instances — ages out. + * The Durable Streams default claim lease (PROTOCOL §7): a webhook wake + * carries no `lease_ttl_ms`, so this is what a claim confirmed through the + * wake callback is recorded with, and what a claim's write token is kept + * alive for. */ -const DEFAULT_TTL_MS = 30_000 + 5_000 +export const DEFAULT_CLAIM_LEASE_MS = 30_000 + +/** + * The grace a backend keeps a write token valid past the lease it was minted + * for. + */ +const WRITE_TOKEN_FENCE_GRACE_MS = 5_000 + +/** + * How long an entry stays valid without a mint or renewal. A heartbeat renews + * the entry, so only a claim that stopped heartbeating — or one whose + * heartbeats and done all reached other server instances — ages out. + */ +const DEFAULT_TTL_MS = DEFAULT_CLAIM_LEASE_MS + WRITE_TOKEN_FENCE_GRACE_MS export interface ClaimWriteTokenStoreOptions { ttlMs?: number diff --git a/packages/agents-server/src/electric-agents-types.ts b/packages/agents-server/src/electric-agents-types.ts index 82afc3af1f..0a2ad73e1a 100644 --- a/packages/agents-server/src/electric-agents-types.ts +++ b/packages/agents-server/src/electric-agents-types.ts @@ -630,3 +630,5 @@ export const ErrCodeEntityPersistFailed = `ENTITY_PERSIST_FAILED` export const ErrCodeAgentUiNotFound = `AGENT_UI_NOT_FOUND` export const ErrCodeSubscriptionNotFound = `SUBSCRIPTION_NOT_FOUND` export const ErrCodeWakeCallbackNotFound = `WAKE_CALLBACK_NOT_FOUND` +export const ErrCodeWakeAlreadyClaimed = `WAKE_ALREADY_CLAIMED` +export const ErrCodeWriteTokenUnavailable = `WRITE_TOKEN_UNAVAILABLE` diff --git a/packages/agents-server/src/entity-registry.ts b/packages/agents-server/src/entity-registry.ts index 7dfb1ceb0a..e6537bacb0 100644 --- a/packages/agents-server/src/entity-registry.ts +++ b/packages/agents-server/src/entity-registry.ts @@ -149,6 +149,22 @@ export interface MaterializeActiveClaimInput { runnerId?: string claimedAt?: Date leaseExpiresAt?: Date + /** + * Leave an existing active claim for the same (consumerId, epoch) alone + * and report false instead of replacing it. The claim callback sets this: + * the backend accepts every callback for a wake it still holds, so a + * redelivered wake's second claim is indistinguishable there, and this + * row is where the duplicate is fenced before a second runtime runs it. + * + * The fence is on `status` alone — `lease_expires_at` is informational, + * because heartbeats do not extend it and a wake outliving one lease would + * otherwise stop fencing its own duplicates. The cost: a claim whose wake + * never sends a done stays `active` until a reconciler sweeps it, and a + * redelivery of that wake is fenced until the backend re-arms it with a + * new generation (one lease of delay, no loss — the re-armed wake carries + * a new `(wake_id, generation)` and so a new row). + */ + unlessActive?: boolean } export interface MaterializeHeartbeatClaimInput { @@ -425,12 +441,18 @@ export class PostgresRegistry { return rows[0] ? this.rowToRunner(rows[0]) : null } + /** + * Returns false only with `unlessActive`, when an active claim for the + * same (consumerId, epoch) already exists. Postgres serialises the + * conflicting upserts, so two concurrent duplicate claims cannot both + * succeed. + */ async materializeActiveClaim( input: MaterializeActiveClaimInput - ): Promise { + ): Promise { const claimedAt = input.claimedAt ?? new Date() - await this.db.transaction(async (tx) => { - await tx + return await this.db.transaction(async (tx) => { + const claimed = await tx .insert(consumerClaims) .values({ tenantId: this.tenantId, @@ -462,7 +484,12 @@ export class PostgresRegistry { releasedAt: null, updatedAt: claimedAt, }, + ...(input.unlessActive + ? { setWhere: ne(consumerClaims.status, `active`) } + : {}), }) + .returning({ consumerId: consumerClaims.consumerId }) + if (claimed.length === 0) return false await tx .insert(entityDispatchState) @@ -489,6 +516,7 @@ export class PostgresRegistry { updatedAt: claimedAt, }, }) + return true }) } @@ -1403,6 +1431,48 @@ export class PostgresRegistry { .where(whereClause) } + /** + * Settles the entity's status once its wake is done: `running` → `idle`, + * `stopping` → `stopped`. One conditional statement rather than a read + * followed by `updateStatus`, because a done can land after the backend + * has already delivered the next wake — which set `running` again — and a + * read-then-write would clobber that. Returns the status written, or null + * when the entity was in neither state. + */ + async updateStatusAfterDone(entityUrl: string): Promise { + const rows = await this.db + .update(entities) + .set({ + status: sql`CASE WHEN ${entities.status} = 'stopping' THEN 'stopped' ELSE 'idle' END`, + updatedAt: Date.now(), + }) + .where( + and( + this.entityWhere(entityUrl), + inArray(entities.status, [`running`, `stopping`]) + ) + ) + .returning({ status: entities.status }) + return rows[0] ? (rows[0].status as EntityStatus) : null + } + + /** + * Reverts the `running` set for a wake the runtime never took: a refused or + * failed forward. Conditional, because the wake that set `running` is not + * necessarily the newest one — a rolling runtime restart refuses wakes for + * entities that are stopping, or that a sibling replica has already woken + * again, and an unconditional `idle` would report those wrong. Returns true + * when the status was reverted. + */ + async revertRunningStatus(entityUrl: string): Promise { + const rows = await this.db + .update(entities) + .set({ status: `idle`, updatedAt: Date.now() }) + .where(and(this.entityWhere(entityUrl), eq(entities.status, `running`))) + .returning({ url: entities.url }) + return rows.length > 0 + } + async updateStatusWithTxid( entityUrl: string, status: EntityStatus diff --git a/packages/agents-server/src/routing/internal-router.ts b/packages/agents-server/src/routing/internal-router.ts index d03403002d..3b403e57e6 100644 --- a/packages/agents-server/src/routing/internal-router.ts +++ b/packages/agents-server/src/routing/internal-router.ts @@ -16,14 +16,17 @@ import { } from '../electric-agents-http.js' import { consumerCallbacks, subscriptionWebhooks } from '../db/schema.js' import { + ErrCodeWakeAlreadyClaimed, ErrCodeWakeCallbackNotFound, ErrCodeForkInProgress, ErrCodeSubscriptionNotFound, ErrCodeUnauthorized, + ErrCodeWriteTokenUnavailable, } from '../electric-agents-types.js' import { ATTR, tracer } from '../tracing.js' import { decodeJsonObject } from '../utils/server-utils.js' import { serverLog } from '../utils/log.js' +import { DEFAULT_CLAIM_LEASE_MS } from '../claim-write-token-store.js' import { applyDurableStreamsBearer } from '../stream-client.js' import { getDefaultWebhookSigner } from '../webhook-signing.js' import { resolveDurableStreamsRoutingAdapter } from './durable-streams-routing-adapter.js' @@ -41,6 +44,7 @@ import type { WebhookSourceContract, WebhookSignatureVerifierConfig, } from '@electric-ax/agents-runtime' +import type { ElectricAgentsEntity } from '../electric-agents-types.js' import type { TenantContext } from './context.js' import type { DurableStreamsRoutingAdapter } from './durable-streams-routing-adapter.js' import type { WebhookSigner } from '../webhook-signing.js' @@ -670,7 +674,7 @@ async function subscriptionWebhook( ) } catch (err) { if (runningEntityUrl) { - await ctx.entityManager.registry.updateStatus(runningEntityUrl, `idle`) + await ctx.entityManager.registry.revertRunningStatus(runningEntityUrl) } return apiError( 502, @@ -682,6 +686,13 @@ async function subscriptionWebhook( const responseBytes = upstream.body ? new Uint8Array(await upstream.arrayBuffer()) : new Uint8Array() + if (!upstream.ok && runningEntityUrl) { + // The runtime refused the wake (draining, unknown type), so nothing will + // claim or finish it; the backend retries later. Leaving `running` set + // for that whole backoff misreports the entity, exactly as a failed + // forward would. + await ctx.entityManager.registry.revertRunningStatus(runningEntityUrl) + } return responseFromUpstream(upstream, responseBytes) } @@ -723,28 +734,44 @@ async function wakeCallback( ) if (!parsedBodyResult.ok) return parsedBodyResult.response const requestBody = parsedBodyResult.value as WakeCallbackBody | undefined - const isClaimRequest = - requestBody?.wakeId !== undefined || requestBody?.wake_id !== undefined const isDoneRequest = requestBody?.done === true + // A done carries `wake_id` too (PROTOCOL §7.1) — only today's runtime + // omits it. Routed as a claim it would be answered + // WRITE_TOKEN_UNAVAILABLE (a backend mints no token for a done) and would + // never reach the release below, leaving the claim and the entity's + // dispatch state pointing at a wake that finished. + const isClaimCallback = + !isDoneRequest && + (requestBody?.wakeId !== undefined || requestBody?.wake_id !== undefined) + const epoch = requestBody?.generation ?? requestBody?.epoch + const subscriptionId = durableStreamsSubscriptionCallback(target.callbackUrl) const headers = forwardHeadersFromRequest(request) headers.delete(`content-length`) - if (isClaimRequest && !isDoneRequest) { - let responseBody: Record = { ok: true } - if (target.primaryStream) { - const writeToken = await mintClaimWriteToken( - ctx, - target.primaryStream, - consumerId - ) - if (writeToken) { - responseBody = { ...responseBody, writeToken } - } - } - return json(responseBody) + // Without stream fencing this store is the only write authority, so a claim + // the backend refused — or never answered — is still safe to run locally: + // the round trip is there to fetch the backend's token, not to gate the + // wake, and answering the runtime with the backend's error would drop a + // wake that used to run and make it wait out a lease. With fencing on there + // is no such fallback: a token minted here is refused on the first append. + const mintClaimLocally = async (reason: string): Promise => { + serverLog.warn( + `[wake-callback] claim not confirmed by the backend (${reason}); minting locally for consumer=${consumerId} (stream fencing off)` + ) + const writeToken = target.primaryStream + ? await mintClaimWriteToken(ctx, target.primaryStream, consumerId) + : undefined + return json(writeToken ? { ok: true, writeToken } : { ok: true }) } + // A claim is forwarded like a heartbeat rather than answered here: as a + // non-done callback it renews the lease and, on a backend with the Write + // Fencing extension, returns the wake's write token (PROTOCOL §7.1). The + // server instance handling the claim is not necessarily the one that + // received the delivery, so the token the delivery carried may live in + // another process's memory; the backend is the only party that can hand it + // to whichever instance the runtime reached. const upstreamBody = encodeWakeCallbackBody( ctx.service, consumerId, @@ -757,9 +784,6 @@ async function wakeCallback( let upstream: Response try { - const subscriptionId = durableStreamsSubscriptionCallback( - target.callbackUrl - ) if (subscriptionId) { const token = claimTokenFromRequest(request) if (!token) { @@ -784,24 +808,74 @@ async function wakeCallback( }) } } catch (err) { - return apiError( - 502, - `WAKE_CALLBACK_FAILED`, - err instanceof Error ? err.message : String(err) - ) + const message = err instanceof Error ? err.message : String(err) + if (isClaimCallback && !ctx.entityManager.fencedSessionStreams) { + return await mintClaimLocally(message) + } + return apiError(502, `WAKE_CALLBACK_FAILED`, message) } let responseBytes: Uint8Array = upstream.body ? new Uint8Array(await upstream.arrayBuffer()) : new Uint8Array() - if (isClaimRequest && upstream.ok && target.primaryStream) { + if ( + isClaimCallback && + !upstream.ok && + !ctx.entityManager.fencedSessionStreams + ) { + return await mintClaimLocally(`answered ${upstream.status}`) + } + + if (isClaimCallback && upstream.ok && target.primaryStream) { const responseBody = decodeJsonObject(responseBytes) if (responseBody?.ok === true) { + const entity = await ctx.entityManager.registry.getEntityByStream( + target.primaryStream + ) + const backendToken = backendWriteToken(responseBody) + if ( + backendToken === undefined && + ctx.entityManager.fencedSessionStreams + ) { + // The backend is the write authority for fenced streams + // (stream-append.ts), so a token minted here would be refused on the + // runtime's first append, and `ok` without one lets the runtime + // claim, fail every write and then ack the message away. Refuse the + // claim instead: the runtime drops the wake unclaimed and the backend + // re-wakes it when the lease lapses. + return apiError( + 502, + ErrCodeWriteTokenUnavailable, + `Backend issued no write token for the claim` + ) + } + if (entity && epoch !== undefined && !subscriptionId) { + // Pull-wake claims were materialised when the runner claimed them + // (runners-router); webhook wakes only now, once the backend has + // confirmed the claim is live, so no row is left behind for a wake + // that was delivered but never claimed. + const claimed = await materializeCallbackClaim( + ctx, + entity, + target.primaryStream, + consumerId, + epoch + ) + if (!claimed) { + return apiError( + 409, + ErrCodeWakeAlreadyClaimed, + `Wake is already claimed` + ) + } + } const writeToken = await mintClaimWriteToken( ctx, target.primaryStream, - consumerId + consumerId, + backendToken, + entity ) if (writeToken) { responseBody.writeToken = writeToken @@ -814,12 +888,8 @@ async function wakeCallback( // 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 refreshedToken = backendWriteToken(responseBody) + if (responseBody?.ok === true && refreshedToken !== undefined) { const writeToken = await mintClaimWriteToken( ctx, target.primaryStream, @@ -830,11 +900,12 @@ async function wakeCallback( responseBody.writeToken = writeToken responseBytes = new TextEncoder().encode(JSON.stringify(responseBody)) } + } else if (responseBody?.ok === true) { + ctx.runtime.claimWriteTokens.renew(ctx.service, consumerId) } } try { - const epoch = requestBody?.generation ?? requestBody?.epoch if ( upstream.ok && !isDoneRequest && @@ -888,23 +959,33 @@ async function wakeCallback( const entity = await ctx.entityManager.registry.getEntityByStream(releasedClaimStream) - // Transition entity back to idle when either signal says it's safe: + // Transition the entity out of its run when either signal says it's + // safe: // - entityCleared: our release just cleared the entity's active - // dispatch state, so no in-flight wake remains. + // dispatch state, so no in-flight wake remains. This is the durable + // signal, and the one that holds on a server instance other than + // the one that minted the claim's write token. // - stillOwnsClaim: this consumer is still the in-memory write-token - // owner, so no newer wake has displaced it. Covers two cases: - // (a) retry of a failed done (first attempt cleared the DB state - // but failed to update status), (b) server restart scenarios where - // the token is intact even though entityDispatchState may diverge. + // owner, so no newer wake has displaced it. Covers a retry of a + // failed done (first attempt cleared the DB state but failed to + // update status) on the minting instance. // If both are false, a newer wake owns the entity — leave status as-is. + // + // The settle is a statement of its own, not part of the release: the + // next wake can be delivered in between and set `running` again, and + // this write would then report `idle` under it. It is self-correcting + // (that wake's own done settles the status, and `updateStatusAfterDone` + // never touches a terminal status), which is why the release keeps the + // narrower job of clearing the dispatch state. if (entity && (entityCleared || stillOwnsClaim)) { - await ctx.entityManager.registry.updateStatus( - entity.url, - entity.status === `stopping` ? `stopped` : `idle` + const status = await ctx.entityManager.registry.updateStatusAfterDone( + entity.url ) - await ctx.entityBridgeManager.onEntityChanged(entity.url) + if (status !== null) { + await ctx.entityBridgeManager.onEntityChanged(entity.url) + } serverLog.info( - `[wake-callback] status updated after done for ${entity.url}` + `[wake-callback] status updated after done for ${entity.url}: ${status ?? `unchanged (${entity.status})`}` ) } else if (!entity) { serverLog.warn( @@ -912,17 +993,15 @@ async function wakeCallback( ) } - // Clear the in-memory write token only if this consumer still owns it. - // If a newer wake has taken over, that newer wake owns the token now - // and we must not clear it out from under it. - if (stillOwnsClaim) { - ctx.runtime.claimWriteTokens.clearStream( - ctx.service, - releasedClaimStream - ) - } else if (entity) { + // Everything this consumer holds is finished: its claim entries (if + // this instance minted or refreshed them) and its delivered token (if + // this instance received the delivery but another one handled the + // claim). Consumer-scoped, so a newer wake's entry for the same stream + // is untouched. + ctx.runtime.claimWriteTokens.clearConsumer(ctx.service, consumerId) + if (!stillOwnsClaim && entity) { serverLog.info( - `[wake-callback] done arrived after in-memory token evicted (stream=${releasedClaimStream} consumer=${consumerId})` + `[wake-callback] done for a claim this instance did not mint (stream=${releasedClaimStream} consumer=${consumerId})` ) } } else if (requestBody?.done === true) { @@ -943,29 +1022,69 @@ async function wakeCallback( return responseFromUpstream(upstream, responseBytes) } +function backendWriteToken( + responseBody: Record | null +): string | undefined { + const token = responseBody?.write_token + return typeof token === `string` && token !== `` ? token : undefined +} + +async function materializeCallbackClaim( + ctx: TenantContext, + entity: ElectricAgentsEntity, + streamPath: string, + consumerId: string, + epoch: number +): Promise { + return await ctx.entityManager.registry.materializeActiveClaim({ + consumerId, + epoch, + wakeId: consumerId, + entityUrl: entity.url, + streamPath, + // A webhook notification carries no `lease_ttl_ms` (PROTOCOL §7.1), so + // the row records the subscription default the backend applied. + leaseExpiresAt: new Date(Date.now() + DEFAULT_CLAIM_LEASE_MS), + // The backend fences a callback for a lapsed or re-armed wake itself, + // before this runs; the only way to meet an active row here is a second + // delivery of a wake another runtime is still processing. + unlessActive: true, + }) +} + async function mintClaimWriteToken( ctx: TenantContext, streamPath: string, consumerId: string, - backendToken?: string + backendToken?: string, + // The claim path has already looked the entity up for the durable row; + // pass it in rather than reading it again on the wake's hot path. + knownEntity?: ElectricAgentsEntity | null ): Promise { - const entity = await ctx.entityManager.registry.getEntityByStream(streamPath) + const entity = + knownEntity !== undefined + ? knownEntity + : await ctx.entityManager.registry.getEntityByStream(streamPath) if (!entity) return undefined // 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. + // mint. When the backend supplied none, the store mints its own token — + // from the delivery, if this instance received it — and behaviour is + // byte-for-byte what it is today. (With `fencedSessionStreams` the claim + // path has already refused a claim the backend issued no token for.) // // 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. + // ever appears, so the store mints and nothing changes — including for a + // claim the backend refused, which falls back to a local mint. + // `fencedSessionStreams` is not available against a base backend at all: + // stream creation fails on the missing fencing support + // (`assertWriteFenced`), and a claim it issues no token for is refused + // rather than minted for. // - 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. diff --git a/packages/agents-server/test/consumer-claim-registry.test.ts b/packages/agents-server/test/consumer-claim-registry.test.ts index e41166dd78..277ea3e0d3 100644 --- a/packages/agents-server/test/consumer-claim-registry.test.ts +++ b/packages/agents-server/test/consumer-claim-registry.test.ts @@ -95,3 +95,136 @@ describe(`PostgresRegistry consumer-claim heartbeat (regression for #4341)`, () expect(await readLease(`wake-extend`, 1)).toEqual(extendedLease) }) }) + +describe(`PostgresRegistry claim fencing and post-done status`, () => { + let registry: PostgresRegistry + let db: ReturnType[`db`] + let client: ReturnType[`client`] + + beforeAll(async () => { + await resetElectricAgentsTestBackend() + const connection = createDb(TEST_POSTGRES_URL) + db = connection.db + client = connection.client + registry = new PostgresRegistry(db) + }, 120_000) + + beforeEach(async () => { + await resetElectricAgentsTestBackend() + }, 120_000) + + afterAll(async () => { + await client?.end() + }, 120_000) + + async function createEntity( + url: string, + status: `idle` | `running` | `stopping` + ) { + const now = Date.now() + await registry.createEntity({ + url, + type: `test`, + status, + streams: { main: `${url}/main` }, + subscription_id: `sub-${url}`, + write_token: `wt-${url}`, + tags: {}, + created_at: now, + updated_at: now, + }) + } + + it(`refuses to replace an active claim with unlessActive, and allows it again once released`, async () => { + const claim = { + consumerId: `wake-dup`, + epoch: 3, + entityUrl: `/horton/dup`, + streamPath: `/horton/dup/main`, + unlessActive: true, + } + const firstClaimedAt = new Date(`2026-05-19T10:00:00Z`) + + expect( + await registry.materializeActiveClaim({ + ...claim, + claimedAt: firstClaimedAt, + }) + ).toBe(true) + // The redelivered wake's second claim: same (consumerId, epoch), row + // still active. + expect( + await registry.materializeActiveClaim({ + ...claim, + claimedAt: new Date(`2026-05-19T10:00:05Z`), + }) + ).toBe(false) + const rows = await db + .select() + .from(consumerClaims) + .where( + and( + eq(consumerClaims.consumerId, `wake-dup`), + eq(consumerClaims.epoch, 3) + ) + ) + expect(rows).toHaveLength(1) + expect(rows[0]!.claimedAt).toEqual(firstClaimedAt) + + await registry.materializeReleasedClaim({ + consumerId: `wake-dup`, + epoch: 3, + }) + expect(await registry.materializeActiveClaim(claim)).toBe(true) + }) + + it(`replaces an active claim without unlessActive (pull-wake reclaim)`, async () => { + const claim = { + consumerId: `wake-reclaim`, + epoch: 1, + entityUrl: `/horton/reclaim`, + streamPath: `/horton/reclaim/main`, + } + expect(await registry.materializeActiveClaim(claim)).toBe(true) + expect( + await registry.materializeActiveClaim({ ...claim, runnerId: `runner-2` }) + ).toBe(true) + }) + + it(`settles running to idle and stopping to stopped after done, and leaves other statuses alone`, async () => { + await createEntity(`/horton/running`, `running`) + await createEntity(`/horton/stopping`, `stopping`) + await createEntity(`/horton/idle`, `idle`) + + expect(await registry.updateStatusAfterDone(`/horton/running`)).toBe(`idle`) + expect(await registry.updateStatusAfterDone(`/horton/stopping`)).toBe( + `stopped` + ) + // A done that raced the next wake's delivery finds the entity idle only + // between wakes; a status the done did not cause is not overwritten. + expect(await registry.updateStatusAfterDone(`/horton/idle`)).toBeNull() + + expect((await registry.getEntity(`/horton/running`))?.status).toBe(`idle`) + expect((await registry.getEntity(`/horton/stopping`))?.status).toBe( + `stopped` + ) + expect((await registry.getEntity(`/horton/idle`))?.status).toBe(`idle`) + }) + + it(`reverts only a running entity when the runtime refuses the wake`, async () => { + await createEntity(`/horton/refused`, `running`) + await createEntity(`/horton/refused-stopping`, `stopping`) + + expect(await registry.revertRunningStatus(`/horton/refused`)).toBe(true) + // A wake refused while the entity is stopping (a rolling runtime restart + // refuses every wake it is handed) must not report the entity idle. + expect(await registry.revertRunningStatus(`/horton/refused-stopping`)).toBe( + false + ) + + expect((await registry.getEntity(`/horton/refused`))?.status).toBe(`idle`) + expect((await registry.getEntity(`/horton/refused-stopping`))?.status).toBe( + `stopping` + ) + }) +}) 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 662522e589..3248412b2b 100644 --- a/packages/agents-server/test/server-claim-write-token.test.ts +++ b/packages/agents-server/test/server-claim-write-token.test.ts @@ -729,14 +729,15 @@ describe(`Claim-scoped write tokens`, () => { await registry.updateStatus(entity.url, `running`) expect(await getEntityStatus(entity.url)).toBe(`running`) - const origUpdateStatus = registry.updateStatus.bind(registry) + const origUpdateStatusAfterDone = + registry.updateStatusAfterDone.bind(registry) let shouldFail = true - registry.updateStatus = async (...args: [string, string]) => { + registry.updateStatusAfterDone = async (entityUrl: string) => { if (shouldFail) { shouldFail = false throw new Error(`simulated DB failure`) } - return origUpdateStatus(...args) + return origUpdateStatusAfterDone(entityUrl) } const firstDone = await sendDone({ @@ -747,7 +748,7 @@ describe(`Claim-scoped write tokens`, () => { expect(firstDone.status).toBe(200) expect(await getEntityStatus(entity.url)).toBe(`running`) - registry.updateStatus = origUpdateStatus + registry.updateStatusAfterDone = origUpdateStatusAfterDone const retryDone = await sendDone({ consumerId: `consumer-done-retry`, diff --git a/packages/agents-server/test/subscription-webhooks-routing.test.ts b/packages/agents-server/test/subscription-webhooks-routing.test.ts index 541264ae6d..be8116e7d2 100644 --- a/packages/agents-server/test/subscription-webhooks-routing.test.ts +++ b/packages/agents-server/test/subscription-webhooks-routing.test.ts @@ -129,12 +129,16 @@ function buildContext(overrides: Partial = {}): TenantContext { registry: { getEntityByStream: vi.fn().mockResolvedValue(entity), updateStatus: vi.fn().mockResolvedValue(undefined), + updateStatusAfterDone: vi.fn().mockResolvedValue(`idle`), + revertRunningStatus: vi.fn().mockResolvedValue(true), + materializeActiveClaim: vi.fn().mockResolvedValue(true), materializeReleasedClaim: vi.fn().mockResolvedValue({ claim: null, entityCleared: true, }), materializeHeartbeatClaim: vi.fn().mockResolvedValue(undefined), }, + fencedSessionStreams: false, enrichPayload: vi.fn(async (payload: Record) => ({ ...payload, entity: { @@ -593,57 +597,470 @@ describe(`subscription webhooks for Durable Streams subscriptions`, () => { } }) - it(`claims new webhook wakes locally and returns a tenant-scoped claim write token`, async () => { + describe(`wake-callback claims round-trip to the backend`, () => { + const callbackUrl = `http://durable.local/v1/stream/tenant-a/__ds/subscriptions/horton-handler/callback` + + function claimContext( + overrides: Partial = {} + ): TenantContext { + const select = selectDb([ + { callbackUrl, primaryStream: `/horton/demo/main` }, + ]) + return buildContext({ + pgDb: { select: select.select } as any, + ...overrides, + }) + } + + function claimRequest(): Request { + return new Request( + `http://agents.local/_electric/wake-callbacks/wake-1`, + { + method: `POST`, + headers: { + 'content-type': `application/json`, + authorization: `Bearer callback-token`, + }, + body: JSON.stringify({ wakeId: `wake-1`, epoch: 7 }), + } + ) + } + + function backendAnswers(body: Record, status = 200) { + return vi.spyOn(globalThis, `fetch`).mockResolvedValue( + new Response(JSON.stringify(body), { + status, + headers: { 'content-type': `application/json` }, + }) + ) + } + + it(`forwards the claim as a lease-renewing callback and mints a claim write token when the backend issues none`, async () => { + const fetchSpy = backendAnswers({ ok: true, next_wake: false }) + const ctx = claimContext() + + try { + const response = await globalRouter.fetch(claimRequest(), ctx) + + expect(response.status).toBe(200) + const body = await responseJson(response) + expect(body.ok).toBe(true) + expect(body.writeToken).toEqual(expect.any(String)) + expect(body.writeToken).not.toBe(`write-token`) + expect( + ctx.runtime.claimWriteTokens.isValid( + `tenant-a`, + `/horton/demo/main`, + body.writeToken + ) + ).toBe(true) + + const [url, init] = fetchSpy.mock.calls[0]! + expect(String(url)).toBe(callbackUrl) + expect((init?.headers as Headers).get(`authorization`)).toBe( + `Bearer callback-token` + ) + expect(requestBodyJson(init?.body)).toEqual({ + wake_id: `wake-1`, + generation: 7, + acks: [], + }) + } finally { + fetchSpy.mockRestore() + } + }) + + it(`adopts the write token the backend re-mints on the claim callback over the one it delivered`, async () => { + const fetchSpy = backendAnswers({ + ok: true, + next_wake: false, + write_token: `backend-token-on-claim`, + }) + const ctx = claimContext() + // Another server instance received the delivery, so this one holds + // no delivered token; the backend's answer is what carries it here. + ctx.runtime.claimWriteTokens.recordDelivered( + `tenant-a`, + `wake-1`, + `backend-token-on-delivery` + ) + + try { + const response = await globalRouter.fetch(claimRequest(), ctx) + + expect(response.status).toBe(200) + const body = await responseJson(response) + expect(body.writeToken).toBe(`backend-token-on-claim`) + expect( + ctx.runtime.claimWriteTokens.isValid( + `tenant-a`, + `/horton/demo/main`, + `backend-token-on-claim` + ) + ).toBe(true) + } finally { + fetchSpy.mockRestore() + } + }) + + it(`refuses a fenced claim the backend issued no write token for, without minting one`, async () => { + const fetchSpy = backendAnswers({ ok: true, next_wake: false }) + const ctx = claimContext() + ;(ctx.entityManager as any).fencedSessionStreams = true + + try { + const response = await globalRouter.fetch(claimRequest(), ctx) + + expect(response.status).toBe(502) + await expect(responseJson(response)).resolves.toMatchObject({ + error: { code: `WRITE_TOKEN_UNAVAILABLE` }, + }) + expect( + ctx.runtime.claimWriteTokens.owns( + `tenant-a`, + `/horton/demo/main`, + `wake-1` + ) + ).toBe(false) + expect( + ctx.entityManager.registry.materializeActiveClaim + ).not.toHaveBeenCalled() + } finally { + fetchSpy.mockRestore() + } + }) + + it(`materializes the durable claim once the backend has accepted it`, async () => { + const fetchSpy = backendAnswers({ ok: true, next_wake: false }) + const ctx = claimContext() + + try { + const before = Date.now() + const response = await globalRouter.fetch(claimRequest(), ctx) + + expect(response.status).toBe(200) + expect( + ctx.entityManager.registry.materializeActiveClaim + ).toHaveBeenCalledWith({ + consumerId: `wake-1`, + epoch: 7, + wakeId: `wake-1`, + entityUrl: `/horton/demo`, + streamPath: `/horton/demo/main`, + leaseExpiresAt: expect.any(Date), + unlessActive: true, + }) + const { leaseExpiresAt } = vi.mocked( + ctx.entityManager.registry.materializeActiveClaim + ).mock.calls[0]![0] + expect(leaseExpiresAt!.getTime()).toBeGreaterThanOrEqual( + before + 30_000 + ) + } finally { + fetchSpy.mockRestore() + } + }) + + it(`does not materialize a fenced claim the backend rejected`, async () => { + const fetchSpy = backendAnswers( + { error: { code: `FENCED`, message: `stale generation` } }, + 409 + ) + const ctx = claimContext() + ;(ctx.entityManager as any).fencedSessionStreams = true + + try { + const response = await globalRouter.fetch(claimRequest(), ctx) + + expect(response.status).toBe(409) + expect( + ctx.entityManager.registry.materializeActiveClaim + ).not.toHaveBeenCalled() + expect( + ctx.runtime.claimWriteTokens.owns( + `tenant-a`, + `/horton/demo/main`, + `wake-1` + ) + ).toBe(false) + } finally { + fetchSpy.mockRestore() + } + }) + + it(`mints locally when the backend rejects a claim on an unfenced stream`, async () => { + // Without fencing the store is the only write authority, so the round + // trip is an optimisation, not a gate: answering the runtime with the + // backend's error would drop a wake that used to run. + const fetchSpy = backendAnswers( + { error: { code: `FENCED`, message: `stale generation` } }, + 409 + ) + const ctx = claimContext() + + try { + const response = await globalRouter.fetch(claimRequest(), ctx) + + expect(response.status).toBe(200) + const body = await responseJson(response) + expect(body.ok).toBe(true) + expect( + ctx.runtime.claimWriteTokens.isValid( + `tenant-a`, + `/horton/demo/main`, + body.writeToken + ) + ).toBe(true) + expect( + ctx.entityManager.registry.materializeActiveClaim + ).not.toHaveBeenCalled() + } finally { + fetchSpy.mockRestore() + } + }) + + it(`mints locally when the claim never reaches an unfenced backend`, async () => { + const fetchSpy = vi + .spyOn(globalThis, `fetch`) + .mockRejectedValue(new Error(`connect ECONNREFUSED`)) + const ctx = claimContext() + + try { + const response = await globalRouter.fetch(claimRequest(), ctx) + + expect(response.status).toBe(200) + const body = await responseJson(response) + expect( + ctx.runtime.claimWriteTokens.isValid( + `tenant-a`, + `/horton/demo/main`, + body.writeToken + ) + ).toBe(true) + } finally { + fetchSpy.mockRestore() + } + }) + + it(`releases a done that carries its wake id instead of answering it as a claim`, async () => { + // PROTOCOL §7.1 puts `wake_id` on every callback, done included; a done + // routed down the claim path would be refused for the write token no + // backend mints for it, and the claim would never be released. + const fetchSpy = backendAnswers({ ok: true, next_wake: false }) + const ctx = claimContext() + ;(ctx.entityManager as any).fencedSessionStreams = true + + 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({ + wake_id: `wake-1`, + generation: 7, + acks: [{ path: `/horton/demo/main`, offset: `1` }], + done: true, + }), + }), + ctx + ) + + expect(response.status).toBe(200) + expect( + ctx.entityManager.registry.materializeReleasedClaim + ).toHaveBeenCalledWith( + expect.objectContaining({ consumerId: `wake-1`, epoch: 7 }) + ) + expect( + ctx.entityManager.registry.updateStatusAfterDone + ).toHaveBeenCalledWith(`/horton/demo`) + expect( + ctx.entityManager.registry.materializeActiveClaim + ).not.toHaveBeenCalled() + } finally { + fetchSpy.mockRestore() + } + }) + + it(`drops a redelivered wake's second claim with 409 before minting a token`, async () => { + const fetchSpy = backendAnswers({ ok: true, next_wake: false }) + const ctx = claimContext() + vi.mocked( + ctx.entityManager.registry.materializeActiveClaim + ).mockResolvedValue(false) + + try { + const response = await globalRouter.fetch(claimRequest(), ctx) + + expect(response.status).toBe(409) + await expect(responseJson(response)).resolves.toMatchObject({ + error: { code: `WAKE_ALREADY_CLAIMED` }, + }) + expect( + ctx.runtime.claimWriteTokens.owns( + `tenant-a`, + `/horton/demo/main`, + `wake-1` + ) + ).toBe(false) + } finally { + fetchSpy.mockRestore() + } + }) + + it(`leaves pull-wake claims to the runner claim that already materialized them`, async () => { + const select = selectDb([ + { + callbackUrl: `ds-subscription:horton-handler`, + primaryStream: `/horton/demo/main`, + }, + ]) + const ackSubscription = vi.fn().mockResolvedValue({ + ok: true, + next_wake: false, + write_token: `backend-token-on-claim`, + }) + const ctx = buildContext({ + pgDb: { select: select.select } as any, + streamClient: { ackSubscription } as any, + }) + + const response = await globalRouter.fetch(claimRequest(), ctx) + + expect(response.status).toBe(200) + expect(ackSubscription).toHaveBeenCalledWith( + `horton-handler`, + `callback-token`, + { wake_id: `wake-1`, generation: 7, acks: [] } + ) + await expect(responseJson(response)).resolves.toMatchObject({ + writeToken: `backend-token-on-claim`, + }) + expect( + ctx.entityManager.registry.materializeActiveClaim + ).not.toHaveBeenCalled() + }) + }) + + it(`renews the in-memory claim on heartbeats the backend answers without a 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`) - + const fetchSpy = vi.spyOn(globalThis, `fetch`).mockResolvedValue( + new Response(JSON.stringify({ ok: true }), { + headers: { 'content-type': `application/json` }, + }) + ) + const clock = { now: 0 } + const claimWriteTokens = new ClaimWriteTokenStore({ + ttlMs: 1_000, + now: () => clock.now, + }) const ctx = buildContext({ pgDb: { select: select.select } as any, + runtime: { claimWriteTokens } as any, }) + const activeToken = claimWriteTokens.mint( + `tenant-a`, + `/horton/demo/main`, + `wake-1` + ) try { + clock.now = 900 const response = await globalRouter.fetch( - request(`POST`, `/_electric/wake-callbacks/wake-1`, { - wakeId: `wake-1`, - epoch: 7, + 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: [] }), }), ctx ) expect(response.status).toBe(200) - const body = await responseJson(response) - expect(body.ok).toBe(true) - expect(body.writeToken).toEqual(expect.any(String)) - expect(body.writeToken).not.toBe(`write-token`) + clock.now = 1_500 expect( - ctx.runtime.claimWriteTokens.isValid( - `tenant-a`, - `/horton/demo/main`, - body.writeToken - ) + claimWriteTokens.isValid(`tenant-a`, `/horton/demo/main`, activeToken) ).toBe(true) - expect(fetchSpy).not.toHaveBeenCalled() } finally { fetchSpy.mockRestore() } }) - it(`adopts the write token a backend delivers with a webhook wake`, async () => { - const claimWriteTokens = new ClaimWriteTokenStore() - const webhookSelect = selectDb([ + it(`resets the entity to idle when the runtime refuses the forwarded wake`, async () => { + const select = 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 }), { + new Response(JSON.stringify({ error: `draining` }), { + status: 503, headers: { 'content-type': `application/json` }, }) ) + const ctx = buildContext({ + pgDb: { select: select.select, insert: insert.insert } as any, + }) + + try { + const response = await globalRouter.fetch( + request(`POST`, `/_electric/subscription-webhooks/horton-handler`, { + subscription_id: `horton-handler`, + wake_id: `wake-refused`, + generation: 7, + 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`, + }), + ctx + ) + + expect(response.status).toBe(503) + expect( + ctx.entityManager.registry.updateStatus + ).toHaveBeenCalledExactlyOnceWith(`/horton/demo`, `running`) + // Conditional: a rolling runtime restart refuses wakes in bulk, and an + // unconditional `idle` would report a stopping entity — or one a + // sibling replica has already woken again — wrong. + expect( + ctx.entityManager.registry.revertRunningStatus + ).toHaveBeenCalledExactlyOnceWith(`/horton/demo`) + } finally { + fetchSpy.mockRestore() + } + }) + + it(`adopts the write token a backend delivers with a webhook wake when the claim callback returns none`, async () => { + const claimWriteTokens = new ClaimWriteTokenStore() + const webhookSelect = selectDb([ + { webhookUrl: `http://runtime.local/_electric/builtin-agent-handler` }, + ]) + const insert = insertDb() + // Answers both the forward to the runtime and the claim's round trip to + // the backend; the latter carries no write_token, as a backend without + // the Write Fencing extension answers. + const fetchSpy = vi.spyOn(globalThis, `fetch`).mockImplementation( + async () => + new Response(JSON.stringify({ ok: true }), { + headers: { 'content-type': `application/json` }, + }) + ) try { const webhookResponse = await globalRouter.fetch( @@ -1047,10 +1464,9 @@ describe(`subscription webhooks for Durable Streams subscriptions`, () => { acks: [{ stream: `horton/demo/main`, offset: `1` }], done: true, }) - expect(ctx.entityManager.registry.updateStatus).toHaveBeenCalledWith( - `/horton/demo`, - `idle` - ) + expect( + ctx.entityManager.registry.updateStatusAfterDone + ).toHaveBeenCalledWith(`/horton/demo`) expect(ctx.entityBridgeManager.onEntityChanged).toHaveBeenCalledWith( `/horton/demo` ) @@ -1165,10 +1581,9 @@ describe(`subscription webhooks for Durable Streams subscriptions`, () => { ackedStreams: [{ path: `/horton/demo/main`, offset: `1` }], }) ) - expect(ctx.entityManager.registry.updateStatus).toHaveBeenCalledWith( - `/horton/demo`, - `idle` - ) + expect( + ctx.entityManager.registry.updateStatusAfterDone + ).toHaveBeenCalledWith(`/horton/demo`) } finally { vi.mocked(globalThis.fetch).mockRestore() } @@ -1212,10 +1627,9 @@ describe(`subscription webhooks for Durable Streams subscriptions`, () => { epoch: 7, }) ) - expect(ctx.entityManager.registry.updateStatus).toHaveBeenCalledWith( - `/horton/demo`, - `idle` - ) + expect( + ctx.entityManager.registry.updateStatusAfterDone + ).toHaveBeenCalledWith(`/horton/demo`) } finally { vi.mocked(globalThis.fetch).mockRestore() } @@ -1292,10 +1706,12 @@ describe(`subscription webhooks for Durable Streams subscriptions`, () => { expect( ctx.entityManager.registry.getEntityByStream ).toHaveBeenCalledWith(`/horton/demo-a/main`) - expect(ctx.entityManager.registry.updateStatus).toHaveBeenCalledWith( - `/horton/demo-a`, - `idle` - ) + expect( + ctx.entityManager.registry.updateStatusAfterDone + ).toHaveBeenCalledWith(`/horton/demo-a`) + // The done finishes the consumer, so every claim it holds is + // released — a consumer id names one wake, and one wake cannot + // outlive its own done on another stream. expect( ctx.runtime.claimWriteTokens.owns( `tenant-a`, @@ -1309,7 +1725,7 @@ describe(`subscription webhooks for Durable Streams subscriptions`, () => { `/horton/demo-b/main`, `wake-1` ) - ).toBe(true) + ).toBe(false) } finally { vi.mocked(globalThis.fetch).mockRestore() } @@ -1370,7 +1786,9 @@ describe(`subscription webhooks for Durable Streams subscriptions`, () => { ) // The newer consumer owns the entity now — its status must stay // as-is and its in-memory write token must remain intact. - expect(ctx.entityManager.registry.updateStatus).not.toHaveBeenCalled() + expect( + ctx.entityManager.registry.updateStatusAfterDone + ).not.toHaveBeenCalled() expect( ctx.runtime.claimWriteTokens.owns( `tenant-a`, @@ -1382,5 +1800,150 @@ describe(`subscription webhooks for Durable Streams subscriptions`, () => { vi.mocked(globalThis.fetch).mockRestore() } }) + + it(`settles the entity status on a server instance that did not mint the claim`, async () => { + // Two server instances share the registry (Postgres) but each has its + // own ClaimWriteTokenStore. The claim lands on one, the done on the + // other; the durable release, not process memory, must settle status. + const entity = { + url: `/horton/demo`, + type: `horton`, + status: `running`, + streams: { main: `/horton/demo/main` }, + tags: {}, + spawn_args: {}, + write_token: `write-token`, + } + const claimRow = { + consumer_id: `wake-1`, + epoch: 7, + entity_url: `/horton/demo`, + stream_path: `/horton/demo/main`, + } + const registry = { + getEntityByStream: vi.fn().mockResolvedValue(entity), + updateStatus: vi.fn().mockResolvedValue(undefined), + updateStatusAfterDone: vi.fn(async () => { + entity.status = entity.status === `stopping` ? `stopped` : `idle` + return entity.status + }), + materializeActiveClaim: vi.fn().mockResolvedValue(true), + materializeHeartbeatClaim: vi.fn().mockResolvedValue(undefined), + materializeReleasedClaim: vi.fn().mockResolvedValue({ + claim: claimRow, + entityCleared: true, + }), + } + const instance = (): TenantContext => { + const select = selectDb([ + { + callbackUrl: `http://durable.local/v1/stream/tenant-a/__ds/subscriptions/horton-handler/callback`, + primaryStream: `/horton/demo/main`, + }, + ]) + const ctx = buildContext({ pgDb: { select: select.select } as any }) + ;(ctx.entityManager as any).registry = registry + return ctx + } + const claimingInstance = instance() + const doneInstance = instance() + vi.spyOn(globalThis, `fetch`).mockImplementation( + async () => + new Response(JSON.stringify({ ok: true, next_wake: false }), { + headers: { 'content-type': `application/json` }, + }) + ) + + try { + const claimResponse = await globalRouter.fetch( + request(`POST`, `/_electric/wake-callbacks/wake-1`, { + wakeId: `wake-1`, + epoch: 7, + }), + claimingInstance + ) + expect(claimResponse.status).toBe(200) + expect( + claimingInstance.runtime.claimWriteTokens.owns( + `tenant-a`, + `/horton/demo/main`, + `wake-1` + ) + ).toBe(true) + + const doneResponse = 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` }], + done: true, + }), + }), + doneInstance + ) + + expect(doneResponse.status).toBe(200) + expect(registry.materializeReleasedClaim).toHaveBeenCalledWith( + expect.objectContaining({ consumerId: `wake-1`, epoch: 7 }) + ) + expect(registry.updateStatusAfterDone).toHaveBeenCalledWith( + `/horton/demo` + ) + expect(entity.status).toBe(`idle`) + expect( + doneInstance.entityBridgeManager.onEntityChanged + ).toHaveBeenCalledWith(`/horton/demo`) + } finally { + vi.mocked(globalThis.fetch).mockRestore() + } + }) + + it(`drops the delivered token the instance held for a claim another instance minted`, async () => { + const select = selectDb([ + { + callbackUrl: `http://durable.local/v1/stream-meta/subscriptions/horton-handler/callback?opaque=tenant-a`, + primaryStream: `/horton/demo/main`, + }, + ]) + upstreamOk() + const ctx = buildContext({ + pgDb: { select: select.select } as any, + }) + ctx.runtime.claimWriteTokens.recordDelivered( + `tenant-a`, + `wake-1`, + `backend-token-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` }], + done: true, + }), + }), + ctx + ) + + expect(response.status).toBe(200) + expect( + ctx.runtime.claimWriteTokens.takeDelivered(`tenant-a`, `wake-1`) + ).toBeUndefined() + } finally { + vi.mocked(globalThis.fetch).mockRestore() + } + }) }) }) From 38b9d260be1fd3e9fd5ecb2e650ecbb1b5e18d58 Mon Sep 17 00:00:00 2001 From: Aditya Kumarakrishnan Date: Wed, 2 Sep 2026 15:11:55 +0530 Subject: [PATCH 5/6] feat(agents-runtime): leave a wake pending when its failure never reached the stream MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A wake whose appends fail (a deposed or lapsed write token, a lost connection) records WRITE_FAILED and appends an error event — through the same producer, so when the failure is a dead token the error event is lost too. The done then still acked safeAckOffset, and the backend consumed the triggering message with nothing on the stream to say why no run answered it. When a producer batch fails after the error event was appended and no append has landed since, the done now carries no acks: PROTOCOL §7.1 idles the claim without moving the cursor, and the pending work re-wakes the entity — the same shape as a crash mid-wake, and a duplicate run rather than a silent loss. Batches are sent concurrently, so a batch queued before the failure can report its own error afterwards; the producer's last written offset is what separates that from the error event itself being lost, and a wake whose error event did land keeps acking as before. The ack-less done is logged as a warning, because it is a full re-run of the handler and a write failure that is not clearing turns it into a wake loop. Two smaller guards on the same path: a redelivery of a wake already in flight for the same stream and generation (the backend's retry of a delivery whose 2xx it never saw) is acknowledged without starting a second wake, which would otherwise claim the same generation — the backend accepts that while the lease is live — and run the handler twice; and a `write_token` delivered with the wake notification is adopted when the claim callback returns none, for a server that passes the field through without adopting it. --- ...agents-runtime-unrecorded-write-failure.md | 5 + packages/agents-runtime/src/create-handler.ts | 36 ++++- packages/agents-runtime/src/process-wake.ts | 46 ++++++- packages/agents-runtime/src/types.ts | 7 + .../test/create-handler.test.ts | 61 ++++++++ .../agents-runtime/test/process-wake.test.ts | 130 ++++++++++++++++++ 6 files changed, 274 insertions(+), 11 deletions(-) create mode 100644 .changeset/agents-runtime-unrecorded-write-failure.md diff --git a/.changeset/agents-runtime-unrecorded-write-failure.md b/.changeset/agents-runtime-unrecorded-write-failure.md new file mode 100644 index 0000000000..1f8ce0ccc7 --- /dev/null +++ b/.changeset/agents-runtime-unrecorded-write-failure.md @@ -0,0 +1,5 @@ +--- +'@electric-ax/agents-runtime': patch +--- + +Stop a wake whose failure never reached its stream from consuming its trigger: when the wake's writes failed and the error event recording that failure was itself lost with the producer, the done is sent without acks so the backend re-wakes the entity instead of marking a message answered that no run answered. The runtime also ignores a redelivery of a wake already in flight for the same stream and generation (the backend's retry of a delivery whose 2xx it never saw), and adopts a `write_token` delivered with the wake notification when the claim callback returns none. diff --git a/packages/agents-runtime/src/create-handler.ts b/packages/agents-runtime/src/create-handler.ts index 7bc8069ec5..b8807f345e 100644 --- a/packages/agents-runtime/src/create-handler.ts +++ b/packages/agents-runtime/src/create-handler.ts @@ -148,8 +148,11 @@ export interface RuntimeRouter { options?: Pick ) => void - /** True when a wake for the stream path is already in flight. */ - isWakeActive: (streamPath: string) => boolean + /** + * True when a wake for the stream path is already in flight — with + * `epoch`, only one for that generation. + */ + isWakeActive: (streamPath: string, epoch?: number) => boolean /** Dispatch an already-parsed webhook wake notification. */ dispatchWebhookWake: (notification: WebhookNotification) => void @@ -372,6 +375,7 @@ export function createRuntimeRouter( process.env.ELECTRIC_AGENTS_DEBUG_REGISTRATION_TIMING === `1` const pendingWakes = new Set>() const pendingWakeLabels = new Map, string>() + const pendingWakeEpochs = new Map, number>() const pendingWakeControllers = new Map, AbortController>() const wakeErrors: Array = [] const debugCleanup = process.env.ELECTRIC_AGENTS_DEBUG_CLEANUP === `1` @@ -434,15 +438,26 @@ export function createRuntimeRouter( .finally(() => { pendingWakes.delete(wake) pendingWakeLabels.delete(wake) + pendingWakeEpochs.delete(wake) pendingWakeControllers.delete(wake) }) pendingWakes.add(wake) pendingWakeLabels.set(wake, wakeLabel) + pendingWakeEpochs.set(wake, notification.epoch) pendingWakeControllers.set(wake, controller) } - const isWakeActive: RuntimeRouter[`isWakeActive`] = (streamPath) => - [...pendingWakeLabels.values()].includes(streamPath) + const isWakeActive: RuntimeRouter[`isWakeActive`] = (streamPath, epoch) => { + for (const [wake, label] of pendingWakeLabels) { + if ( + label === streamPath && + (epoch === undefined || pendingWakeEpochs.get(wake) === epoch) + ) { + return true + } + } + return false + } const dispatchWebhookWake: RuntimeRouter[`dispatchWebhookWake`] = dispatchWake @@ -524,6 +539,19 @@ export function createRuntimeRouter( ) } + if (isWakeActive(notification.streamPath, notification.epoch)) { + // The backend redelivers a wake whose 2xx it never saw; a second wake + // for the same generation would claim it again (the backend accepts + // that — the lease is live) and run the handler twice. The in-flight + // wake's done settles the delivery, so acknowledge without starting + // another. + runtimeLog.warn( + `[agent-runtime]`, + `wake for ${notification.streamPath} (epoch=${notification.epoch}) is already in flight; ignoring redelivery` + ) + return json({ ok: true }, 200) + } + dispatchWebhookWake(notification) return json({ ok: true }, 200) } diff --git a/packages/agents-runtime/src/process-wake.ts b/packages/agents-runtime/src/process-wake.ts index cf8cf6992d..f68d334c2f 100644 --- a/packages/agents-runtime/src/process-wake.ts +++ b/packages/agents-runtime/src/process-wake.ts @@ -548,6 +548,11 @@ export async function processWake( return globalThis.fetch(input, { ...init, headers }) }, onError: (error) => { + // A batch that fails once this wake's failure is on record may be the + // one carrying the error event failBackgroundWake appended; whether it + // was is settled at done time against the producer's last written + // offset. See errorEventUnrecorded. + if (liveProcessError) writeFailedAfterErrorEvent = true failBackgroundWake(error, `WRITE_FAILED`) }, }) @@ -634,6 +639,16 @@ export async function processWake( close: () => void }> = [] let liveProcessError: Error | null = null + // Set when a producer batch fails after failBackgroundWake appended this + // wake's error event — possibly the batch carrying it, since batches are + // sent concurrently. `producer.lastSuccessfulOffset` as it stood when the + // event was appended decides which: an append that landed after it proves + // the error event landed too. When it did not, nothing durable records that + // the wake failed, so the done carries no acks — the trigger stays pending + // and the backend re-wakes the entity instead of consuming a message no run + // answered, the same shape as a crash mid-wake. + let writeFailedAfterErrorEvent = false + let offsetBeforeErrorEvent: string | undefined let acceptLiveInputs = false const handledSignalKeys = new Set() @@ -701,6 +716,7 @@ export async function processWake( liveProcessError = toError(err) log.error(`wake background task failed for ${entityUrl}:`, liveProcessError) + offsetBeforeErrorEvent = producer.lastSuccessfulOffset writeEvent( entityStateSchema.errors.insert({ key: `error-${epoch}-${crypto.randomUUID()}`, @@ -1222,7 +1238,7 @@ export async function processWake( return null } claimedWake = true - writeToken = claimed.writeToken ?? `` + writeToken = claimed.writeToken ?? notification.write_token ?? `` handleRuntimeSideEffectEvents(catchUpEvents) @@ -2561,7 +2577,12 @@ export async function processWake( cleanupErrors.push(toError(err)) } } - const doneOffset = safeAckOffset + // Every batch has been flushed by now, so an offset past the one the + // error event was queued behind is an append that succeeded after it. + const errorEventUnrecorded = + writeFailedAfterErrorEvent && + producer.lastSuccessfulOffset === offsetBeforeErrorEvent + const doneOffset = errorEventUnrecorded ? `-1` : safeAckOffset for (const sdb of secondaryDbs) { try { await sdb.flushWrites?.() @@ -2600,11 +2621,22 @@ export async function processWake( } } if (claimedWake) { - log.info( - doneOffset === `-1` - ? `done without ack (no consumed offset)` - : `done acking ${streamPath} at ${doneOffset}` - ) + if (errorEventUnrecorded) { + // Warn, not info: the trigger stays pending, so the backend re-wakes + // this entity immediately and the whole handler — model and tool + // calls included — runs again. One line is a duplicated run; a stream + // of them for the same entity is a wake loop over a write failure + // that is not clearing, and the writes are what to fix. + log.warn( + `done without ack for ${streamPath} (epoch=${epoch}): the write failure was never recorded on the stream, leaving the wake pending for redelivery` + ) + } else { + log.info( + doneOffset === `-1` + ? `done without ack (no consumed offset)` + : `done acking ${streamPath} at ${doneOffset}` + ) + } if (shutdownRequested) { log.info(`shutdown requested, sending done callback at checkpoint`) } diff --git a/packages/agents-runtime/src/types.ts b/packages/agents-runtime/src/types.ts index 8cee5d4722..0706d304ec 100644 --- a/packages/agents-runtime/src/types.ts +++ b/packages/agents-runtime/src/types.ts @@ -704,6 +704,13 @@ export interface WebhookNotification { triggeredBy?: Array callback: string claimToken: string + /** + * The wake's write token as the Durable Streams backend delivered it + * (Write Fencing extension), passed through by the server. The claim + * callback's `writeToken` is authoritative; this is the fallback for a + * server that adopts none. + */ + write_token?: string triggerEvent?: string wakeEvent?: WakeEvent entity?: { diff --git a/packages/agents-runtime/test/create-handler.test.ts b/packages/agents-runtime/test/create-handler.test.ts index 81432345a3..a4689ef95a 100644 --- a/packages/agents-runtime/test/create-handler.test.ts +++ b/packages/agents-runtime/test/create-handler.test.ts @@ -244,6 +244,67 @@ describe(`createRuntimeHandler`, () => { }) }) + it(`acknowledges a redelivered wake for an in-flight (stream, epoch) without starting it again`, async () => { + defineEntity(`test-agent`, { handler: async () => {} }) + + const resolvers: Array<() => void> = [] + processWakeMock.mockImplementation( + () => + new Promise((resolve) => { + resolvers.push(resolve) + }) + ) + + const notification = { + consumerId: `wake-1`, + epoch: 1, + wakeId: `wake-1`, + streamPath: `/streams/entity:test-1`, + streams: [{ path: `/streams/entity:test-1`, offset: `0_0` }], + callback: `http://localhost:3000/_electric/wakes/wake-1`, + claimToken: `tok-1`, + entity: { + type: `test-agent`, + status: `active`, + url: `http://localhost:3000/test-agent/test-1`, + streams: { + main: `/streams/entity:test-1`, + }, + }, + } + + const handler = createRuntimeHandler({ + baseUrl: `http://localhost:3000`, + handlerUrl: `http://localhost:4000/electric-agents`, + webhookSignature: false, + }) + const deliver = (body: unknown): Promise => + handler.handleWebhookRequest( + new Request(`http://localhost/electric-agents`, { + method: `POST`, + headers: { 'content-type': `application/json` }, + body: JSON.stringify(body), + }) + ) + + expect((await deliver(notification)).status).toBe(200) + // The backend's retry of the same wake (its 2xx was lost) must be + // acknowledged — the in-flight wake's done settles it — not run twice. + expect((await deliver(notification)).status).toBe(200) + expect(processWakeMock).toHaveBeenCalledTimes(1) + + // A later generation for the same stream is a new wake, not a redelivery. + expect( + (await deliver({ ...notification, epoch: 2, wakeId: `wake-2` })).status + ).toBe(200) + expect(processWakeMock).toHaveBeenCalledTimes(2) + expect(handler.debugState()).toMatchObject({ pendingWakeCount: 2 }) + + for (const resolve of resolvers) resolve() + await handler.waitForSettled() + expect(handler.isWakeActive(`/streams/entity:test-1`, 1)).toBe(false) + }) + it(`records wake errors in debugState() until drained`, async () => { defineEntity(`test-agent`, { handler: async () => {} }) processWakeMock.mockRejectedValueOnce(new Error(`wake failed`)) diff --git a/packages/agents-runtime/test/process-wake.test.ts b/packages/agents-runtime/test/process-wake.test.ts index 0c09cbc8a2..9e04598efd 100644 --- a/packages/agents-runtime/test/process-wake.test.ts +++ b/packages/agents-runtime/test/process-wake.test.ts @@ -29,6 +29,7 @@ const { mockProducerAppend, mockProducerFlush, mockProducerDetach, + mockProducerLastSuccessfulOffset, mockConstructedProducers, mockDbClose, mockDbPreload, @@ -50,6 +51,9 @@ const { mockProducerAppend: vi.fn(), mockProducerFlush: vi.fn().mockResolvedValue(undefined), mockProducerDetach: vi.fn().mockResolvedValue(undefined), + mockProducerLastSuccessfulOffset: { + value: undefined as string | undefined, + }, mockConstructedProducers: [] as Array<{ producerId: string opts?: Record @@ -131,6 +135,10 @@ vi.mock(`@durable-streams/client`, async (importOriginal) => { append = mockProducerAppend flush = mockProducerFlush detach = mockProducerDetach + + get lastSuccessfulOffset(): string | undefined { + return mockProducerLastSuccessfulOffset.value + } } return { ...actual, @@ -439,6 +447,7 @@ describe(`processWake`, () => { vi.useRealTimers() clearRegistry() mockConstructedProducers.length = 0 + mockProducerLastSuccessfulOffset.value = undefined mockDbPreload.mockResolvedValue(undefined) mockSourceDbPreload.mockResolvedValue(undefined) mockCreateStreamDB.mockImplementation((options) => { @@ -1900,6 +1909,127 @@ describe(`processWake`, () => { ) }) + it(`falls back to the notification's write_token when the claim callback returns none`, async () => { + defineEntity(`test-agent`, { + handler: () => {}, + }) + + await processWake( + makeNotification({ write_token: `delivered-token` }), + BASE_CONFIG + ) + + 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 delivered-token` + ) + }) + + describe(`done after a write failure`, () => { + const entityProducerOnError = (): ((error: Error) => void) => { + const producer = mockConstructedProducers.find( + (constructed) => + constructed.producerId === + `entity-http://localhost:3000/test-agent/agent-1` + ) + return producer!.opts!.onError as (error: Error) => void + } + + const doneBody = (): { acks: Array; done: boolean } => { + const doneCalls = fetchMock.mock.calls.filter( + ([url, opts]) => + String(url).includes(`/_electric/wakes/wake-abc`) && + String((opts as RequestInit | undefined)?.body).includes( + `"done":true` + ) + ) + expect(doneCalls).toHaveLength(1) + return JSON.parse(doneCalls[0]![1]!.body as string) + } + + it(`acks the consumed offset when the error event reached the stream`, async () => { + defineEntity(`test-agent`, { + handler: () => { + entityProducerOnError()(new Error(`append rejected (401)`)) + }, + }) + + await expect( + processWake(makeNotification(), BASE_CONFIG) + ).rejects.toThrow(`append rejected (401)`) + + expect(mockProducerAppend).toHaveBeenCalledWith( + expect.stringContaining(`WRITE_FAILED`) + ) + const body = doneBody() + expect(body.done).toBe(true) + expect(body.acks).toEqual([ + { path: `/streams/entity:agent-1`, offset: expect.any(String) }, + ]) + }) + + it(`acks when a later batch failure was not the error event's`, async () => { + defineEntity(`test-agent`, { + handler: () => { + const onError = entityProducerOnError() + // One in-flight batch reports the append failure ... + onError(new Error(`append rejected (401)`)) + mockProducerFlush.mockImplementationOnce(async () => { + // ... the error event's own batch then lands, so the failure is + // on the stream ... + mockProducerLastSuccessfulOffset.value = `10_200` + // ... and a second batch queued before the failure reports it + // too. Batches are sent concurrently, so a failure after the + // error event is not evidence the error event was lost. + onError(new Error(`append rejected (401)`)) + }) + }, + }) + + await expect( + processWake(makeNotification(), BASE_CONFIG) + ).rejects.toThrow(`append rejected (401)`) + + const body = doneBody() + expect(body.acks).toEqual([ + { path: `/streams/entity:agent-1`, offset: `10_200` }, + ]) + }) + + it(`sends done without acks when the error event was lost with the producer`, async () => { + defineEntity(`test-agent`, { + handler: () => { + const onError = entityProducerOnError() + // The stream rejected an append (a deposed or lapsed write token): + // the runtime records WRITE_FAILED and appends an error event ... + onError(new Error(`append rejected (401)`)) + // ... whose batch the producer then fails to deliver as well, so + // nothing on the stream says this wake failed. + mockProducerFlush.mockImplementationOnce(async () => { + onError(new Error(`append rejected (401)`)) + }) + }, + }) + + await expect( + processWake(makeNotification(), BASE_CONFIG) + ).rejects.toThrow(`append rejected (401)`) + + const body = doneBody() + expect(body.done).toBe(true) + expect(body.acks).toEqual([]) + }) + }) + it(`flushes producer on completion`, async () => { defineEntity(`test-agent`, { handler: () => {}, From 3124fdc071769eab58f66dcac0889afbb390a413 Mon Sep 17 00:00:00 2001 From: Aditya Kumarakrishnan Date: Wed, 2 Sep 2026 16:38:59 +0530 Subject: [PATCH 6/6] fix(agents-server): replace a manifest entry's wake registration in place MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Syncing a manifest entry's wake registration unregistered the entry's rows and then registered the replacement. A source event evaluated between the two statements matched nothing, and nothing ever re-evaluates it: a `runFinished` wake was lost whenever a spawned child's run completed in the few milliseconds after its parent's run wrote the child's manifest entry. The parent's end-of-run batch and the child's `run` update reach the server together, so the two race every time; forwarding claims to the backend (01825eb) moved the child by a claim round trip, enough to land its completion in that gap in a quarter to a half of runtime-dsl F1 runs. The registry now replaces the registration: the replacement is upserted first and only then are the entry's other rows removed, so an event in between matches the old row or the new one, never neither. A manifest that re-describes the registration spawn already made — the usual case — resolves to that same row through `uq_wake_registration`, and nothing is deleted or re-created at all; `include_response`, the one field outside the constraint, is adopted on the row. --- .../agents-manifest-wake-replacement.md | 5 + packages/agents-server/src/runtime.ts | 30 ++-- packages/agents-server/src/wake-registry.ts | 119 ++++++++++++++-- .../agents-server/test/wake-registry.test.ts | 132 +++++++++++++++++- 4 files changed, 261 insertions(+), 25 deletions(-) create mode 100644 .changeset/agents-manifest-wake-replacement.md diff --git a/.changeset/agents-manifest-wake-replacement.md b/.changeset/agents-manifest-wake-replacement.md new file mode 100644 index 0000000000..c09ce9185a --- /dev/null +++ b/.changeset/agents-manifest-wake-replacement.md @@ -0,0 +1,5 @@ +--- +'@electric-ax/agents-server': patch +--- + +Fix a lost `runFinished` wake when a child finishes while its parent's manifest lands. Syncing a manifest entry's wake registration used to unregister the entry's rows and then register the replacement, and a source event evaluated between the two statements — the spawned child's run completing a few milliseconds after the parent's run wrote the child's manifest entry — matched nothing and was never re-evaluated. The registry now replaces a manifest entry's registration in place: the replacement is upserted first (a manifest that re-describes the registration spawn already made resolves to that same row, so nothing is deleted or re-created), and only then are the entry's other rows removed. diff --git a/packages/agents-server/src/runtime.ts b/packages/agents-server/src/runtime.ts index 164ac965e6..5f77ef022c 100644 --- a/packages/agents-server/src/runtime.ts +++ b/packages/agents-server/src/runtime.ts @@ -252,23 +252,29 @@ export class ElectricAgentsTenantRuntime { continue } - await this.manager.wakeRegistry.unregisterByManifestKey( + const reg = buildManifestWakeRegistration( subscriberUrl, - manifestKey, - this.serviceId + value, + manifestKey ) - - if (value) { - const reg = buildManifestWakeRegistration( + if (reg) { + // Replace rather than unregister-then-register: the spawn that wrote + // this manifest entry registered the same wake up front, and its + // child can finish in the gap between the two. + await this.manager.wakeRegistry.replaceByManifestKey({ + ...reg, + manifestKey, + tenantId: this.serviceId, + }) + } else { + await this.manager.wakeRegistry.unregisterByManifestKey( subscriberUrl, - value, - manifestKey + manifestKey, + this.serviceId ) - if (reg) { - reg.tenantId = this.serviceId - await this.manager.wakeRegistry.register(reg) - } + } + if (value) { const cronSpec = extractManifestCronSpec(value) if (cronSpec) { void this.manager diff --git a/packages/agents-server/src/wake-registry.ts b/packages/agents-server/src/wake-registry.ts index ff8ef6e75b..89bd2bf658 100644 --- a/packages/agents-server/src/wake-registry.ts +++ b/packages/agents-server/src/wake-registry.ts @@ -3,7 +3,7 @@ import { isChangeMessage, isControlMessage, } from '@electric-sql/client' -import { and, eq } from 'drizzle-orm' +import { and, eq, ne, sql } from 'drizzle-orm' import { wakeRegistrations } from './db/schema.js' import { serverLog } from './utils/log.js' import { electricUrlWithPath } from './utils/electric-url.js' @@ -329,17 +329,7 @@ export class WakeRegistry { const tenantId = this.resolveTenantId(reg.tenantId) const result = await this.db .insert(wakeRegistrations) - .values({ - tenantId, - subscriberUrl: reg.subscriberUrl, - sourceUrl: reg.sourceUrl, - condition: reg.condition, - debounceMs: reg.debounceMs ?? 0, - timeoutMs: reg.timeoutMs ?? 0, - oneShot: reg.oneShot, - includeResponse: reg.includeResponse !== false, - manifestKey: reg.manifestKey ?? null, - }) + .values(this.registrationRow(reg, tenantId)) .onConflictDoNothing() .returning({ id: wakeRegistrations.id }) @@ -360,6 +350,111 @@ export class WakeRegistry { }) } + /** + * Makes `reg` the one registration its manifest entry owns, without a + * moment in which the entry owns none. The replacement is upserted first + * and only then are the entry's other rows removed, so a source event + * evaluated in between — a child's run finishing while its parent's + * manifest lands — matches the old row or the new one, never neither. An + * unregister-then-register sequence leaves exactly that gap, and a wake + * missed there is never re-evaluated. A manifest that re-describes the + * registration spawn already made (the usual case) resolves to that same + * row, so nothing is deleted or re-created at all. + */ + async replaceByManifestKey( + reg: WakeRegistration & { manifestKey: string } + ): Promise { + const tenantId = this.resolveTenantId(reg.tenantId) + const kept = await this.upsertRegistration(reg, tenantId) + + await this.db + .delete(wakeRegistrations) + .where( + and( + eq(wakeRegistrations.tenantId, tenantId), + eq(wakeRegistrations.subscriberUrl, reg.subscriberUrl), + eq(wakeRegistrations.manifestKey, reg.manifestKey), + ne(wakeRegistrations.id, kept.dbId) + ) + ) + + const stale = Array.from(this.registrationCache.values()).flatMap((regs) => + regs + .filter( + (r) => + r.tenantId === tenantId && + r.subscriberUrl === reg.subscriberUrl && + r.manifestKey === reg.manifestKey && + r.dbId !== kept.dbId + ) + .map((r) => r.dbId) + ) + for (const dbId of stale) { + this.removeCachedRegistrationByDbId(dbId) + } + } + + /** + * Inserts the registration, or adopts the row `uq_wake_registration` says + * already exists for it, and caches whichever it is. `include_response` is + * outside that constraint, so the conflict branch writes it — that is what + * makes RETURNING yield the existing row, and it is the one field a + * re-described registration can legitimately change. + */ + private async upsertRegistration( + reg: WakeRegistration, + tenantId: string + ): Promise { + const rows = await this.db + .insert(wakeRegistrations) + .values(this.registrationRow(reg, tenantId)) + .onConflictDoUpdate({ + target: [ + wakeRegistrations.tenantId, + wakeRegistrations.subscriberUrl, + wakeRegistrations.sourceUrl, + wakeRegistrations.oneShot, + wakeRegistrations.debounceMs, + wakeRegistrations.timeoutMs, + wakeRegistrations.condition, + wakeRegistrations.manifestKey, + ], + set: { includeResponse: sql`excluded.include_response` }, + }) + .returning({ + id: wakeRegistrations.id, + createdAt: wakeRegistrations.createdAt, + timeoutConsumed: wakeRegistrations.timeoutConsumed, + }) + const row = rows[0]! + const cached: CachedWakeRegistration = { + ...reg, + tenantId, + dbId: row.id, + createdAt: row.createdAt, + timeoutConsumed: row.timeoutConsumed, + } + this.upsertCachedRegistration(cached) + return cached + } + + private registrationRow( + reg: WakeRegistration, + tenantId: string + ): typeof wakeRegistrations.$inferInsert { + return { + tenantId, + subscriberUrl: reg.subscriberUrl, + sourceUrl: reg.sourceUrl, + condition: reg.condition, + debounceMs: reg.debounceMs ?? 0, + timeoutMs: reg.timeoutMs ?? 0, + oneShot: reg.oneShot, + includeResponse: reg.includeResponse !== false, + manifestKey: reg.manifestKey ?? null, + } + } + private startTimeoutTimer(reg: CachedWakeRegistration, dbId: number): void { if (reg.timeoutMs == null || reg.timeoutMs <= 0) return this.startTimeoutTimerWithDuration(reg, dbId, reg.timeoutMs) diff --git a/packages/agents-server/test/wake-registry.test.ts b/packages/agents-server/test/wake-registry.test.ts index c879a124ca..7933ac62ca 100644 --- a/packages/agents-server/test/wake-registry.test.ts +++ b/packages/agents-server/test/wake-registry.test.ts @@ -12,7 +12,10 @@ import { createServer } from 'node:http' import { DurableStreamTestServer } from '@durable-streams/server' +import { eq } from 'drizzle-orm' import { afterAll, beforeAll, describe, expect, it, vi } from 'vitest' +import { createDb } from '../src/db/index' +import { wakeRegistrations } from '../src/db/schema' import { EntityManager } from '../src/entity-manager' import { ElectricAgentsServer } from '../src/server' import { WakeRegistry } from '../src/wake-registry' @@ -27,7 +30,7 @@ import { resetElectricAgentsTestBackend, } from './test-backend' import type { Server } from 'node:http' -import type { WakeEvalResult } from '../src/wake-registry' +import type { WakeEvalResult, WakeRegistration } from '../src/wake-registry' let nextDbId = 1 function createMockDb(): any { @@ -1600,3 +1603,130 @@ describe(`Wake Registry Integration`, () => { expect(changes[0]!.collection).toBe(`runs`) }, 15_000) }) + +describe(`Wake Registry manifest replacement`, () => { + let db: ReturnType[`db`] + let client: ReturnType[`client`] + + beforeAll(async () => { + await resetElectricAgentsTestBackend() + const connection = createDb(TEST_POSTGRES_URL) + db = connection.db + client = connection.client + }, 120_000) + + afterAll(async () => { + vi.restoreAllMocks() + await client?.end() + }, 120_000) + + const runCompleted = { + type: `run`, + key: `run-0`, + value: { status: `completed` }, + headers: { operation: `update` }, + } + + // The registration spawn makes for an observed child, keyed by the child's + // manifest entry — the one the parent's manifest sync re-describes. + function spawnWake(name: string): WakeRegistration & { manifestKey: string } { + return { + subscriberUrl: `/parent/${name}`, + sourceUrl: `/child/${name}`, + condition: `runFinished`, + oneShot: false, + includeResponse: true, + manifestKey: `child:child:${name}`, + } + } + + async function registrationIds( + subscriberUrl: string + ): Promise> { + const rows = await db + .select({ id: wakeRegistrations.id }) + .from(wakeRegistrations) + .where(eq(wakeRegistrations.subscriberUrl, subscriberUrl)) + return rows.map((row) => row.id) + } + + it(`keeps the row spawn registered when the manifest re-describes it`, async () => { + const registry = new WakeRegistry(db) + const reg = spawnWake(`same`) + await registry.register(reg) + const before = registry.evaluate(reg.sourceUrl, runCompleted) + expect(before).toHaveLength(1) + + await registry.replaceByManifestKey(reg) + + const after = registry.evaluate(reg.sourceUrl, runCompleted) + expect(after.map((result) => result.registrationDbId)).toEqual([ + before[0]!.registrationDbId, + ]) + expect(await registrationIds(reg.subscriberUrl)).toEqual([ + before[0]!.registrationDbId, + ]) + }) + + it(`adopts a changed includeResponse on the existing row`, async () => { + const registry = new WakeRegistry(db) + const reg = spawnWake(`include-response`) + await registry.register(reg) + const [before] = registry.evaluate(reg.sourceUrl, runCompleted) + + await registry.replaceByManifestKey({ ...reg, includeResponse: false }) + + const after = registry.evaluate(reg.sourceUrl, runCompleted) + expect(after).toHaveLength(1) + expect(after[0]!.registrationDbId).toBe(before!.registrationDbId) + expect(after[0]!.includeResponse).toBe(false) + const rows = await db + .select({ includeResponse: wakeRegistrations.includeResponse }) + .from(wakeRegistrations) + .where(eq(wakeRegistrations.id, before!.registrationDbId)) + expect(rows).toEqual([{ includeResponse: false }]) + }) + + it(`never leaves the manifest key without a registration while replacing it`, async () => { + const registry = new WakeRegistry(db) + const reg = spawnWake(`gap`) + await registry.register(reg) + + // A child finishing while its parent's manifest lands has its run event + // evaluated somewhere between the registry's statements. Evaluate at + // each one and require a match every time: unregister-then-register + // would answer with nothing once the delete had landed. + const matches: Array = [] + const observe = (): void => { + matches.push(registry.evaluate(reg.sourceUrl, runCompleted).length) + } + const originalInsert = db.insert.bind(db) + const originalDelete = db.delete.bind(db) + vi.spyOn(db, `insert`).mockImplementation(((...args: Array) => { + observe() + return (originalInsert as any)(...args) + }) as any) + vi.spyOn(db, `delete`).mockImplementation(((...args: Array) => { + observe() + return (originalDelete as any)(...args) + }) as any) + + // A replacement that changes the unique shape, so the old row must go. + const replacement = { ...reg, condition: { on: `change` } as const } + try { + await registry.replaceByManifestKey(replacement) + } finally { + vi.restoreAllMocks() + } + + expect(matches.length).toBeGreaterThanOrEqual(2) + expect(matches.every((count) => count >= 1)).toBe(true) + + const after = registry.evaluate(reg.sourceUrl, runCompleted) + expect(after).toHaveLength(1) + expect(after[0]!.runFinishedStatus).toBeUndefined() + expect(await registrationIds(reg.subscriberUrl)).toEqual([ + after[0]!.registrationDbId, + ]) + }) +})