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

2026-01-10 00:28:21 +08:00
parent 21641b469d
commit eb803ab76e
10 changed files with 615 additions and 6 deletions
@@ -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)
@@ -0,0 +1 @@
+export * from './route'
@@ -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)
+  })
+})
@@ -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 }
@@ -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))
+}
@@ -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
 }
@@ -0,0 +1,7 @@
+import { defineConfig } from 'vitest/config'
+
+export default defineConfig({
+  test: {
+    include: ['src/**/*.test.ts'],
+  },
+})