feat(server-*): extend server capability, routing, heartbeat, and better type

packages/server-runtime/src/index.ts

@@ -1,12 +1,24 @@
-import type { WebSocketEvent } from '@proj-airi/server-shared/types'
+import type { MetadataEventSource, WebSocketEvent } from '@proj-airi/server-shared/types'
 
+import type {
+  RouteContext,
+  RouteDecision,
+  RouteMiddleware,
+  RoutingPolicy,
+} from './middlewares'
 import type { AuthenticatedPeer, Peer } from './types'
 
 import { availableLogLevelStrings, Format, LogLevelString, logLevelStringToLogLevelMap, useLogg } from '@guiiai/logg'
-import { WebSocketEventSource } from '@proj-airi/server-shared/types'
+import { MessageHeartbeat, MessageHeartbeatMark, WebSocketEventSource } from '@proj-airi/server-shared/types'
 import { defineWebSocketHandler, H3 } from 'h3'
 
 import { optionOrEnv } from './config'
+import {
+  collectDestinations,
+  createPolicyMiddleware,
+  isDevtoolsPeer,
+  matchesDestinations,
+} from './middlewares'
 
 // pre-stringified responses
 const RESPONSES = {
@@ -14,6 +26,8 @@ const RESPONSES = {
   notAuthenticated: JSON.stringify({ type: 'error', data: { message: 'not authenticated' }, source: WebSocketEventSource.Server } satisfies WebSocketEvent),
 }
 
+const DEFAULT_HEARTBEAT_TTL_MS = 60_000
+
 // helper send function
 function send(peer: Peer, event: WebSocketEvent<Record<string, unknown>> | string) {
   peer.send(typeof event === 'string' ? event : JSON.stringify(event))
@@ -27,6 +41,14 @@ export function setupApp(options?: {
     app?: { level?: LogLevelString, format?: Format }
     websocket?: { level?: LogLevelString, format?: Format }
   }
+  routing?: {
+    middleware?: RouteMiddleware[]
+    allowBypass?: boolean
+    policy?: RoutingPolicy
+  }
+  heartbeat?: {
+    readTimeout?: number
+  }
 }): H3 {
   const authToken = optionOrEnv(options?.auth?.token, 'AUTHENTICATION_TOKEN', '')
 
@@ -44,6 +66,32 @@ export function setupApp(options?: {
 
   const peers = new Map<string, AuthenticatedPeer>()
   const peersByModule = new Map<string, Map<number | undefined, AuthenticatedPeer>>()
+  const heartbeatTtlMs = options?.heartbeat?.readTimeout ?? DEFAULT_HEARTBEAT_TTL_MS
+  const routingMiddleware = [
+    ...(options?.routing?.policy ? [createPolicyMiddleware(options.routing.policy)] : []),
+    ...(options?.routing?.middleware ?? []),
+  ]
+
+  setInterval(() => {
+    const now = Date.now()
+    for (const [id, peerInfo] of peers.entries()) {
+      if (!peerInfo.lastHeartbeatAt) {
+        continue
+      }
+
+      if (now - peerInfo.lastHeartbeatAt > heartbeatTtlMs) {
+        logger.withFields({ peer: id, peerName: peerInfo.name }).debug('heartbeat expired, dropping peer')
+        try {
+          (peerInfo.peer as Peer & { close?: () => void }).close?.()
+        }
+        catch (error) {
+          logger.withFields({ peer: id, peerName: peerInfo.name }).withError(error as Error).debug('failed to close expired peer')
+        }
+        peers.delete(id)
+        unregisterModulePeer(peerInfo)
+      }
+    }
+  }, Math.max(5_000, Math.floor(heartbeatTtlMs / 2)))
 
   function registerModulePeer(p: AuthenticatedPeer, name: string, index?: number) {
     if (!peersByModule.has(name)) {
@@ -76,11 +124,11 @@ export function setupApp(options?: {
   app.get('/ws', defineWebSocketHandler({
     open: (peer) => {
       if (authToken) {
-        peers.set(peer.id, { peer, authenticated: false, name: '' })
+        peers.set(peer.id, { peer, authenticated: false, name: '', lastHeartbeatAt: Date.now() })
       }
       else {
         peer.send(RESPONSES.authenticated)
-        peers.set(peer.id, { peer, authenticated: true, name: '' })
+        peers.set(peer.id, { peer, authenticated: true, name: '', lastHeartbeatAt: Date.now() })
       }
 
       logger.withFields({ peer: peer.id, activePeers: peers.size }).log('connected')
@@ -106,7 +154,35 @@ export function setupApp(options?: {
         peerModuleIndex: authenticatedPeer?.index,
       }).debug('received event')
 
+      if (authenticatedPeer) {
+        authenticatedPeer.lastHeartbeatAt = Date.now()
+        if (event.metadata?.source) {
+          authenticatedPeer.identity = event.metadata.source
+        }
+      }
+
       switch (event.type) {
+        case 'transport:connection:heartbeat': {
+          const p = peers.get(peer.id)
+          if (p) {
+            p.lastHeartbeatAt = Date.now()
+          }
+
+          if (event.data.message === MessageHeartbeat.Ping) {
+            send(peer, {
+              type: 'transport:connection:heartbeat',
+              data: {
+                message: MessageHeartbeat.Pong,
+                mark: MessageHeartbeatMark.Pong,
+                at: Date.now(),
+              },
+              source: WebSocketEventSource.Server,
+            })
+          }
+
+          return
+        }
+
         case 'module:authenticate': {
           if (authToken && event.data.token !== authToken) {
             logger.withFields({ peer: peer.id, peerRemote: peer.remoteAddress, peerRequest: peer.request.url }).log('authentication failed')
@@ -137,7 +213,7 @@ export function setupApp(options?: {
           unregisterModulePeer(p)
 
           // verify
-          const { name, index } = event.data as { name: string, index?: number }
+          const { name, index, identity } = event.data as { name: string, index?: number, identity?: MetadataEventSource }
           if (!name || typeof name !== 'string') {
             send(peer, {
               type: 'error',
@@ -170,6 +246,9 @@ export function setupApp(options?: {
 
           p.name = name
           p.index = index
+          if (identity) {
+            p.identity = identity
+          }
 
           registerModulePeer(p, name, index)
 
@@ -231,6 +310,33 @@ export function setupApp(options?: {
       }
 
       const payload = JSON.stringify(event)
+      const allowBypass = options?.routing?.allowBypass !== false
+      const shouldBypass = Boolean(event.route?.bypass && allowBypass && isDevtoolsPeer(p))
+      const destinations = shouldBypass ? undefined : collectDestinations(event)
+      const routingContext: RouteContext = {
+        event,
+        fromPeer: p,
+        peers,
+        destinations,
+      }
+
+      let decision: RouteDecision | undefined
+      for (const middleware of routingMiddleware) {
+        const result = middleware(routingContext)
+        if (result) {
+          decision = result
+          break
+        }
+      }
+
+      if (decision?.type === 'drop') {
+        logger.withFields({ peer: peer.id, peerName: p.name, event }).debug('routing dropped event')
+        return
+      }
+
+      const targetIds = decision?.type === 'targets' ? decision.targetIds : undefined
+      const shouldBroadcast = decision?.type === 'broadcast' || !targetIds
+
       logger.withFields({ peer: peer.id, peerName: p.name, event }).debug('broadcasting event to peers')
 
       for (const [id, other] of peers.entries()) {
@@ -239,6 +345,14 @@ export function setupApp(options?: {
           continue
         }
 
+        if (!shouldBroadcast && targetIds && !targetIds.has(id)) {
+          continue
+        }
+
+        if (shouldBroadcast && destinations && destinations.length > 0 && !matchesDestinations(destinations, other)) {
+          continue
+        }
+
         try {
           logger.withFields({ fromPeer: peer.id, fromPeerName: p.name, toPeer: other.peer.id, toPeerName: other.name, event }).debug('sending event to peer')
           other.peer.send(payload)

packages/server-runtime/src/middlewares/index.ts

@@ -0,0 +1 @@
+export * from './route'

packages/server-runtime/src/middlewares/route.test.ts

@@ -0,0 +1,139 @@
+import type { RouteTargetExpression, WebSocketBaseEvent, WebSocketEvents } from '@proj-airi/server-shared/types'
+
+import type { AuthenticatedPeer } from '../types'
+
+import { describe, expect, it } from 'vitest'
+
+import { collectDestinations, createPolicyMiddleware, isDevtoolsPeer, matchesDestinations } from './route'
+import { matchesLabelSelector, matchesLabelSelectors, matchesRouteExpression } from './route/match-expression'
+
+function createPeer(options: {
+  id: string
+  name: string
+  plugin?: string
+  instanceId?: string
+  labels?: Record<string, string>
+}): AuthenticatedPeer {
+  return {
+    peer: { id: options.id, send: () => 0 },
+    authenticated: true,
+    name: options.name,
+    identity: options.plugin && options.instanceId
+      ? { plugin: options.plugin, instanceId: options.instanceId, labels: options.labels }
+      : undefined,
+  }
+}
+
+function createSparkNotifyEvent(overrides?: Partial<WebSocketBaseEvent<'spark:notify', WebSocketEvents['spark:notify'], any>>): WebSocketBaseEvent<'spark:notify', WebSocketEvents['spark:notify'], any> {
+  return {
+    type: 'spark:notify',
+    data: {
+      id: 'evt-1',
+      eventId: 'spark-1',
+      kind: 'ping',
+      urgency: 'soon',
+      headline: 'hello',
+      destinations: ['module:character'],
+    },
+    source: 'proj-airi:server-runtime',
+    ...overrides,
+  }
+}
+
+describe('match-expression', () => {
+  it('matches label selectors', () => {
+    expect(matchesLabelSelector('env=prod', { env: 'prod' })).toBe(true)
+    expect(matchesLabelSelector('env=prod', { env: 'dev' })).toBe(false)
+    expect(matchesLabelSelector('feature', { feature: 'on' })).toBe(true)
+    expect(matchesLabelSelector('missing', { env: 'prod' })).toBe(false)
+  })
+
+  it('matches label selector list', () => {
+    expect(matchesLabelSelectors(['env=prod', 'tier=backend'], { env: 'prod', tier: 'backend' })).toBe(true)
+    expect(matchesLabelSelectors(['env=prod', 'tier=backend'], { env: 'prod', tier: 'frontend' })).toBe(false)
+  })
+
+  it('matches route expressions', () => {
+    const peer = createPeer({
+      id: 'peer-1',
+      name: 'stage-ui',
+      plugin: 'stage-ui',
+      instanceId: 'stage-ui-1',
+      labels: { env: 'prod' },
+    })
+
+    const expression: RouteTargetExpression = { type: 'label', selectors: ['env=prod'] }
+    expect(matchesRouteExpression(expression, peer)).toBe(true)
+
+    const globExpression: RouteTargetExpression = { type: 'glob', glob: 'stage-*' }
+    expect(matchesRouteExpression(globExpression, peer)).toBe(true)
+  })
+})
+
+describe('route middleware', () => {
+  it('collects destinations from route before data', () => {
+    const event = createSparkNotifyEvent({
+      data: {
+        id: 'evt-2',
+        eventId: 'spark-2',
+        kind: 'ping',
+        urgency: 'soon',
+        headline: 'hello',
+        destinations: ['module:character'],
+      },
+      route: { destinations: ['label:env=prod'] },
+    })
+
+    expect(collectDestinations(event)).toEqual(['label:env=prod'])
+  })
+
+  it('matches destinations by label selector', () => {
+    const peer = createPeer({
+      id: 'peer-2',
+      name: 'telegram-bot',
+      plugin: 'telegram-bot',
+      instanceId: 'telegram-1',
+      labels: { app: 'telegram', env: 'prod' },
+    })
+
+    expect(matchesDestinations(['label:app=telegram'], peer)).toBe(true)
+    expect(matchesDestinations(['label:env=dev'], peer)).toBe(false)
+  })
+
+  it('policy middleware filters targets', () => {
+    const peers = new Map<string, AuthenticatedPeer>([
+      ['peer-1', createPeer({ id: 'peer-1', name: 'telegram', plugin: 'telegram-bot', instanceId: 'telegram-1', labels: { env: 'prod' } })],
+      ['peer-2', createPeer({ id: 'peer-2', name: 'stage-ui', plugin: 'stage-ui', instanceId: 'stage-ui-1', labels: { env: 'dev' } })],
+    ])
+
+    const policy = createPolicyMiddleware({ allowLabels: ['env=prod'] })
+    const decision = policy({
+      event: createSparkNotifyEvent(),
+      fromPeer: peers.get('peer-1')!,
+      peers,
+      destinations: undefined,
+    })
+
+    expect(decision).toBeDefined()
+    if (!decision)
+      return
+
+    expect(decision?.type).toBe('targets')
+    if (decision.type !== 'targets')
+      return
+
+    expect([...decision!.targetIds]).toEqual(['peer-1'])
+  })
+
+  it('devtools peer detection uses label', () => {
+    const peer = createPeer({
+      id: 'peer-3',
+      name: 'debug-ui',
+      plugin: 'debug-ui',
+      instanceId: 'debug-ui-1',
+      labels: { devtools: 'true' },
+    })
+
+    expect(isDevtoolsPeer(peer)).toBe(true)
+  })
+})

packages/server-runtime/src/middlewares/route.ts

@@ -0,0 +1,85 @@
+import type { RouteTargetExpression, WebSocketEvent } from '@proj-airi/server-shared/types'
+
+import type { AuthenticatedPeer } from '../types'
+
+import { matchesDestinations, matchesLabelSelectors } from './route/match-expression'
+
+export type RouteDecision
+  = | { type: 'drop' }
+    | { type: 'broadcast' }
+    | { type: 'targets', targetIds: Set<string> }
+
+export interface RoutingPolicy {
+  allowPlugins?: string[]
+  denyPlugins?: string[]
+  allowLabels?: string[]
+  denyLabels?: string[]
+}
+
+export interface RouteContext {
+  event: WebSocketEvent
+  fromPeer: AuthenticatedPeer
+  peers: Map<string, AuthenticatedPeer>
+  destinations?: Array<string | RouteTargetExpression>
+}
+
+export type RouteMiddleware = (context: RouteContext) => RouteDecision | void
+
+export function isDevtoolsPeer(peer: AuthenticatedPeer) {
+  const devtoolsLabel = peer.identity?.labels?.devtools
+  const isDevtoolsLabel = devtoolsLabel === 'true' || devtoolsLabel === '1'
+  return Boolean(isDevtoolsLabel || peer.name.includes('devtools'))
+}
+
+export function peerMatchesPolicy(peer: AuthenticatedPeer, policy: RoutingPolicy) {
+  if (policy.allowPlugins?.length && !policy.allowPlugins.includes(peer.identity?.plugin ?? '')) {
+    return false
+  }
+
+  if (policy.denyPlugins?.length && policy.denyPlugins.includes(peer.identity?.plugin ?? '')) {
+    return false
+  }
+
+  const labels = peer.identity?.labels ?? {}
+  if (policy.allowLabels?.length && !matchesLabelSelectors(policy.allowLabels, labels)) {
+    return false
+  }
+
+  if (policy.denyLabels?.length && matchesLabelSelectors(policy.denyLabels, labels)) {
+    return false
+  }
+
+  return true
+}
+
+export function createPolicyMiddleware(policy: RoutingPolicy): RouteMiddleware {
+  return ({ event, peers }) => {
+    if (event.route?.bypass) {
+      return
+    }
+
+    const targetIds = new Set<string>()
+    for (const [id, peer] of peers.entries()) {
+      if (peerMatchesPolicy(peer, policy)) {
+        targetIds.add(id)
+      }
+    }
+
+    return { type: 'targets', targetIds }
+  }
+}
+
+export function collectDestinations(event: WebSocketEvent) {
+  if (event.route?.destinations?.length) {
+    return event.route.destinations
+  }
+
+  const data = event.data as { destinations?: Array<string | RouteTargetExpression> } | undefined
+  if (data?.destinations?.length) {
+    return data.destinations
+  }
+
+  return undefined
+}
+
+export { matchesDestinations }

packages/server-runtime/src/middlewares/route/match-expression.ts

@@ -0,0 +1,112 @@
+import type { RouteTargetExpression } from '@proj-airi/server-shared/types'
+
+import type { AuthenticatedPeer } from '../../types'
+
+function globToRegExp(glob: string) {
+  const escaped = glob.replace(/[.+^${}()|[\]\\]/g, '\\$&')
+  const pattern = `^${escaped.replace(/\*/g, '.*')}$`
+  return new RegExp(pattern)
+}
+
+function matchesGlob(glob: string, value?: string) {
+  if (!value) {
+    return false
+  }
+
+  return globToRegExp(glob).test(value)
+}
+
+export function matchesLabelSelector(selector: string, labels: Record<string, string>) {
+  const [key, value] = selector.split('=', 2)
+  if (!key) {
+    return false
+  }
+
+  if (typeof value === 'undefined') {
+    return key in labels
+  }
+
+  return labels[key] === value
+}
+
+export function matchesLabelSelectors(selectors: string[], labels: Record<string, string>) {
+  return selectors.every(selector => matchesLabelSelector(selector, labels))
+}
+
+export function matchesRouteExpression(expression: RouteTargetExpression, peer: AuthenticatedPeer) {
+  switch (expression.type) {
+    case 'and':
+      return expression.all.every(expr => matchesRouteExpression(expr, peer))
+    case 'or':
+      return expression.any.some(expr => matchesRouteExpression(expr, peer))
+    case 'glob': {
+      const matched = matchesGlob(expression.glob, peer.name)
+        || matchesGlob(expression.glob, peer.identity?.plugin)
+        || matchesGlob(expression.glob, peer.identity?.instanceId)
+
+      return expression.inverted ? !matched : matched
+    }
+    case 'ids': {
+      const matched = expression.ids.includes(peer.peer.id)
+      return expression.inverted ? !matched : matched
+    }
+    case 'plugin': {
+      const matched = expression.plugins.includes(peer.identity?.plugin ?? '')
+      return expression.inverted ? !matched : matched
+    }
+    case 'instance': {
+      const matched = expression.instances.includes(peer.identity?.instanceId ?? '')
+      return expression.inverted ? !matched : matched
+    }
+    case 'label': {
+      const matched = matchesLabelSelectors(expression.selectors, peer.identity?.labels ?? {})
+      return expression.inverted ? !matched : matched
+    }
+    case 'module': {
+      const matched = expression.modules.includes(peer.name)
+      return expression.inverted ? !matched : matched
+    }
+    case 'source': {
+      const matched = expression.sources.includes(peer.name)
+      return expression.inverted ? !matched : matched
+    }
+    default:
+      return false
+  }
+}
+
+export function matchesDestination(destination: string | RouteTargetExpression, peer: AuthenticatedPeer) {
+  if (typeof destination !== 'string') {
+    return matchesRouteExpression(destination, peer)
+  }
+
+  if (destination === '*') {
+    return true
+  }
+
+  const [prefix, rawValue] = destination.split(':', 2)
+  const value = rawValue ?? ''
+
+  switch (prefix) {
+    case 'plugin':
+      return peer.identity?.plugin === value
+    case 'instance':
+      return peer.identity?.instanceId === value
+    case 'label':
+      return matchesLabelSelectors([value], peer.identity?.labels ?? {})
+    case 'peer':
+      return peer.peer.id === value
+    case 'module':
+      return peer.name === value
+    case 'source':
+      return peer.name === value
+    default:
+      return matchesGlob(destination, peer.name)
+        || matchesGlob(destination, peer.identity?.plugin)
+        || matchesGlob(destination, peer.identity?.instanceId)
+  }
+}
+
+export function matchesDestinations(destinations: Array<string | RouteTargetExpression>, peer: AuthenticatedPeer) {
+  return destinations.some(destination => matchesDestination(destination, peer))
+}

packages/server-runtime/src/types/conn.ts

@@ -1,3 +1,5 @@
+import type { MetadataEventSource } from '@proj-airi/server-shared/types'
+
 export interface Peer {
   /**
    * Unique random [uuid v4](https://developer.mozilla.org/en-US/docs/Glossary/UUID) identifier for the peer.
@@ -27,4 +29,6 @@ export enum WebSocketReadyState {
 
 export interface AuthenticatedPeer extends NamedPeer {
   authenticated: boolean
+  identity?: MetadataEventSource
+  lastHeartbeatAt?: number
 }

packages/server-runtime/vitest.config.ts

@@ -0,0 +1,7 @@
+import { defineConfig } from 'vitest/config'
+
+export default defineConfig({
+  test: {
+    include: ['src/**/*.test.ts'],
+  },
+})

packages/server-sdk/src/client.ts

@@ -1,4 +1,7 @@
 import type {
+  MessageHeartbeat,
+  MessageHeartbeatMark,
+  MetadataEventSource,
   WebSocketBaseEvent,
   WebSocketEvent,
   WebSocketEventOptionalSource,
@@ -15,6 +18,11 @@ export interface ClientOptions<C = undefined> {
   name: string
   possibleEvents?: Array<keyof WebSocketEvents<C>>
   token?: string
+  identity?: MetadataEventSource
+  heartbeat?: {
+    readTimeout?: number
+    message?: MessageHeartbeat | string
+  }
   onError?: (error: unknown) => void
   onClose?: () => void
   autoConnect?: boolean
@@ -22,6 +30,10 @@ export interface ClientOptions<C = undefined> {
   maxReconnectAttempts?: number
 }
 
+function createInstanceId() {
+  return `${Date.now().toString(36)}-${Math.random().toString(36).slice(2, 8)}`
+}
+
 export class Client<C = undefined> {
   private connected = false
   private connecting = false
@@ -29,6 +41,8 @@ export class Client<C = undefined> {
   private shouldClose = false
   private connectAttempt?: Promise<void>
   private connectTask?: Promise<void>
+  private heartbeatTimer?: ReturnType<typeof setInterval>
+  private readonly identity: MetadataEventSource
 
   private readonly opts: Required<Omit<ClientOptions<C>, 'token'>> & Pick<ClientOptions<C>, 'token'>
   private readonly eventListeners = new Map<
@@ -37,6 +51,11 @@ export class Client<C = undefined> {
   >()
 
   constructor(options: ClientOptions<C>) {
+    const identity = options.identity ?? {
+      plugin: options.name,
+      instanceId: createInstanceId(),
+    }
+
     this.opts = {
       url: 'ws://localhost:6121/ws',
       possibleEvents: [],
@@ -45,9 +64,16 @@ export class Client<C = undefined> {
       autoConnect: true,
       autoReconnect: true,
       maxReconnectAttempts: -1,
+      heartbeat: {
+        readTimeout: 30_000,
+        message: MessageHeartbeat.Ping,
+      },
       ...options,
+      identity,
     }
 
+    this.identity = identity
+
     // Authentication listener is registered once only
     this.onEvent('module:authenticated', async (event) => {
       if (event.data.authenticated) {
@@ -64,6 +90,12 @@ export class Client<C = undefined> {
       }
     })
 
+    this.onEvent('transport:connection:heartbeat', (event) => {
+      if (event.data.message === MessageHeartbeat.Ping) {
+        this.sendHeartbeatPong()
+      }
+    })
+
     if (this.opts.autoConnect) {
       void this.connect()
     }
@@ -145,6 +177,7 @@ export class Client<C = undefined> {
 
         if (this.connected) {
           this.connected = false
+          this.stopHeartbeat()
           this.opts.onClose?.()
         }
         if (this.opts.autoReconnect && !this.shouldClose) {
@@ -155,6 +188,8 @@ export class Client<C = undefined> {
         settle(() => {
           this.connected = true
 
+          this.startHeartbeat()
+
           if (this.opts.token)
             this.tryAuthenticate()
           else
@@ -186,6 +221,7 @@ export class Client<C = undefined> {
       type: 'module:announce',
       data: {
         name: this.opts.name,
+        identity: this.identity,
         possibleEvents: this.opts.possibleEvents,
       },
     })
@@ -261,7 +297,11 @@ export class Client<C = undefined> {
 
   send(data: WebSocketEventOptionalSource<C>): void {
     if (this.websocket && this.connected) {
-      this.websocket.send(JSON.stringify({ source: this.opts.name as WebSocketEventSource | string, ...data } as WebSocketEvent<C>))
+      this.websocket.send(JSON.stringify({
+        source: this.opts.name as WebSocketEventSource | string,
+        metadata: { source: this.identity },
+        ...data,
+      } as WebSocketEvent<C>))
     }
   }
 
@@ -273,12 +313,71 @@ export class Client<C = undefined> {
 
   close(): void {
     this.shouldClose = true
+    this.stopHeartbeat()
     if (this.websocket) {
       this.websocket.close()
       this.connected = false
     }
   }
 
+  private startHeartbeat() {
+    if (!this.opts.heartbeat?.readTimeout) {
+      return
+    }
+
+    this.stopHeartbeat()
+
+    const ping = () => this.sendHeartbeatPing()
+
+    ping()
+    this.heartbeatTimer = setInterval(ping, this.opts.heartbeat.readTimeout)
+  }
+
+  private stopHeartbeat() {
+    if (this.heartbeatTimer) {
+      clearInterval(this.heartbeatTimer)
+      this.heartbeatTimer = undefined
+    }
+  }
+
+  private sendNativeHeartbeat(kind: 'ping' | 'pong') {
+    const websocket = this.websocket as WebSocket & {
+      ping?: () => void
+      pong?: () => void
+    }
+
+    if (kind === 'ping') {
+      websocket.ping?.()
+    }
+    else {
+      websocket.pong?.()
+    }
+  }
+
+  private sendHeartbeatPing() {
+    this.send({
+      type: 'transport:connection:heartbeat',
+      data: {
+        message: this.opts.heartbeat?.message ?? MessageHeartbeat.Ping,
+        mark: MessageHeartbeatMark.Ping,
+        at: Date.now(),
+      },
+    })
+    this.sendNativeHeartbeat('ping')
+  }
+
+  private sendHeartbeatPong() {
+    this.send({
+      type: 'transport:connection:heartbeat',
+      data: {
+        message: MessageHeartbeat.Pong,
+        mark: MessageHeartbeatMark.Pong,
+        at: Date.now(),
+      },
+    })
+    this.sendNativeHeartbeat('pong')
+  }
+
   private async _reconnectDueToUnauthorized() {
     if (this.shouldClose)
       return

packages/server-shared/src/types/websocket/events.ts

@@ -12,6 +12,39 @@ export interface Discord {
   channelId?: string
 }
 
+export interface MetadataEventSource {
+  plugin: string
+  instanceId: string
+  version?: string
+  labels?: Record<string, string>
+}
+
+export type RouteTargetExpression
+  = | { type: 'and', all: RouteTargetExpression[] }
+    | { type: 'or', any: RouteTargetExpression[] }
+    | { type: 'glob', glob: string, inverted?: boolean }
+    | { type: 'ids', ids: string[], inverted?: boolean }
+    | { type: 'plugin', plugins: string[], inverted?: boolean }
+    | { type: 'instance', instances: string[], inverted?: boolean }
+    | { type: 'label', selectors: string[], inverted?: boolean }
+    | { type: 'module', modules: string[], inverted?: boolean }
+    | { type: 'source', sources: string[], inverted?: boolean }
+
+export interface RouteConfig {
+  destinations?: Array<string | RouteTargetExpression>
+  bypass?: boolean
+}
+
+export enum MessageHeartbeat {
+  Ping = 'ping',
+  Pong = 'pong',
+}
+
+export enum MessageHeartbeatMark {
+  Ping = '🩵',
+  Pong = '💛',
+}
+
 export enum WebSocketEventSource {
   Server = 'proj-airi:server-runtime',
   StageWeb = 'proj-airi:stage-web',
@@ -74,7 +107,14 @@ export interface ContextUpdate<
 export interface WebSocketBaseEvent<T, D, S extends string = string> {
   type: T
   data: D
+  /**
+   * @deprecated Prefer metadata.source.
+   */
   source: WebSocketEventSource | S
+  metadata?: {
+    source: MetadataEventSource
+  }
+  route?: RouteConfig
 }
 
 export type WithInputSource<Source extends keyof InputSource> = {
@@ -102,6 +142,7 @@ export interface WebSocketEvents<C = undefined> {
   }
   'module:announce': {
     name: string
+    identity?: MetadataEventSource
     possibleEvents: Array<(keyof WebSocketEvents<C>)>
   }
   'module:configure': {
@@ -245,6 +286,12 @@ export interface WebSocketEvents<C = undefined> {
     destinations: Array<string>
   }
 
+  'transport:connection:heartbeat': {
+    message: MessageHeartbeat | string
+    mark?: '🩵' | '💛'
+    at?: number
+  }
+
   'context:update': ContextUpdate
 }
 

vitest.config.ts

@@ -7,6 +7,7 @@ export default defineConfig({
       'packages/stage-ui',
       'packages/vite-plugin-warpdrive',
       'packages/audio-pipelines-transcribe',
+      'packages/server-runtime',
     ],
   },
 })