@electric-ax/agents-server 0.4.5 → 0.4.7

package/src/routing/internal-router.ts

@@ -2,7 +2,10 @@
  * Sub-router for /_electric/* control-plane routes.
  */
 
-import { appendPathToUrl } from '@electric-ax/agents-runtime'
+import {
+  appendPathToUrl,
+  verifyWebhookSignature,
+} from '@electric-ax/agents-runtime'
 import { Type, type Static } from '@sinclair/typebox'
 import { and, eq } from 'drizzle-orm'
 import { Router, json, status } from 'itty-router'
@@ -16,10 +19,13 @@ import {
   ErrCodeCallbackNotFound,
   ErrCodeForkInProgress,
   ErrCodeSubscriptionNotFound,
+  ErrCodeUnauthorized,
 } 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 { applyDurableStreamsBearer } from '../stream-client.js'
+import { getDefaultWebhookSigner } from '../webhook-signing.js'
 import { cronRouter } from './cron-router.js'
 import { resolveDurableStreamsRoutingAdapter } from './durable-streams-routing-adapter.js'
 import { electricProxyRouter } from './electric-proxy-router.js'
@@ -30,8 +36,10 @@ import { runnersRouter } from './runners-router.js'
 import { routeBody, validateOptionalJsonBody, withSchema } from './schema.js'
 import { withLeadingSlash } from './tenant-stream-paths.js'
 import type { IRequest, RouterType } from 'itty-router'
+import type { WebhookSignatureVerifierConfig } from '@electric-ax/agents-runtime'
 import type { TenantContext } from './context.js'
 import type { DurableStreamsRoutingAdapter } from './durable-streams-routing-adapter.js'
