@electric-ax/agents-server 0.3.0

package/src/schema-validation.ts

@@ -0,0 +1,15 @@
+import Ajv from 'ajv'
+import type { TSchema as TypeBoxSchema } from '@sinclair/typebox'
+import type { ValidateFunction } from 'ajv'
+
+const jsonBodyAjv = new Ajv({ allErrors: true })
+const schemaValidators = new WeakMap<TypeBoxSchema, ValidateFunction>()
+
+export function schemaValidator(schema: TypeBoxSchema): ValidateFunction {
+  let validate = schemaValidators.get(schema)
+  if (!validate) {
+    validate = jsonBodyAjv.compile(schema)
+    schemaValidators.set(schema, validate)
+  }
+  return validate
+}

package/src/server.ts

@@ -0,0 +1,374 @@
+import { createServer } from 'node:http'
+import { createServerAdapter } from '@whatwg-node/server'
+import { Agent } from 'undici'
+import {
+  appendPathToUrl,
+  createEntityRegistry,
+  createRuntimeHandler,
+} from '@electric-ax/agents-runtime'
+import { createDb, runMigrations } from './db/index.js'
+import { pathPrefixedSingleTenantDurableStreamsRoutingAdapter } from './routing/durable-streams-routing-adapter.js'
+import { ossServerRouter } from './routing/oss-server-router.js'
+import { startStandaloneAgentsRuntime } from './standalone-runtime.js'
+import { StreamClient, durableStreamsServiceUrl } from './stream-client.js'
+import { DEFAULT_TENANT_ID } from './tenant.js'
+import { serverLog } from './utils/log.js'
+import type { DrizzleDB, PgClient } from './db/index.js'
+import type { Server } from 'node:http'
+import type { DurableStreamTestServer } from '@durable-streams/server'
+import type { StreamFn } from '@mariozechner/pi-agent-core'
+import type {
+  AgentModel,
+  EntityRegistry,
+  RuntimeHandler,
+} from '@electric-ax/agents-runtime'
+import type { AuthenticateRequest } from './electric-agents-types.js'
+import type { EntityBridgeCoordinator } from './entity-bridge-manager.js'
+import type { DurableStreamsRoutingAdapter } from './routing/durable-streams-routing-adapter.js'
+import type { OssServerContext } from './routing/oss-server-router.js'
+import type { StartedStandaloneAgentsRuntime } from './standalone-runtime.js'
+import type { DurableStreamsBearerProvider } from './stream-client.js'
+
+const MOCK_AGENT_HANDLER_PATH = `/_electric/mock-agent-handler`
+
+export interface ElectricAgentsServerOptions {
+  service?: string
+  tenantId?: string
+  baseUrl?: string
+  durableStreamsUrl?: string
+  durableStreamsBearer?: DurableStreamsBearerProvider
+  durableStreamsRouting?: DurableStreamsRoutingAdapter
+  durableStreamsServer?: DurableStreamTestServer
+  port: number
+  host?: string
+  workingDirectory?: string
+  mockStreamFn?: StreamFn
+  postgresUrl: string
+  electricUrl?: string
+  electricSecret?: string
+  authenticateRequest?: AuthenticateRequest
+  /**
+   * Disabled by default. When set to a positive interval, periodically
+   * recovers expired dispatch claims and stale outstanding wakes.
+   */
+  dispatchRecoveryIntervalMs?: number
+  /**
+   * Age threshold for outstanding wakes recovered by the periodic loop.
+   * Defaults to dispatchRecoveryIntervalMs when periodic recovery is enabled.
+   */
+  staleOutstandingWakeAfterMs?: number
+}
+
+interface MockAgentBootstrap {
+  runtime: RuntimeHandler
+  registry: EntityRegistry
+}
+
+const MOCK_CHAT_MODEL: AgentModel = {
+  id: `mock-chat`,
+  name: `Mock Chat`,
+  api: `anthropic-messages`,
+  provider: `anthropic`,
+  baseUrl: `http://mock`,
+  reasoning: false,
+  input: [`text`],
+  cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 },
+  contextWindow: 200_000,
+  maxTokens: 4_096,
+}
+
+function createMockAgentBootstrap(options: {
+  agentServerUrl: string
+  workingDirectory?: string
+  streamFn: StreamFn
+}): MockAgentBootstrap {
+  const registry = createEntityRegistry()
+
+  registry.define(`chat`, {
+    description: `Mock chat entity for conformance and end-to-end tests.`,
+    handler: async (ctx) => {
+      ctx.useAgent({
+        systemPrompt: `You are a concise test assistant.`,
+        model: MOCK_CHAT_MODEL,
+        tools: [],
+        streamFn: options.streamFn,
+      })
+      await ctx.agent.run()
+    },
+  })
+
+  const runtime = createRuntimeHandler({
+    baseUrl: options.agentServerUrl,
+    serveEndpoint: appendPathToUrl(
+      options.agentServerUrl,
+      MOCK_AGENT_HANDLER_PATH
+    ),
+    registry,
+    subscriptionPathForType: (name) => `/${name}/*/main`,
+    idleTimeout: 100,
+  })
+
+  return { runtime, registry }
+}
+
+export class ElectricAgentsServer {
+  private server?: Server
+  private electricAgentsManager?: StartedStandaloneAgentsRuntime[`manager`]
+  private pgDb?: DrizzleDB
+  private pgClient?: PgClient
+  private entityBridgeManager?: EntityBridgeCoordinator
+  private mockAgentBootstrap?: MockAgentBootstrap
+  private _url?: string
+  private shuttingDown = false
+  private streamsAgent?: Agent
+  private standaloneRuntime?: StartedStandaloneAgentsRuntime
+
+  streamClient: StreamClient
+  readonly options: ElectricAgentsServerOptions
+
+  constructor(options: ElectricAgentsServerOptions) {
+    if (!options.durableStreamsUrl && !options.durableStreamsServer) {
+      throw new Error(
+        `Either durableStreamsUrl or durableStreamsServer is required`
+      )
+    }
+    this.options = options
+    this.streamClient = options.durableStreamsUrl
+      ? new StreamClient(
+          durableStreamsServiceUrl(options.durableStreamsUrl, this.tenantId),
+          { bearer: options.durableStreamsBearer }
+        )
+      : null!
+  }
+
+  get url(): string {
+    if (!this._url) {
+      throw new Error(`Server not started`)
+    }
+    return this._url
+  }
+
+  private get publicUrl(): string {
+    if (!this._url) {
+      throw new Error(`Server not started`)
+    }
+    return this.options.baseUrl?.replace(/\/+$/, ``) ?? this._url
+  }
+
+  private get tenantId(): string {
+    return this.options.service ?? this.options.tenantId ?? DEFAULT_TENANT_ID
+  }
+
+  async start(): Promise<string> {
+    if (this.server) {
+      throw new Error(`Server already started`)
+    }
+
+    try {
+      if (
+        this.options.durableStreamsServer &&
+        !this.options.durableStreamsUrl
+      ) {
+        serverLog.info(`[agent-server] starting durable streams server...`)
+        const streamsUrl = await this.options.durableStreamsServer.start()
+        serverLog.info(
+          `[agent-server] durable streams server started at ${streamsUrl}`
+        )
+        this.options.durableStreamsUrl = streamsUrl
+        this.streamClient = new StreamClient(
+          durableStreamsServiceUrl(streamsUrl, this.tenantId),
+          { bearer: this.options.durableStreamsBearer }
+        )
+      }
+
+      this.streamsAgent = new Agent({
+        keepAliveTimeout: 60_000,
+        keepAliveMaxTimeout: 600_000,
+        connections: 256,
+        pipelining: 1,
+        bodyTimeout: 0,
+        headersTimeout: 0,
+      })
+
+      serverLog.info(`[agent-server] running migrations...`)
+      await runMigrations(this.options.postgresUrl)
+      serverLog.info(`[agent-server] migrations complete`)
+      const { db, client } = createDb(this.options.postgresUrl)
+      this.pgDb = db
+      this.pgClient = client
+
+      this.standaloneRuntime = await startStandaloneAgentsRuntime({
+        service: this.tenantId,
+        db,
+        pgClient: client,
+        streamClient: this.streamClient,
+        electricUrl: this.options.electricUrl,
+        electricSecret: this.options.electricSecret,
+      })
+      this.electricAgentsManager = this.standaloneRuntime.manager
+      this.entityBridgeManager = this.standaloneRuntime.entityBridgeManager
+
+      const serverAdapter = createServerAdapter((request) =>
+        this.handleRequest(request)
+      )
+      const server = createServer(serverAdapter)
+      this.server = server
+
+      const host = this.options.host ?? `127.0.0.1`
+      await this.listen(server, host)
+
+      if (this.options.mockStreamFn) {
+        this.mockAgentBootstrap = createMockAgentBootstrap({
+          agentServerUrl: this.publicUrl,
+          workingDirectory: this.options.workingDirectory,
+          streamFn: this.options.mockStreamFn,
+        })
+        await this.mockAgentBootstrap.runtime.registerTypes()
+        serverLog.info(
+          `[agent-server] mock chat agent registered at ${MOCK_AGENT_HANDLER_PATH}`
+        )
+      }
+
+      return this.url
+    } catch (err) {
+      await this.stop().catch(() => {})
+      throw err
+    }
+  }
+
+  async stop(): Promise<void> {
+    this.shuttingDown = true
+
+    if (this.server) {
+      const server = this.server
+      await new Promise<void>((resolve) => {
+        server.close(() => resolve())
+        server.closeIdleConnections?.()
+        server.closeAllConnections?.()
+      })
+      this.server = undefined
+      this._url = undefined
+    }
+
+    if (this.mockAgentBootstrap) {
+      this.mockAgentBootstrap.runtime.abortWakes()
+      await Promise.race([
+        this.mockAgentBootstrap.runtime.drainWakes().catch(() => {}),
+        new Promise<void>((resolve) => setTimeout(resolve, 5_000)),
+      ])
+      this.mockAgentBootstrap = undefined
+    }
+
+    if (this.standaloneRuntime) {
+      await this.standaloneRuntime.stop()
+      this.standaloneRuntime = undefined
+      this.entityBridgeManager = undefined
+      this.electricAgentsManager = undefined
+    }
+
+    if (this.options.durableStreamsServer) {
+      await this.options.durableStreamsServer.stop()
+    }
+
+    if (this.pgClient) {
+      await this.pgClient.end()
+      this.pgClient = undefined
+      this.pgDb = undefined
+    }
+
+    if (this.streamsAgent) {
+      await this.streamsAgent.close().catch(() => {})
+      this.streamsAgent = undefined
+    }
+
+    this.shuttingDown = false
+  }
+
+  private async listen(server: Server, host: string): Promise<void> {
+    await new Promise<void>((resolve, reject) => {
+      const onError = (err: Error): void => {
+        server.off?.(`listening`, onListening)
+        reject(err)
+      }
+      const onListening = (): void => {
+        server.off?.(`error`, onError)
+        resolve()
+      }
+
+      server.once?.(`error`, onError) ?? server.on(`error`, onError)
+      server.listen(this.options.port, host, onListening)
+    })
+
+    const addr = server.address()
+    if (typeof addr === `string`) {
+      this._url = addr
+    } else if (addr) {
+      const resolvedHost = host === `0.0.0.0` ? `127.0.0.1` : host
+      this._url = `http://${resolvedHost}:${addr.port}`
+    } else {
+      throw new Error(`Could not determine server address`)
+    }
+  }
+
+  private async handleRequest(request: Request): Promise<Response> {
+    if (
+      !this._url ||
+      !this.standaloneRuntime ||
+      !this.electricAgentsManager ||
+      !this.entityBridgeManager ||
+      !this.pgDb ||
+      !this.streamsAgent ||
+      !this.options.durableStreamsUrl
+    ) {
+      return new Response(null, { status: 503 })
+    }
+
+    return await ossServerRouter.fetch(
+      request as Parameters<typeof ossServerRouter.fetch>[0],
+      await this.buildTenantContext(request)
+    )
+  }
+
+  private async buildTenantContext(
+    request: Request
+  ): Promise<OssServerContext> {
+    if (
+      !this.standaloneRuntime ||
+      !this.electricAgentsManager ||
+      !this.entityBridgeManager ||
+      !this.pgDb ||
+      !this.streamsAgent ||
+      !this.options.durableStreamsUrl
+    ) {
+      throw new Error(`agents-server runtime is not started`)
+    }
+
+    return {
+      service: this.tenantId,
+      authenticatedUser:
+        (await this.options.authenticateRequest?.(request)) ?? undefined,
+      publicUrl: this.publicUrl,
+      localUrl: this._url,
+      durableStreamsUrl: this.options.durableStreamsUrl,
+      durableStreamsBearer: this.options.durableStreamsBearer,
+      durableStreamsRouting:
+        this.options.durableStreamsRouting ??
+        pathPrefixedSingleTenantDurableStreamsRoutingAdapter,
+      durableStreamsDispatcher: this.streamsAgent,
+      electricUrl: this.options.electricUrl,
+      electricSecret: this.options.electricSecret,
+      ownAgentHandlerPaths: this.mockAgentBootstrap
+        ? [MOCK_AGENT_HANDLER_PATH]
+        : undefined,
+      pgDb: this.pgDb,
+      entityManager: this.electricAgentsManager,
+      streamClient: this.streamClient,
+      runtime: this.standaloneRuntime.runtime,
+      entityBridgeManager: this.entityBridgeManager,
+      isShuttingDown: () => this.shuttingDown,
+      mockAgent: this.mockAgentBootstrap
+        ? { runtime: this.mockAgentBootstrap.runtime }
+        : undefined,
+    }
+  }
+}

