Adopt backend-issued claim write tokens and opt-in fenced session streams

.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.

packages/agents-runtime/src/process-wake.ts

@@ -1248,6 +1248,7 @@ export async function processWake(
             ok: boolean
             claimToken?: string
             token?: string
+            writeToken?: string
           }
           if (!data.ok) {
             failBackgroundWake(
@@ -1258,6 +1259,10 @@ export async function processWake(
           }
           if (data.claimToken) activeClaimToken = data.claimToken
           if (data.token) activeClaimToken = data.token
+          // The server may re-mint the write token on each heartbeat (its
+          // backend's token TTL tracks the claim lease); adopt the refresh
+          // so appends from a long-running activation keep a live token.
+          if (data.writeToken) writeToken = data.writeToken
         })
         .catch((err: unknown) => {
           failBackgroundWake(err, `HEARTBEAT_FAILED`)

packages/agents-runtime/test/process-wake.test.ts

@@ -1847,6 +1847,59 @@ describe(`processWake`, () => {
     setIntervalSpy.mockRestore()
   })
 
+  it(`adopts a refreshed write token from heartbeat responses`, async () => {
+    defineEntity(`test-agent`, {
+      handler: async () => {
+        // Keep the wake alive long enough for several heartbeat intervals.
+        await new Promise((resolve) => setTimeout(resolve, 150))
+      },
+    })
+
+    fetchMock.mockImplementation(async (url, opts) => {
+      if (String(url).includes(`/_electric/wakes/wake-abc`)) {
+        const body = JSON.parse(String(opts?.body ?? `{}`)) as Record<
+          string,
+          unknown
+        >
+        const isClaim = body.wakeId !== undefined
+        const isHeartbeat = !isClaim && body.done === undefined
+        return new Response(
+          JSON.stringify({
+            ok: true,
+            ...(isClaim ? { writeToken: `wt-initial` } : {}),
+            ...(isHeartbeat ? { writeToken: `wt-refreshed` } : {}),
+          }),
+          { status: 200, headers: { 'content-type': `application/json` } }
+        )
+      }
+      return new Response(JSON.stringify({ ok: true }), {
+        status: 200,
+        headers: { 'content-type': `application/json` },
+      })
+    })
+
+    await processWake(makeNotification(), {
+      ...BASE_CONFIG,
+      heartbeatInterval: 20,
+    })
+
+    // The producer reads the write token per request, so appends issued
+    // after a heartbeat refresh must carry the refreshed token.
+    const producer = mockConstructedProducers.find(
+      (constructed) =>
+        constructed.producerId ===
+        `entity-http://localhost:3000/test-agent/agent-1`
+    )
+    const producerFetch = producer!.opts!.fetch as typeof fetch
+    await producerFetch(`http://localhost:3000/streams/entity:agent-1`, {
+      method: `POST`,
+    })
+    const [, init] = fetchMock.mock.calls.at(-1)!
+    expect(new Headers(init?.headers).get(`authorization`)).toBe(
+      `Bearer wt-refreshed`
+    )
+  })
+
   it(`flushes producer on completion`, async () => {
     defineEntity(`test-agent`, {
       handler: () => {},

packages/agents-server/src/claim-write-token-store.ts

@@ -2,14 +2,27 @@ import { randomUUID } from 'node:crypto'
 
 interface ActiveClaimWriteToken {
   token: string
+  /**
+   * The token this consumer's previous mint for the same stream issued.
+   * Kept valid so an append already in flight with the pre-refresh token
+   * is not rejected when a heartbeat re-mints. A different consumer's mint
+   * evicts the claim outright, previous token included.
+   */
+  previousToken?: string
   consumerId: string
 }
 
 export class ClaimWriteTokenStore {
   private readonly claimsByStream = new Map<string, ActiveClaimWriteToken>()
   private readonly streamKeysByConsumer = new Map<string, Set<string>>()
+  private readonly deliveredTokensByConsumer = new Map<string, string>()
 
-  mint(service: string, streamPath: string, consumerId: string): string {
+  mint(
+    service: string,
+    streamPath: string,
+    consumerId: string,
+    token: string = randomUUID()
+  ): string {
     const streamKey = this.streamKey(service, streamPath)
     const consumerKey = this.consumerKey(service, consumerId)
     const previousClaimForStream = this.claimsByStream.get(streamKey)
@@ -20,16 +33,46 @@ export class ClaimWriteTokenStore {
       )
     }
 
-    const token = randomUUID()
-    this.claimsByStream.set(streamKey, { token, consumerId })
+    this.claimsByStream.set(streamKey, {
+      token,
+      consumerId,
+      ...(previousClaimForStream?.consumerId === consumerId
+        ? { previousToken: previousClaimForStream.token }
+        : {}),
+    })
     this.addConsumerStream(consumerKey, streamKey)
     return token
   }
 
+  /**
+   * Remembers a write token the Durable Streams backend issued with a wake
+   * delivery (webhook notification or pull-wake claim), keyed by the wake's
+   * consumer id, until that consumer's claim callback mints it as the active
+   * claim write token.
+   */
+  recordDelivered(service: string, consumerId: string, token: string): void {
+    this.deliveredTokensByConsumer.set(
+      this.consumerKey(service, consumerId),
+      token
+    )
+  }
+
+  takeDelivered(service: string, consumerId: string): string | undefined {
+    const consumerKey = this.consumerKey(service, consumerId)
+    const token = this.deliveredTokensByConsumer.get(consumerKey)
+    if (token !== undefined) {
+      this.deliveredTokensByConsumer.delete(consumerKey)
+    }
+    return token
+  }
+
   isValid(service: string, streamPath: string, token: string): boolean {
+    const activeClaim = this.claimsByStream.get(
+      this.streamKey(service, streamPath)
+    )
     return (
-      this.claimsByStream.get(this.streamKey(service, streamPath))?.token ===
-      token
+      activeClaim !== undefined &&
+      (activeClaim.token === token || activeClaim.previousToken === token)
     )
   }
 
@@ -54,6 +97,7 @@ export class ClaimWriteTokenStore {
 
   clearConsumer(service: string, consumerId: string): void {
     const consumerKey = this.consumerKey(service, consumerId)
+    this.deliveredTokensByConsumer.delete(consumerKey)
     const streamKeys = this.streamKeysByConsumer.get(consumerKey)
     if (!streamKeys) return
 

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)
         }

packages/agents-server/src/routing/internal-router.ts

@@ -84,6 +84,7 @@ const subscriptionWebhookBodySchema = Type.Object(
     streams: Type.Optional(Type.Array(Type.Record(Type.String(), Type.Any()))),
     callback_url: Type.Optional(Type.String()),
     callback_token: Type.Optional(Type.String()),
+    write_token: Type.Optional(Type.String()),
     primary_stream: Type.Optional(Type.String()),
     primaryStream: Type.Optional(Type.String()),
     streamPath: Type.Optional(Type.String()),
@@ -315,6 +316,7 @@ function newWebhookPayload(body: SubscriptionWebhookBody | undefined): {
   tailOffset: string
   callbackUrl: string
   callbackToken: string
+  writeToken?: string
 } | null {
   if (
     !body ||
@@ -358,6 +360,7 @@ function newWebhookPayload(body: SubscriptionWebhookBody | undefined): {
     tailOffset: selectedStream.tail_offset,
     callbackUrl: body.callback_url,
     callbackToken: body.callback_token,
+    ...(body.write_token ? { writeToken: body.write_token } : {}),
   }
 }
 
@@ -585,6 +588,22 @@ async function subscriptionWebhook(
         )
       }
 
+      if (newWebhook?.writeToken) {
+        // The backend minted this wake's write token before delivery (Write
+        // Fencing extension); hold it for the runtime's claim callback to
+        // adopt. Recorded only after the stopped/paused auto-ack and
+        // fork-lock rejections above, so a wake that never reaches the
+        // runtime — and therefore never claims — leaves no store entry
+        // behind. It must be recorded before the forward below, because the
+        // runtime's claim callback can arrive while the forward is still in
+        // flight.
+        ctx.runtime.claimWriteTokens.recordDelivered(
+          ctx.service,
+          newWebhook.wakeId,
+          newWebhook.writeToken
+        )
+      }
+
       if (entity) {
         rootSpan?.setAttribute(ATTR.ENTITY_URL, entity.url)
         await tracer.startActiveSpan(
@@ -789,6 +808,29 @@ async function wakeCallback(
         responseBytes = new TextEncoder().encode(JSON.stringify(responseBody))
       }
     }
+  } else if (!isDoneRequest && upstream.ok && target.primaryStream) {
+    // Heartbeat: a backend implementing the Write Fencing extension re-mints
+    // the claim's write token on every ack. Adopt the refresh so the active
+    // token tracks the backend's, and surface it to the runtime as
+    // `writeToken` alongside the passed-through ack body.
+    const responseBody = decodeJsonObject(responseBytes)
+    const refreshedToken = responseBody?.write_token
+    if (
+      responseBody?.ok === true &&
+      typeof refreshedToken === `string` &&
+      refreshedToken !== ``
+    ) {
+      const writeToken = await mintClaimWriteToken(
+        ctx,
+        target.primaryStream,
+        consumerId,
+        refreshedToken
+      )
+      if (writeToken) {
+        responseBody.writeToken = writeToken
+        responseBytes = new TextEncoder().encode(JSON.stringify(responseBody))
+      }
+    }
   }
 
   try {
@@ -904,12 +946,45 @@ async function wakeCallback(
 async function mintClaimWriteToken(
   ctx: TenantContext,
   streamPath: string,
-  consumerId: string
+  consumerId: string,
+  backendToken?: string
 ): Promise<string | undefined> {
   const entity = await ctx.entityManager.registry.getEntityByStream(streamPath)
   if (!entity) return undefined
 
-  return ctx.runtime.claimWriteTokens.mint(ctx.service, streamPath, consumerId)
+  // When the Durable Streams backend issued a write token for this claim
+  // (Write Fencing extension), adopt it as the claim's write token so the
+  // store stays the single validation authority while the backend is the
+  // mint. When the backend supplied none, the store mints its own token and
+  // behaviour is byte-for-byte what it is today.
+  //
+  // Version-skew matrix — adoption is data-driven (it follows `write_token`
+  // fields wherever the backend sends them), while stream fencing is a
+  // separate opt-in (`fencedSessionStreams`):
+  // - Base backend (no Write Fencing), any server/runtime: no `write_token`
+  //   ever appears, so the store mints and nothing changes. With
+  //   `fencedSessionStreams` on, the fencing headers are sent but a base
+  //   backend ignores them (additive headers, base spec §11), so
+  //   enforcement remains this store only.
+  // - Token-minting backend, older server (no adoption): the optional
+  //   fields are ignored and the old server mints its own tokens; it never
+  //   creates streams fenced, so the backend enforces nothing.
+  // - Token-minting backend, this server, older runtime (one that ignores
+  //   the heartbeat response's `writeToken`): adoption still happens here,
+  //   so a backend that rotates the token on ack refresh ages the runtime's
+  //   original token out of the store's one-refresh grace, and the
+  //   runtime's appends start failing 401 mid-activation. Upgrade runtimes
+  //   before pointing this server at a token-minting backend (and before
+  //   enabling `fencedSessionStreams`).
+  const token =
+    backendToken ??
+    ctx.runtime.claimWriteTokens.takeDelivered(ctx.service, consumerId)
+  return ctx.runtime.claimWriteTokens.mint(
+    ctx.service,
+    streamPath,
+    consumerId,
+    token
+  )
 }
 
 function encodeWakeCallbackBody(

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,