+import type { WebhookSigner } from '../webhook-signing.js'
 
 const wakeRegistrationBodySchema = Type.Object({
   subscriberUrl: Type.String(),
@@ -137,13 +145,22 @@ function bodyFromBytes(body: Uint8Array): ArrayBuffer {
 }
 
 function responseFromUpstream(response: Response, body?: Uint8Array): Response {
-  return new Response(body ? bodyFromBytes(body) : response.body, {
+  const responseBody = forbidsResponseBody(response.status)
+    ? null
+    : body !== undefined
+      ? bodyFromBytes(body)
+      : response.body
+  return new Response(responseBody, {
     status: response.status,
     statusText: response.statusText,
     headers: responseHeaders(response),
   })
 }
 
+function forbidsResponseBody(status: number): boolean {
+  return status === 204 || status === 205 || status === 304
+}
+
 function forwardHeadersFromRequest(request: IRequest): Headers {
   const headers = new Headers(request.headers)
   headers.delete(`host`)
@@ -156,6 +173,87 @@ function durableStreamsSubscriptionCallback(value: string): string | null {
     : null
 }
 
+function resolveWebhookSigner(ctx: TenantContext): WebhookSigner {
+  return ctx.webhookSigner ?? getDefaultWebhookSigner()
+}
+
+function durableStreamsWebhookJwksUrl(ctx: TenantContext): string {
+  if (!ctx.durableStreamsRouting) {
+    return appendPathToUrl(ctx.durableStreamsUrl, `/__ds/jwks.json`)
+  }
+
+  return resolveDurableStreamsRoutingAdapter(
+    ctx.durableStreamsRouting,
+    ctx.durableStreamsUrl
+  )
+    .controlUrl({
+      durableStreamsUrl: ctx.durableStreamsUrl,
+      serviceId: ctx.service,
+      requestUrl: appendPathToUrl(ctx.publicUrl, `/__ds/jwks.json`),
+    })
+    .toString()
+}
+
+function durableStreamsJwksFetchClient(ctx: TenantContext): typeof fetch {
+  return async (input, init) => {
+    const headers = new Headers(init?.headers)
+    await applyDurableStreamsBearer(headers, ctx.durableStreamsBearer, {
+      overwrite: false,
+    })
+    const nextInit: RequestInit & {
+      dispatcher?: TenantContext[`durableStreamsDispatcher`]
+    } = {
+      ...(init ?? {}),
+      headers,
+    }
+    if (ctx.durableStreamsDispatcher) {
+      nextInit.dispatcher = ctx.durableStreamsDispatcher
+    }
+    return await fetch(input, nextInit as RequestInit)
+  }
+}
+
+function resolveDurableStreamsWebhookSignature(
+  ctx: TenantContext
+): false | WebhookSignatureVerifierConfig {
+  if (ctx.durableStreamsWebhookSignature === false) return false
+
+  return {
+    jwksUrl:
+      ctx.durableStreamsWebhookSignature?.jwksUrl ??
+      durableStreamsWebhookJwksUrl(ctx),
+    toleranceSeconds: ctx.durableStreamsWebhookSignature?.toleranceSeconds,
+    cacheTtlMs: ctx.durableStreamsWebhookSignature?.cacheTtlMs,
+    fetchClient:
+      ctx.durableStreamsWebhookSignature?.fetchClient ??
+      durableStreamsJwksFetchClient(ctx),
+  }
+}
+
+async function verifyDurableStreamsWebhook(
+  request: IRequest,
+  ctx: TenantContext,
+  body: Uint8Array
+): Promise<Response | null> {
+  const config = resolveDurableStreamsWebhookSignature(ctx)
+  if (config === false) return null
+
+  const verification = await verifyWebhookSignature(
+    body,
+    request.headers.get(`webhook-signature`),
+    config
+  )
+  if (verification.ok) return null
+
+  return apiError(
+    verification.status,
+    verification.status === 401
+      ? ErrCodeUnauthorized
+      : `WEBHOOK_SIGNATURE_UNAVAILABLE`,
+    verification.error
+  )
+}
+
 function claimTokenFromRequest(request: IRequest): string | undefined {
   const electricClaimToken = request.headers.get(`electric-claim-token`)?.trim()
   if (electricClaimToken) return electricClaimToken
@@ -249,7 +347,11 @@ async function webhookForward(
     subscriptionId
   )
 
-  const lookupPromise: Promise<string | null> = tracer.startActiveSpan(
+  const body = await readRequestBody(request as Request)
+  const signatureError = await verifyDurableStreamsWebhook(request, ctx, body)
+  if (signatureError) return signatureError
+
+  const targetWebhookUrl: string | null = await tracer.startActiveSpan(
     `db.lookupSubscription`,
     async (span) => {
       try {
@@ -270,11 +372,6 @@ async function webhookForward(
     }
   )
 
-  const [targetWebhookUrl, body] = await Promise.all([
-    lookupPromise,
-    readRequestBody(request as Request),
-  ])
-
   if (!targetWebhookUrl) {
     return apiError(
       404,
@@ -382,7 +479,7 @@ async function webhookForward(
         enrichPromise,
       ])
 
-      if (entity?.status === `stopped`) {
+      if (entity?.status === `stopped` || entity?.status === `paused`) {
         if (upsertPromise) await upsertPromise
         return json({ done: true })
       }
@@ -439,6 +536,10 @@ async function webhookForward(
   const headers = forwardHeadersFromRequest(request)
   headers.set(`content-type`, `application/json`)
   headers.delete(`content-length`)
+  headers.set(
+    `webhook-signature`,
+    await resolveWebhookSigner(ctx).sign(forwardBody)
+  )
 
   let upstream: Response
   try {
@@ -624,8 +725,15 @@ async function callbackForward(
       const entity = await ctx.entityManager.registry.getEntityByStream(
         target.primaryStream
       )
-      if (entity && stillOwnsClaim) {
-        if (epoch !== undefined) {
+
+      // Release the consumer_claims row by its DB identity (consumerId,
+      // epoch). The in-memory write token is a separate concern (write
+      // authorization during the run); release of the durable row must
+      // succeed even if the token was lost (server restart) or evicted
+      // (a later wake re-minted for the same stream).
+      let entityCleared = false
+      if (epoch !== undefined) {
+        const result =
           await ctx.entityManager.registry.materializeReleasedClaim?.({
             consumerId,
             epoch,
@@ -643,28 +751,44 @@ async function callbackForward(
                 })
               : undefined,
           })
-        }
-        await ctx.entityManager.registry.updateStatus(entity.url, `idle`)
-        ctx.runtime.claimWriteTokens.clearStream(
-          ctx.service,
-          target.primaryStream
+        entityCleared = result?.entityCleared ?? false
+      }
+
+      // Transition entity back to idle when either signal says it's safe:
+      // - entityCleared: our release just cleared the entity's active
+      //   dispatch state, so no in-flight wake remains.
+      // - 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.
+      // If both are false, a newer wake owns the entity — leave status as-is.
+      if (entity && (entityCleared || stillOwnsClaim)) {
+        await ctx.entityManager.registry.updateStatus(
+          entity.url,
+          entity.status === `stopping` ? `stopped` : `idle`
         )
         await ctx.entityBridgeManager.onEntityChanged(entity.url)
         serverLog.info(
-          `[callback-forward] status updated to idle for ${entity.url}`
+          `[callback-forward] status updated after done for ${entity.url}`
+        )
+      } else if (!entity) {
+        serverLog.warn(
+          `[callback-forward] done received but no entity found for stream=${target.primaryStream}`
         )
-      } else if (stillOwnsClaim) {
+      }
+
+      // 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,
           target.primaryStream
         )
       } else if (entity) {
         serverLog.info(
-          `[callback-forward] done ignored for stale claim stream=${target.primaryStream} consumer=${consumerId}`
-        )
-      } else {
-        serverLog.warn(
-          `[callback-forward] done received but no entity found for stream=${target.primaryStream}`
+          `[callback-forward] done arrived after in-memory token evicted (stream=${target.primaryStream} consumer=${consumerId})`
         )
       }
     } else if (requestBody?.done === true) {

package/src/routing/runners-router.ts

@@ -575,7 +575,7 @@ async function notificationFromClaim(
       404
     )
   }
-  if (entity.status === `stopped`) {
+  if (entity.status === `stopped` || entity.status === `paused`) {
     await ctx.streamClient.releaseSubscription(
       input.subscriptionId,
       input.claim.token,

package/src/runtime.ts

@@ -533,7 +533,11 @@ export class ElectricAgentsTenantRuntime {
       return
     }
 
-    await this.manager.registry.updateStatus(entityUrl, `idle`)
+    const entity = await this.manager.registry.getEntity(entityUrl)
+    await this.manager.registry.updateStatus(
+      entityUrl,
+      entity?.status === `stopping` ? `stopped` : `idle`
+    )
     await this.entityBridgeManager.onEntityChanged(entityUrl)
   }
 }

package/src/server.ts

@@ -19,6 +19,7 @@ import {
 } from './electric-agents-types.js'
 import { ElectricAgentsError } from './entity-manager.js'
 import { serverLog } from './utils/log.js'
+import { createEd25519WebhookSigner } from './webhook-signing.js'
 import type { DrizzleDB, PgClient } from './db/index.js'
 import type { Server } from 'node:http'
 import type { DurableStreamTestServer } from '@durable-streams/server'
@@ -27,6 +28,7 @@ import type {
   AgentModel,
   EntityRegistry,
   RuntimeHandler,
+  WebhookSignatureVerifierConfig,
 } from '@electric-ax/agents-runtime'
 import type { Principal } from './principal.js'
 import type { EntityBridgeCoordinator } from './entity-bridge-manager.js'
@@ -34,6 +36,10 @@ import type { DurableStreamsRoutingAdapter } from './routing/durable-streams-rou
 import type { OssServerContext } from './routing/oss-server-router.js'
 import type { StartedStandaloneAgentsRuntime } from './standalone-runtime.js'
 import type { DurableStreamsBearerProvider } from './stream-client.js'
+import type {
+  WebhookSigner,
+  WebhookSigningKeyInput,
+} from './webhook-signing.js'
 
 const MOCK_AGENT_HANDLER_PATH = `/_electric/mock-agent-handler`
 
@@ -45,6 +51,11 @@ export interface ElectricAgentsServerOptions {
   durableStreamsBearer?: DurableStreamsBearerProvider
   durableStreamsRouting?: DurableStreamsRoutingAdapter
   durableStreamsServer?: DurableStreamTestServer
+  durableStreamsWebhookSignature?:
+    | false
+    | Partial<WebhookSignatureVerifierConfig>
+  webhookSigner?: WebhookSigner
+  webhookSigningKey?: WebhookSigningKeyInput
   port: number
   host?: string
   workingDirectory?: string
@@ -141,6 +152,7 @@ export class ElectricAgentsServer {
   private shuttingDown = false
   private streamsAgent?: Agent
   private standaloneRuntime?: StartedStandaloneAgentsRuntime
+  private readonly webhookSigner: WebhookSigner
 
   streamClient: StreamClient
   readonly options: ElectricAgentsServerOptions
@@ -152,6 +164,9 @@ export class ElectricAgentsServer {
       )
     }
     this.options = options
+    this.webhookSigner =
+      options.webhookSigner ??
+      createEd25519WebhookSigner({ privateKey: options.webhookSigningKey })
     this.streamClient = options.durableStreamsUrl
       ? new StreamClient(options.durableStreamsUrl, {
           bearer: options.durableStreamsBearer,
@@ -413,6 +428,9 @@ export class ElectricAgentsServer {
       durableStreamsBearer: this.options.durableStreamsBearer,
       durableStreamsRouting: this.options.durableStreamsRouting,
       durableStreamsDispatcher: this.streamsAgent,
+      durableStreamsWebhookSignature:
+        this.options.durableStreamsWebhookSignature,
+      webhookSigner: this.webhookSigner,
       electricUrl: this.options.electricUrl,
       electricSecret: this.options.electricSecret,
       ownAgentHandlerPaths: this.mockAgentBootstrap

package/src/stream-client.ts

@@ -44,11 +44,17 @@ export interface SubscriptionResponse {
   type?: `webhook` | `pull-wake`
   pattern?: string
   streams?: Array<string | SubscriptionStreamInfo>
-  webhook?: { url?: string }
+  webhook?: {
+    url?: string
+    signing?: {
+      alg?: string
+      kid?: string
+      jwks_url?: string
+    }
+  }
   wake_stream?: string
   callback_url?: string
   callback_token?: string
-  webhook_secret?: string
 }
 
 export interface SubscriptionCreateInput {
@@ -535,14 +541,14 @@ export class StreamClient {
     subscriptionId: string,
     webhookUrl: string,
     description?: string
-  ): Promise<{ subscription_id: string; webhook_secret?: string }> {
+  ): Promise<SubscriptionResponse> {
     const res = await this.putSubscription(subscriptionId, {
       type: `webhook`,
       pattern: normalizeSubscriptionPattern(pattern),
       webhook: { url: webhookUrl },
       ...(description ? { description } : {}),
     })
-    return res as { subscription_id: string; webhook_secret?: string }
+    return res
   }
 
   async putSubscription(

package/src/webhook-signing.ts

@@ -0,0 +1,173 @@
+import {
+  createHash,
+  createPrivateKey,
+  createPublicKey,
+  generateKeyPairSync,
+  sign,
+} from 'node:crypto'
+import { appendPathToUrl } from '@electric-ax/agents-runtime'
+import type { JsonWebKey as NodeJsonWebKey, KeyObject } from 'node:crypto'
+
+export interface WebhookPublicJwk {
+  kty: `OKP`
+  crv: `Ed25519`
+  x: string
+  kid: string
+  use: `sig`
+  alg: `EdDSA`
+}
+
+export interface WebhookJwks {
+  keys: Array<WebhookPublicJwk>
+}
+
+export interface WebhookSigningMetadata {
+  alg: `ed25519`
+  kid: string
+  jwks_url: string
+}
+
+export interface WebhookSigner {
+  sign: (body: Uint8Array | string) => string | Promise<string>
+  jwks: () => WebhookJwks | Promise<WebhookJwks>
+}
+
+export type WebhookSigningKeyInput =
+  | string
+  | Buffer
+  | NodeJsonWebKey
+  | KeyObject
+
+export interface Ed25519WebhookSignerOptions {
+  privateKey?: WebhookSigningKeyInput
+  kid?: string
+}
+
+const encoder = new TextEncoder()
+const defaultWebhookSigner = createEd25519WebhookSigner()
+
+export function createEd25519WebhookSigner(
+  options: Ed25519WebhookSignerOptions = {}
+): WebhookSigner {
+  const privateKey = options.privateKey
+    ? importPrivateKey(options.privateKey)
+    : generateKeyPairSync(`ed25519`).privateKey
+
+  if (privateKey.asymmetricKeyType !== `ed25519`) {
+    throw new Error(`Webhook signing key must be an Ed25519 private key`)
+  }
+
+  const publicJwk = buildPublicJwk(privateKey, options.kid)
+
+  return {
+    sign: (body) => signWebhookBody(privateKey, publicJwk.kid, body),
+    jwks: () => ({ keys: [{ ...publicJwk }] }),
+  }
+}
+
+export function getDefaultWebhookSigner(): WebhookSigner {
+  return defaultWebhookSigner
+}
+
+export async function webhookSigningMetadata(
+  signer: WebhookSigner,
+  streamRootUrl: string
+): Promise<WebhookSigningMetadata> {
+  const jwks = await signer.jwks()
+  const key = jwks.keys[0]
+  if (!key) {
+    throw new Error(`Webhook signer did not provide any public keys`)
+  }
+
+  return {
+    alg: `ed25519`,
+    kid: key.kid,
+    jwks_url: appendPathToUrl(streamRootUrl, `/__ds/jwks.json`),
+  }
+}
+
+export function signWebhookBody(
+  privateKey: KeyObject,
+  kid: string,
+  body: Uint8Array | string
+): string {
+  const timestamp = Math.floor(Date.now() / 1000)
+  const payload = bytesWithTimestamp(timestamp, body)
+  const signature = sign(null, payload, privateKey).toString(`base64url`)
+  return `t=${timestamp},kid=${kid},ed25519=${signature}`
+}
+
+export function bytesWithTimestamp(
+  timestamp: number | string,
+  body: Uint8Array | string
+): Buffer {
+  const prefix = encoder.encode(`${timestamp}.`)
+  const bodyBytes = typeof body === `string` ? encoder.encode(body) : body
+  return Buffer.concat([Buffer.from(prefix), Buffer.from(bodyBytes)])
+}
+
+function importPrivateKey(input: WebhookSigningKeyInput): KeyObject {
+  if (isKeyObject(input)) return input
+
+  if (typeof input === `string`) {
+    const trimmed = input.trim()
+    if (trimmed.startsWith(`{`)) {
+      return createPrivateKey({
+        key: JSON.parse(trimmed) as NodeJsonWebKey,
+        format: `jwk`,
+      })
+    }
+    return createPrivateKey(trimmed.replace(/\\n/g, `\n`))
+  }
+
+  if (Buffer.isBuffer(input)) {
+    return createPrivateKey(input)
+  }
+
+  return createPrivateKey({ key: input, format: `jwk` })
+}
+
+function isKeyObject(input: WebhookSigningKeyInput): input is KeyObject {
+  return (
+    typeof input === `object` && `type` in input && input.type === `private`
+  )
+}
+
+function buildPublicJwk(
+  privateKey: KeyObject,
+  kid: string | undefined
+): WebhookPublicJwk {
+  const exported = createPublicKey(privateKey).export({ format: `jwk` }) as {
+    kty?: string
+    crv?: string
+    x?: string
+  }
+
+  if (exported.kty !== `OKP` || exported.crv !== `Ed25519` || !exported.x) {
+    throw new Error(`Failed to export Ed25519 webhook signing key`)
+  }
+
+  return {
+    kty: `OKP`,
+    crv: `Ed25519`,
+    x: exported.x,
+    kid:
+      kid ??
+      deriveKeyId({
+        kty: exported.kty,
+        crv: exported.crv,
+        x: exported.x,
+      }),
+    use: `sig`,
+    alg: `EdDSA`,
+  }
+}
+
+function deriveKeyId(jwk: { kty: string; crv: string; x: string }): string {
+  const thumbprintInput = JSON.stringify({
+    crv: jwk.crv,
+    kty: jwk.kty,
+    x: jwk.x,
+  })
+  return `ds_${createHash(`sha256`).update(thumbprintInput).digest(`base64url`)}`
+}