diff --git a/.changeset/agents-backend-issued-write-tokens.md b/.changeset/agents-backend-issued-write-tokens.md new file mode 100644 index 0000000000..b8c4f5278c --- /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 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/.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/.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/.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 472cba39b5..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) @@ -1248,6 +1264,7 @@ export async function processWake( ok: boolean claimToken?: string token?: string + writeToken?: string } if (!data.ok) { failBackgroundWake( @@ -1258,6 +1275,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`) @@ -2556,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?.() @@ -2595,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 49c3a5460f..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) => { @@ -1847,6 +1856,180 @@ 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(`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: () => {}, diff --git a/packages/agents-server/src/claim-write-token-store.ts b/packages/agents-server/src/claim-write-token-store.ts index 1975e78b46..7424b13e3e 100644 --- a/packages/agents-server/src/claim-write-token-store.ts +++ b/packages/agents-server/src/claim-write-token-store.ts @@ -2,14 +2,70 @@ 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 + expiresAt: number +} + +interface DeliveredWriteToken { + token: string + expiresAt: number +} + +/** + * 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. + */ +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 + now?: () => number } export class ClaimWriteTokenStore { private readonly claimsByStream = new Map() private readonly streamKeysByConsumer = 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, streamPath: string, consumerId: string): string { + mint( + service: string, + streamPath: string, + 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) @@ -20,23 +76,69 @@ export class ClaimWriteTokenStore { ) } - const token = randomUUID() - this.claimsByStream.set(streamKey, { token, consumerId }) + this.claimsByStream.set(streamKey, { + token, + consumerId, + expiresAt: this.now() + this.ttlMs, + ...(previousClaimForStream?.consumerId === consumerId + ? { previousToken: previousClaimForStream.token } + : {}), + }) this.addConsumerStream(consumerKey, streamKey) 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 + * 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.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 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.liveClaim(this.streamKey(service, streamPath)) return ( - this.claimsByStream.get(this.streamKey(service, streamPath))?.token === - token + activeClaim !== undefined && + (activeClaim.token === token || activeClaim.previousToken === token) ) } 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 ) } @@ -54,6 +156,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 @@ -63,6 +166,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/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-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/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 87b78ee00d..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' @@ -84,6 +88,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 +320,7 @@ function newWebhookPayload(body: SubscriptionWebhookBody | undefined): { tailOffset: string callbackUrl: string callbackToken: string + writeToken?: string } | null { if ( !body || @@ -358,6 +364,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 +592,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( @@ -651,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, @@ -663,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) } @@ -704,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, @@ -738,9 +784,6 @@ async function wakeCallback( let upstream: Response try { - const subscriptionId = durableStreamsSubscriptionCallback( - target.callbackUrl - ) if (subscriptionId) { const token = claimTokenFromRequest(request) if (!token) { @@ -765,34 +808,104 @@ 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 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 = backendWriteToken(responseBody) + if (responseBody?.ok === true && refreshedToken !== undefined) { + const writeToken = await mintClaimWriteToken( + ctx, + target.primaryStream, + consumerId, + refreshedToken + ) + if (writeToken) { + 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 && @@ -846,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( @@ -870,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) { @@ -901,15 +1022,88 @@ 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 + consumerId: 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 - 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 — + // 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 — 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. + // - 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..7daa67413d 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' @@ -103,7 +104,20 @@ async function handleStreamAppend( if (entity) { const token = writeTokenFromHeaders(request.headers) - if (!manager.isValidWriteToken(entity, token)) { + if (manager.fencedSessionStreams) { + // 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)) { diff --git a/packages/agents-server/src/runtime.ts b/packages/agents-server/src/runtime.ts index c545fa714c..5f77ef022c 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 ?? @@ -250,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/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..275a57f6c5 100644 --- a/packages/agents-server/src/stream-client.ts +++ b/packages/agents-server/src/stream-client.ts @@ -11,6 +11,19 @@ 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, 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 + */ +export const WRITE_FENCE_HEADER = `Write-Fence` +export const WRITE_TOKEN_HEADER = `Write-Token` + export interface StreamClientOptions { bearer?: DurableStreamsBearerProvider } @@ -72,6 +85,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 +244,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,18 +258,29 @@ 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, }) + 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) + } }) } 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 +291,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 @@ -275,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()}` @@ -283,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/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/claim-write-token-store.test.ts b/packages/agents-server/test/claim-write-token-store.test.ts index f7ebed3a2a..086a8834cf 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,131 @@ 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`) + }) + + 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() + }) + }) }) 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/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..3248412b2b 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`) @@ -656,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({ @@ -674,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/stream-append.test.ts b/packages/agents-server/test/stream-append.test.ts new file mode 100644 index 0000000000..99fd8659ea --- /dev/null +++ b/packages/agents-server/test/stream-append.test.ts @@ -0,0 +1,148 @@ +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 + knownToken?: boolean +}): ElectricAgentsStreamAppendRuntime { + return { + manager: { + registry: { + getEntityByStream: vi + .fn() + .mockResolvedValue(opts.hasEntity === false ? null : entity), + }, + fencedSessionStreams: opts.fencedSessionStreams, + isAttachmentStreamPath: vi.fn(() => false), + isValidWriteToken: vi.fn(() => opts.knownToken ?? 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(`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 }) => { + 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..94e365762c 100644 --- a/packages/agents-server/test/stream-client-fork.test.ts +++ b/packages/agents-server/test/stream-client-fork.test.ts @@ -35,6 +35,27 @@ describe(`StreamClient.fork`, () => { ).resolves.toEqual({ offset: expect.any(String) }) }) + 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: `[]`, + }) + await expect( + client.fork(`/fork-fenced`, `/source-unfenced`, { writeFence: true }) + ).rejects.toThrow(/did not echo Write-Fence/) + }) + 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..1eb9c19303 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,96 @@ 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`) + + 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` + ) + // 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, headers: { 'Write-Fence': `true` } }) + ) + 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(`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 }), { diff --git a/packages/agents-server/test/subscription-webhooks-routing.test.ts b/packages/agents-server/test/subscription-webhooks-routing.test.ts index 9c2bf2cef7..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,41 +597,640 @@ 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`).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( + 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) + clock.now = 1_500 + expect( + claimWriteTokens.isValid(`tenant-a`, `/horton/demo/main`, activeToken) + ).toBe(true) + } finally { + fetchSpy.mockRestore() + } + }) + + 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({ 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( + 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`) - + 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( - 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: [{ 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).toEqual(expect.any(String)) - expect(body.writeToken).not.toBe(`write-token`) + await expect(responseJson(response)).resolves.toEqual({ ok: true }) expect( ctx.runtime.claimWriteTokens.isValid( `tenant-a`, `/horton/demo/main`, - body.writeToken + activeToken ) ).toBe(true) - expect(fetchSpy).not.toHaveBeenCalled() } finally { fetchSpy.mockRestore() } @@ -686,6 +1289,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([ { @@ -741,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` ) @@ -859,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() } @@ -906,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() } @@ -986,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`, @@ -1003,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() } @@ -1064,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`, @@ -1076,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() + } + }) }) }) 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, + ]) + }) +})