package/src/standalone-runtime.ts

@@ -0,0 +1,188 @@
+import { randomUUID } from 'node:crypto'
+import { PostgresRegistry } from './entity-registry.js'
+import { EntityBridgeManager } from './entity-bridge-manager.js'
+import { serverLog } from './utils/log.js'
+import { ElectricAgentsTenantRuntime } from './runtime.js'
+import { Scheduler } from './scheduler.js'
+import { StreamClient, durableStreamsServiceUrl } from './stream-client.js'
+import { TagStreamOutboxDrainer } from './tag-stream-outbox-drainer.js'
+import { DEFAULT_TENANT_ID } from './tenant.js'
+import { WakeRegistry } from './wake-registry.js'
+import type { DrizzleDB, PgClient } from './db/index.js'
+import type { EntityManager } from './entity-manager.js'
+import type { EntityBridgeCoordinator } from './entity-bridge-manager.js'
+import type { CronTickPayload, DelayedSendPayload } from './scheduler.js'
+import type { DurableStreamsBearerProvider } from './stream-client.js'
+
+export interface StandaloneAgentsRuntimeOptions {
+  service?: string
+  tenantId?: string
+  db: DrizzleDB
+  pgClient: PgClient
+  durableStreamsUrl?: string
+  durableStreamsBearer?: DurableStreamsBearerProvider
+  streamClient?: StreamClient
+  electricUrl?: string
+  electricSecret?: string
+  instanceId?: string
+  wakeRegistry?: WakeRegistry
+  startWakeRegistry?: boolean
+  startScheduler?: boolean
+  startTagStreamOutboxDrainer?: boolean
+  startEntityBridgeManager?: boolean
+  rehydrateOnStart?: boolean
+  entityBridgeManager?: EntityBridgeCoordinator
+}
+
+export interface StartedStandaloneAgentsRuntime {
+  serviceId: string
+  service: string
+  db: DrizzleDB
+  pgClient: PgClient
+  streamClient: StreamClient
+  registry: PostgresRegistry
+  wakeRegistry: WakeRegistry
+  runtime: ElectricAgentsTenantRuntime
+  manager: EntityManager
+  scheduler: Scheduler
+  entityBridgeManager: EntityBridgeCoordinator
+  tagStreamOutboxDrainer: TagStreamOutboxDrainer
+  stop: () => Promise<void>
+}
+
+export async function startStandaloneAgentsRuntime(
+  options: StandaloneAgentsRuntimeOptions
+): Promise<StartedStandaloneAgentsRuntime> {
+  const serviceId = options.service ?? options.tenantId ?? DEFAULT_TENANT_ID
+  const streamClient =
+    options.streamClient ??
+    (options.durableStreamsUrl
+      ? new StreamClient(
+          durableStreamsServiceUrl(options.durableStreamsUrl, serviceId),
+          { bearer: options.durableStreamsBearer }
+        )
+      : undefined)
+  if (!streamClient) {
+    throw new Error(`Either durableStreamsUrl or streamClient is required`)
+  }
+
+  const registry = new PostgresRegistry(options.db, serviceId)
+  const wakeRegistry =
+    options.wakeRegistry ?? new WakeRegistry(options.db, serviceId)
+  let runtime: ElectricAgentsTenantRuntime
+  const scheduler = new Scheduler({
+    pgClient: options.pgClient,
+    instanceId: options.instanceId ?? randomUUID(),
+    tenantId: serviceId,
+    executors: {
+      delayed_send: async (payload: DelayedSendPayload, taskId: number) => {
+        await runtime.executeDelayedSend(payload, taskId)
+      },
+      cron_tick: async (payload: CronTickPayload, tickNumber: number) => {
+        await runtime.executeCronTick(payload, tickNumber)
+      },
+    },
+  })
+  const entityBridgeManager =
+    options.entityBridgeManager ??
+    new EntityBridgeManager(
+      registry,
+      streamClient,
+      options.electricUrl,
+      options.electricSecret,
+      serviceId
+    )
+  const tagStreamOutboxDrainer = new TagStreamOutboxDrainer(
+    registry,
+    streamClient
+  )
+
+  runtime = new ElectricAgentsTenantRuntime({
+    service: serviceId,
+    db: options.db,
+    registry,
+    streamClient,
+    wakeRegistry,
+    scheduler,
+    entityBridgeManager,
+    stopWakeRegistryOnShutdown: options.wakeRegistry ? false : true,
+  })
+
+  const startWakeRegistry = options.startWakeRegistry ?? true
+  const startScheduler = options.startScheduler ?? true
+  const startTagStreamOutboxDrainer =
+    options.startTagStreamOutboxDrainer ?? true
+  const startEntityBridgeManager = options.startEntityBridgeManager ?? true
+  const rehydrateOnStart = options.rehydrateOnStart ?? true
+  let entityBridgeManagerStarted = false
+  let tagStreamOutboxDrainerStarted = false
+  let schedulerStarted = false
+  let stopped = false
+
+  const stop = async (): Promise<void> => {
+    if (stopped) return
+    stopped = true
+    if (schedulerStarted) {
+      await scheduler.stop()
+      schedulerStarted = false
+    }
+    if (tagStreamOutboxDrainerStarted) {
+      await tagStreamOutboxDrainer.stop()
+      tagStreamOutboxDrainerStarted = false
+    }
+    if (entityBridgeManagerStarted) {
+      await entityBridgeManager.stop()
+      entityBridgeManagerStarted = false
+    }
+    await runtime.stop()
+  }
+
+  try {
+    if (startWakeRegistry) {
+      serverLog.info(`[agent-server] rebuilding wake registry...`)
+      await runtime.manager.rebuildWakeRegistry(
+        options.electricUrl,
+        options.electricSecret
+      )
+    }
+    if (rehydrateOnStart) {
+      serverLog.info(`[agent-server] rehydrating cron schedules...`)
+      await runtime.rehydrateCronSchedules()
+    }
+    if (startEntityBridgeManager) {
+      serverLog.info(`[agent-server] starting entity bridge manager...`)
+      await entityBridgeManager.start()
+      entityBridgeManagerStarted = true
+    }
+    if (startTagStreamOutboxDrainer) {
+      serverLog.info(`[agent-server] starting tag stream outbox drainer...`)
+      tagStreamOutboxDrainer.start()
+      tagStreamOutboxDrainerStarted = true
+    }
+    if (startScheduler) {
+      serverLog.info(`[agent-server] starting scheduler...`)
+      schedulerStarted = true
+      await scheduler.start()
+      serverLog.info(`[agent-server] scheduler started`)
+    }
+  } catch (error) {
+    await stop().catch(() => {})
+    throw error
+  }
+
+  return {
+    serviceId,
+    service: serviceId,
+    db: options.db,
+    pgClient: options.pgClient,
+    streamClient,
+    registry,
+    wakeRegistry,
+    runtime,
+    manager: runtime.manager,
+    scheduler,
+    entityBridgeManager,
+    tagStreamOutboxDrainer,
+    stop,
+  }
+}