@electric-ax/agents-server 0.4.5 → 0.4.7

package/src/entity-registry.ts

@@ -15,6 +15,7 @@ import {
   assertEntityStatus,
   assertRunnerAdminStatus,
   assertRunnerKind,
+  isTerminalEntityStatus,
 } from './electric-agents-types.js'
 import { DEFAULT_TENANT_ID } from './tenant.js'
 import type { DrizzleDB } from './db/index.js'
@@ -360,11 +361,17 @@ export class PostgresRegistry {
     input: MaterializeHeartbeatClaimInput
   ): Promise<void> {
     const heartbeatAt = input.heartbeatAt ?? new Date()
+    // Only touch leaseExpiresAt when the caller explicitly provides one.
+    // The lease was set at materializeActiveClaim time from the upstream
+    // lease_ttl_ms and remains the authoritative expiry; heartbeats are
+    // alive-pings, not lease extensions.
     await this.db
       .update(consumerClaims)
       .set({
         lastHeartbeatAt: heartbeatAt,
-        leaseExpiresAt: input.leaseExpiresAt ?? null,
+        ...(input.leaseExpiresAt !== undefined
+          ? { leaseExpiresAt: input.leaseExpiresAt }
+          : {}),
         updatedAt: heartbeatAt,
       })
       .where(
@@ -378,7 +385,7 @@ export class PostgresRegistry {
 
   async materializeReleasedClaim(
     input: MaterializeReleasedClaimInput
-  ): Promise<ConsumerClaim | null> {
+  ): Promise<{ claim: ConsumerClaim | null; entityCleared: boolean }> {
     const releasedAt = input.releasedAt ?? new Date()
     const rows = await this.db
       .update(consumerClaims)
@@ -398,8 +405,13 @@ export class PostgresRegistry {
       .returning()
 
     const claim = rows[0] ? this.rowToConsumerClaim(rows[0]) : null
+    let entityCleared = false
     if (claim) {
-      await this.db
+      // entityCleared distinguishes "we were the active dispatch and now it's
+      // empty" from "a newer claim was already active for this entity." The
+      // WHERE clause matches our (consumerId, epoch) so an evicted-by-newer
+      // case correctly returns zero rows.
+      const cleared = await this.db
         .update(entityDispatchState)
         .set({
           activeConsumerId: null,
@@ -419,8 +431,10 @@ export class PostgresRegistry {
             eq(entityDispatchState.activeEpoch, input.epoch)
           )
         )
+        .returning({ entityUrl: entityDispatchState.entityUrl })
+      entityCleared = cleared.length > 0
     }
-    return claim
+    return { claim, entityCleared }
   }
 
   async getActiveClaimsForRunner(
@@ -713,10 +727,13 @@ export class PostgresRegistry {
   }
 
   async updateStatus(entityUrl: string, status: EntityStatus): Promise<void> {
-    const whereClause =
-      status === `stopped`
-        ? this.entityWhere(entityUrl)
-        : and(this.entityWhere(entityUrl), ne(entities.status, `stopped`))
+    const whereClause = isTerminalEntityStatus(status)
+      ? this.entityWhere(entityUrl)
+      : and(
+          this.entityWhere(entityUrl),
+          ne(entities.status, `stopped`),
+          ne(entities.status, `killed`)
+        )
 
     await this.db
       .update(entities)
@@ -727,21 +744,41 @@ export class PostgresRegistry {
   async updateStatusWithTxid(
     entityUrl: string,
     status: EntityStatus
-  ): Promise<number> {
+  ): Promise<number | null> {
     return await this.db.transaction(async (tx) => {
-      const whereClause =
-        status === `stopped`
-          ? this.entityWhere(entityUrl)
-          : and(this.entityWhere(entityUrl), ne(entities.status, `stopped`))
-
-      await tx
+      const rows = await tx
         .update(entities)
         .set({ status, updatedAt: Date.now() })
-        .where(whereClause)
-      const result = await tx.execute(
-        sql`SELECT pg_current_xact_id()::xid::text AS txid`
-      )
-      return parseInt((result[0] as { txid: string }).txid)
+        .where(
+          and(
+            this.entityWhere(entityUrl),
+            ne(entities.status, `stopped`),
+            ne(entities.status, `killed`)
+          )
+        )
+        .returning({
+          txid: sql<string>`pg_current_xact_id()::xid::text`,
+        })
+      return rows[0] ? parseInt(rows[0].txid) : null
+    })
+  }
+
+  async touchEntityWithTxid(entityUrl: string): Promise<number | null> {
+    return await this.db.transaction(async (tx) => {
+      const rows = await tx
+        .update(entities)
+        .set({ updatedAt: Date.now() })
+        .where(
+          and(
+            eq(entities.url, entityUrl),
+            ne(entities.status, `stopped`),
+            ne(entities.status, `killed`)
+          )
+        )
+        .returning({
+          txid: sql<string>`pg_current_xact_id()::xid::text`,
+        })
+      return rows[0] ? parseInt(rows[0].txid) : null
     })
   }
 

package/src/entrypoint-lib.ts

@@ -139,6 +139,10 @@ export function resolveElectricAgentsEntrypointOptions(
     `ELECTRIC_URL`,
   ])
   const electricSecret = readEnv(env, [`ELECTRIC_AGENTS_ELECTRIC_SECRET`])
+  const webhookSigningKey = readEnv(env, [
+    `ELECTRIC_AGENTS_WEBHOOK_SIGNING_PRIVATE_KEY`,
+    `WEBHOOK_SIGNING_PRIVATE_KEY`,
+  ])
   const baseUrl = readEnv(env, [`ELECTRIC_AGENTS_BASE_URL`, `BASE_URL`])
   return {
     service: readEnv(env, [`ELECTRIC_AGENTS_SERVICE`, `SERVICE`]),
@@ -153,6 +157,7 @@ export function resolveElectricAgentsEntrypointOptions(
       ? validateUrl(`Electric URL`, electricUrl)
       : undefined,
     electricSecret,
+    webhookSigningKey,
     host: readEnv(env, [`ELECTRIC_AGENTS_HOST`, `HOST`]) ?? DEFAULT_HOST,
     port: readPort(env),
     workingDirectory:

package/src/index.ts

@@ -15,6 +15,14 @@ export type {
   SubscriptionResponse,
   SubscriptionStreamInfo,
 } from './stream-client.js'
+export {
+  assertEntitySignal,
+  assertEntityStatus,
+  expectedSignalStatus,
+  isTerminalEntityStatus,
+  rejectsNormalWrites,
+  toPublicEntity,
+} from './electric-agents-types.js'
 export type {
   AuthenticateRequest,
   ConsumerClaim,
@@ -32,6 +40,18 @@ export type {
   RunnerLiveness,
   SourceStreamOffset,
   WakeNotificationRow,
+  ElectricAgentsEntity,
+  ElectricAgentsEntityRow,
+  ElectricAgentsEntityType,
+  EntityStatus,
+  EntitySignal,
+  PublicElectricAgentsEntity,
+  EntityListFilter,
+  RegisterEntityTypeRequest,
+  SendRequest,
+  SignalRequest,
+  SignalResponse,
+  TypedSpawnRequest,
 } from './electric-agents-types.js'
 export type { Principal, PrincipalKind } from './principal.js'
 export { globalRouter } from './routing/global-router.js'
@@ -46,6 +66,19 @@ export type {
   DurableStreamsRoutingAdapter,
   DurableStreamsRoutingInput,
 } from './routing/durable-streams-routing-adapter.js'
+export {
+  createEd25519WebhookSigner,
+  getDefaultWebhookSigner,
+  webhookSigningMetadata,
+} from './webhook-signing.js'
+export type {
+  Ed25519WebhookSignerOptions,
+  WebhookJwks,
+  WebhookPublicJwk,
+  WebhookSigner,
+  WebhookSigningKeyInput,
+  WebhookSigningMetadata,
+} from './webhook-signing.js'
 export type { EntityBridgeCoordinator } from './entity-bridge-manager.js'
 export {
   DEFAULT_TENANT_ID,

package/src/routing/context.ts

@@ -1,4 +1,5 @@
 import type { Agent } from 'undici'
+import type { WebhookSignatureVerifierConfig } from '@electric-ax/agents-runtime'
 import type { DrizzleDB } from '../db/index.js'
 import type { EntityBridgeCoordinator } from '../entity-bridge-manager.js'
 import type { EntityManager } from '../entity-manager.js'
@@ -7,6 +8,7 @@ import type { StreamClient } from '../stream-client.js'
 import type { DurableStreamsRoutingAdapter } from './durable-streams-routing-adapter.js'
 import type { Principal } from '../principal.js'
 import type { DurableStreamsBearerProvider } from '../stream-client.js'
+import type { WebhookSigner } from '../webhook-signing.js'
 
 /**
  * Per-request tenant context passed through every router and handler.
@@ -24,6 +26,10 @@ export interface TenantContext {
   durableStreamsBearer?: DurableStreamsBearerProvider
   durableStreamsRouting?: DurableStreamsRoutingAdapter
   durableStreamsDispatcher: Agent
+  durableStreamsWebhookSignature?:
+    | false
+    | Partial<WebhookSignatureVerifierConfig>
+  webhookSigner?: WebhookSigner
   electricUrl?: string
   electricSecret?: string
   ownAgentHandlerPaths?: ReadonlyArray<string>

package/src/routing/dispatch-policy.ts

@@ -217,6 +217,12 @@ export async function assertDispatchPolicyAllowed(
   }
 }
 
+export function shouldLinkDispatchBeforeInitialMessage(
+  policy: DispatchPolicy | undefined
+): boolean {
+  return policy?.targets[0] !== undefined
+}
+
 export async function linkEntityDispatchSubscription(
   ctx: TenantContext,
   entity: ElectricAgentsEntity

package/src/routing/durable-streams-router.ts

@@ -15,10 +15,15 @@ import {
 import { validateBody } from './schema.js'
 import { rewriteLoopbackWebhookUrl } from '../utils/webhook-url.js'
 import { forwardFetchRequest } from '../utils/server-utils.js'
+import {
+  getDefaultWebhookSigner,
+  webhookSigningMetadata,
+} from '../webhook-signing.js'
 import { resolveDurableStreamsRoutingAdapter } from './durable-streams-routing-adapter.js'
 import type { IRequest, RouterType } from 'itty-router'
 import type { TenantContext } from './context.js'
 import type { DurableStreamsRoutingAdapter } from './durable-streams-routing-adapter.js'
+import type { WebhookSigner } from '../webhook-signing.js'
 
 const subscriptionProxyBodySchema = Type.Object(
   {
@@ -81,6 +86,7 @@ for (const action of subscriptionControlActions) {
     subscriptionAction(action)
   )
 }
+durableStreamsRouter.get(`/__ds/jwks.json`, webhookJwks)
 durableStreamsRouter.all(`/__ds`, controlPassThrough)
 durableStreamsRouter.all(`/__ds/*`, controlPassThrough)
 durableStreamsRouter.post(`*`, streamAppend)
@@ -94,13 +100,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
+}
+
 async function forwardToDurableStreams(
   ctx: TenantContext,
   request: IRequest,
@@ -178,12 +193,12 @@ function rewriteSubscriptionBodyForBackend(
   }
 }
 
-function rewriteSubscriptionResponseForClient(
+async function rewriteSubscriptionResponseForClient(
   bytes: Uint8Array,
   response: Response,
-  service: string,
+  ctx: TenantContext,
   routingAdapter: DurableStreamsRoutingAdapter
-): Uint8Array {
+): Promise<Uint8Array> {
   if (!response.headers.get(`content-type`)?.includes(`application/json`)) {
     return bytes
   }
@@ -192,14 +207,14 @@ function rewriteSubscriptionResponseForClient(
 
   if (typeof payload.pattern === `string`) {
     payload.pattern = routingAdapter.toRuntimeStreamPath(
-      service,
+      ctx.service,
       payload.pattern
     )
   }
   if (Array.isArray(payload.streams)) {
     payload.streams = payload.streams.map((stream) => {
       if (typeof stream === `string`) {
-        return routingAdapter.toRuntimeStreamPath(service, stream)
+        return routingAdapter.toRuntimeStreamPath(ctx.service, stream)
       }
       if (
         stream &&
@@ -209,7 +224,7 @@ function rewriteSubscriptionResponseForClient(
         return {
           ...(stream as Record<string, unknown>),
           path: routingAdapter.toRuntimeStreamPath(
-            service,
+            ctx.service,
             (stream as Record<string, string>).path
           ),
         }
@@ -219,26 +234,43 @@ function rewriteSubscriptionResponseForClient(
   }
   if (typeof payload.wake_stream === `string`) {
     payload.wake_stream = routingAdapter.toRuntimeStreamPath(
-      service,
+      ctx.service,
       payload.wake_stream
     )
   }
   if (typeof payload.stream === `string`) {
-    payload.stream = routingAdapter.toRuntimeStreamPath(service, payload.stream)
+    payload.stream = routingAdapter.toRuntimeStreamPath(
+      ctx.service,
+      payload.stream
+    )
   }
   if (Array.isArray(payload.acks)) {
     payload.acks = payload.acks.map((ack) => {
       if (!ack || typeof ack !== `object`) return ack
       const next = { ...(ack as Record<string, unknown>) }
       if (typeof next.stream === `string`) {
-        next.stream = routingAdapter.toRuntimeStreamPath(service, next.stream)
+        next.stream = routingAdapter.toRuntimeStreamPath(
+          ctx.service,
+          next.stream
+        )
       }
       if (typeof next.path === `string`) {
-        next.path = routingAdapter.toRuntimeStreamPath(service, next.path)
+        next.path = routingAdapter.toRuntimeStreamPath(ctx.service, next.path)
       }
       return next
     })
   }
+  if (
+    payload.webhook &&
+    typeof payload.webhook === `object` &&
+    !Array.isArray(payload.webhook)
+  ) {
+    const webhook = payload.webhook as Record<string, unknown>
+    webhook.signing = await webhookSigningMetadata(
+      resolveWebhookSigner(ctx),
+      ctx.publicUrl
+    )
+  }
 
   return new TextEncoder().encode(JSON.stringify(payload))
 }
@@ -269,6 +301,10 @@ function subscriptionRoutingAdapter(
   )
 }
 
+function resolveWebhookSigner(ctx: TenantContext): WebhookSigner {
+  return ctx.webhookSigner ?? getDefaultWebhookSigner()
+}
+
 async function rewriteSubscriptionRequestBody(
   request: IRequest,
   ctx: TenantContext,
@@ -334,10 +370,10 @@ async function forwardSubscriptionRequest(
   let responseBytes: Uint8Array = upstream.body
     ? new Uint8Array(await upstream.arrayBuffer())
     : new Uint8Array()
-  responseBytes = rewriteSubscriptionResponseForClient(
+  responseBytes = await rewriteSubscriptionResponseForClient(
     responseBytes,
     upstream,
-    ctx.service,
+    ctx,
     routingAdapter
   )
   return {
@@ -530,6 +566,19 @@ async function controlPassThrough(
   return responseFromUpstream(upstream)
 }
 
+async function webhookJwks(
+  _request: IRequest,
+  ctx: TenantContext
+): Promise<Response> {
+  return new Response(JSON.stringify(await resolveWebhookSigner(ctx).jwks()), {
+    status: 200,
+    headers: {
+      'content-type': `application/jwk-set+json`,
+      'cache-control': `public, max-age=300`,
+    },
+  })
+}
+
 async function streamAppend(
   request: IRequest,
   ctx: TenantContext

package/src/routing/entities-router.ts

@@ -18,6 +18,7 @@ import {
   backfillEntityDispatchPolicy,
   linkEntityDispatchSubscription,
   resolveEffectiveDispatchPolicyForSpawn,
+  shouldLinkDispatchBeforeInitialMessage,
   unlinkEntityDispatchSubscription,
 } from './dispatch-policy.js'
 import { routeBody, withSchema } from './schema.js'
@@ -133,6 +134,22 @@ const setTagBodySchema = Type.Object({
   value: Type.String(),
 })
 
+const entitySignalSchema = Type.Union([
+  Type.Literal(`SIGINT`),
+  Type.Literal(`SIGHUP`),
+  Type.Literal(`SIGTERM`),
+  Type.Literal(`SIGKILL`),
+  Type.Literal(`SIGSTOP`),
+  Type.Literal(`SIGCONT`),
+  Type.Literal(`SIGUSR`),
+])
+
+const signalBodySchema = Type.Object({
+  signal: entitySignalSchema,
+  reason: Type.Optional(Type.String()),
+  payload: Type.Optional(Type.Unknown()),
+})
+
 const scheduleBodySchema = Type.Union([
   Type.Object({
     scheduleType: Type.Literal(`cron`),
@@ -161,6 +178,7 @@ type SendBody = Static<typeof sendBodySchema>
 type InboxMessageBody = Static<typeof inboxMessageBodySchema>
 type ForkBody = Static<typeof forkBodySchema>
 type SetTagBody = Static<typeof setTagBodySchema>
+type SignalBody = Static<typeof signalBodySchema>
 type ScheduleBody = Static<typeof scheduleBodySchema>
 type EntitiesRegisterBody = Static<typeof entitiesRegisterBodySchema>
 
@@ -187,6 +205,12 @@ entitiesRouter.put(
 entitiesRouter.get(`/:type/:instanceId`, withExistingEntity, getEntity)
 entitiesRouter.head(`/:type/:instanceId`, withExistingEntity, headEntity)
 entitiesRouter.delete(`/:type/:instanceId`, withExistingEntity, killEntity)
+entitiesRouter.post(
+  `/:type/:instanceId/signal`,
+  withExistingEntity,
+  withSchema(signalBodySchema),
+  signalEntity
+)
 entitiesRouter.post(
   `/:type/:instanceId/send`,
   withExistingEntity,
@@ -614,13 +638,21 @@ async function spawnEntity(
     wake: parsed.wake,
     created_by: principal.url,
   })
-  await linkEntityDispatchSubscription(ctx, entity)
+  const linkBeforeInitialMessage =
+    parsed.initialMessage !== undefined &&
+    shouldLinkDispatchBeforeInitialMessage(dispatchPolicy)
+  if (linkBeforeInitialMessage) {
+    await linkEntityDispatchSubscription(ctx, entity)
+  }
   if (parsed.initialMessage !== undefined) {
     await ctx.entityManager.send(entity.url, {
       from: principal.url,
       payload: parsed.initialMessage,
     })
   }
+  if (!linkBeforeInitialMessage) {
+    await linkEntityDispatchSubscription(ctx, entity)
+  }
 
   return json(
     { ...toPublicEntity(entity), txid: entity.txid },
@@ -655,3 +687,27 @@ async function killEntity(
   ctx.runtime.claimWriteTokens.clearStream(ctx.service, entity.streams.main)
   return json(result)
 }
+
+async function signalEntity(
+  request: AgentsRouteRequest,
+  ctx: TenantContext
+): Promise<Response> {
+  const principalMutationError = rejectPrincipalEntityMutation(
+    request,
+    `signaled`
+  )
+  if (principalMutationError) return principalMutationError
+
+  const parsed = routeBody<SignalBody>(request)
+  const { entityUrl, entity } = requireExistingEntityRoute(request)
+  const result = await ctx.entityManager.signal(entityUrl, {
+    signal: parsed.signal,
+    reason: parsed.reason,
+    payload: parsed.payload,
+  })
+  if (result.new_state === `stopped` || result.new_state === `killed`) {
+    await unlinkEntityDispatchSubscription(ctx, entity)
+    ctx.runtime.claimWriteTokens.clearStream(ctx.service, entity.streams.main)
+  }
+  return json(result)
+}