clawmatrix 0.1.0

package/src/peer-manager.ts

@@ -0,0 +1,333 @@
+import { EventEmitter } from "node:events";
+import type { ServerWebSocket } from "bun";
+import type { ClawMatrixConfig, PeerConfig } from "./config.ts";
+import { Connection } from "./connection.ts";
+import type { WsTransport } from "./connection.ts";
+import { Router } from "./router.ts";
+import type {
+  AnyClusterFrame,
+  ClusterFrame,
+  NodeCapabilities,
+  PeerInfo,
+  PeerSync,
+} from "./types.ts";
+
+const RECONNECT_BASE = 1_000;
+const RECONNECT_MAX = 60_000;
+
+export interface PeerManagerEvents {
+  frame: [frame: AnyClusterFrame, from: Connection];
+  peerConnected: [nodeId: string];
+  peerDisconnected: [nodeId: string];
+}
+
+export class PeerManager extends EventEmitter<PeerManagerEvents> {
+  readonly router: Router;
+  private config: ClawMatrixConfig;
+  private localCapabilities: NodeCapabilities;
+  private wsServer: ReturnType<typeof Bun.serve> | null = null;
+  private reconnectTimers = new Map<string, ReturnType<typeof setTimeout>>();
+  private reconnectAttempts = new Map<string, number>();
+  private stopped = false;
+  /** Map from ServerWebSocket to Connection for inbound connections. */
+  private inboundConnections = new Map<ServerWebSocket<unknown>, Connection>();
+
+  constructor(config: ClawMatrixConfig) {
+    super();
+    this.config = config;
+    this.localCapabilities = {
+      nodeId: config.nodeId,
+      agents: config.agents,
+      models: config.models,
+      tags: config.tags,
+    };
+    this.router = new Router(config.nodeId, {
+      agents: config.agents,
+      models: config.models,
+      tags: config.tags,
+    });
+  }
+
+  // ── Lifecycle ──────────────────────────────────────────────────
+  async start() {
+    if (this.config.listen) {
+      this.startListening();
+    }
+    for (const peer of this.config.peers) {
+      this.connectToPeer(peer);
+    }
+  }
+
+  async stop() {
+    this.stopped = true;
+    for (const timer of this.reconnectTimers.values()) {
+      clearTimeout(timer);
+    }
+    this.reconnectTimers.clear();
+
+    this.router.broadcast({
+      type: "peer_leave",
+      from: this.config.nodeId,
+      timestamp: Date.now(),
+      payload: { nodeId: this.config.nodeId },
+    } as AnyClusterFrame);
+
+    for (const conn of this.router.getDirectConnections()) {
+      conn.close(1000, "shutdown");
+    }
+
+    if (this.wsServer) {
+      this.wsServer.stop();
+      this.wsServer = null;
+    }
+  }
+
+  // ── Inbound WS server (Bun.serve) ─────────────────────────────
+  private startListening() {
+    const port = this.config.listenPort;
+    const hostname = this.config.listenHost;
+    const self = this;
+
+    this.wsServer = Bun.serve<undefined>({
+      port,
+      hostname,
+      fetch(req, server) {
+        if (server.upgrade(req)) {
+          return undefined;
+        }
+        return new Response("WebSocket upgrade required", { status: 426 });
+      },
+      websocket: {
+        open(ws) {
+          self.handleInboundOpen(ws);
+        },
+        message(ws, message) {
+          const conn = self.inboundConnections.get(ws);
+          if (conn) {
+            conn.feedMessage(typeof message === "string" ? message : Buffer.from(message));
+          }
+        },
+        close(ws, code, reason) {
+          const conn = self.inboundConnections.get(ws);
+          if (conn) {
+            conn.feedClose(code, reason);
+            self.inboundConnections.delete(ws);
+          }
+        },
+      },
+    });
+  }
+
+  private handleInboundOpen(ws: ServerWebSocket<unknown>) {
+    // Wrap Bun's ServerWebSocket into our WsTransport interface
+    const transport: WsTransport = {
+      send(data: string) {
+        ws.send(data);
+      },
+      close(code?: number, reason?: string) {
+        ws.close(code, reason);
+      },
+      get readyState() {
+        return ws.readyState;
+      },
+    };
+
+    const conn = new Connection(
+      transport,
+      "inbound",
+      this.config.nodeId,
+      this.config.secret,
+      this.localCapabilities,
+    );
+
+    this.inboundConnections.set(ws, conn);
+
+    conn.on("authenticated", (caps) => {
+      this.onPeerAuthenticated(conn, caps);
+    });
+
+    conn.on("error", () => {
+      // Connection handles cleanup
+    });
+  }
+
+  // ── Outbound connections (standard WebSocket) ──────────────────
+  private connectToPeer(peer: PeerConfig) {
+    if (this.stopped) return;
+
+    const ws = new WebSocket(peer.url);
+
+    ws.addEventListener("open", () => {
+      const conn = new Connection(
+        ws,
+        "outbound",
+        this.config.nodeId,
+        this.config.secret,
+        this.localCapabilities,
+      );
+      conn.bindWebSocket(ws);
+
+      conn.on("authenticated", (caps) => {
+        this.reconnectAttempts.delete(peer.nodeId);
+        this.onPeerAuthenticated(conn, caps);
+      });
+
+      conn.on("error", () => {
+        // will close
+      });
+    });
+
+    ws.addEventListener("error", () => {
+      this.scheduleReconnect(peer);
+    });
+
+    ws.addEventListener("close", () => {
+      this.scheduleReconnect(peer);
+    });
+  }
+
+  private scheduleReconnect(peer: PeerConfig) {
+    if (this.stopped) return;
+    if (this.reconnectTimers.has(peer.nodeId)) return;
+
+    const attempt = this.reconnectAttempts.get(peer.nodeId) ?? 0;
+    const delay = Math.min(RECONNECT_BASE * 2 ** attempt, RECONNECT_MAX);
+    this.reconnectAttempts.set(peer.nodeId, attempt + 1);
+
+    const timer = setTimeout(() => {
+      this.reconnectTimers.delete(peer.nodeId);
+      this.connectToPeer(peer);
+    }, delay);
+    this.reconnectTimers.set(peer.nodeId, timer);
+  }
+
+  // ── Peer lifecycle ─────────────────────────────────────────────
+  private onPeerAuthenticated(conn: Connection, caps: NodeCapabilities) {
+    const nodeId = conn.remoteNodeId!;
+    this.router.addDirectPeer(nodeId, conn, caps);
+
+    conn.on("message", (frame) => this.onFrame(frame, conn));
+    conn.on("close", () => this.onPeerDisconnected(conn));
+
+    this.sendPeerSync(conn);
+
+    // Broadcast peer_join to other peers
+    this.router.broadcast({
+      type: "peer_join",
+      from: this.config.nodeId,
+      timestamp: Date.now(),
+      payload: {
+        nodeId,
+        agents: caps.agents,
+        models: caps.models,
+        tags: caps.tags,
+      },
+    } as AnyClusterFrame);
+
+    // Re-sync with all existing peers so they learn about the new node
+    for (const existingConn of this.router.getDirectConnections()) {
+      if (existingConn !== conn && existingConn.isOpen) {
+        this.sendPeerSync(existingConn);
+      }
+    }
+
+    this.emit("peerConnected", nodeId);
+  }
+
+  private onPeerDisconnected(conn: Connection) {
+    const nodeId = conn.remoteNodeId;
+    if (!nodeId) return;
+
+    this.router.removePeer(nodeId);
+
+    this.router.broadcast({
+      type: "peer_leave",
+      from: this.config.nodeId,
+      timestamp: Date.now(),
+      payload: { nodeId },
+    } as AnyClusterFrame);
+
+    this.emit("peerDisconnected", nodeId);
+  }
+
+  // ── Message handling ───────────────────────────────────────────
+  private onFrame(frame: AnyClusterFrame, from: Connection) {
+    if (frame.id && this.router.isDuplicate(frame.id)) return;
+
+    if (frame.type === "peer_sync") {
+      this.handlePeerSync(frame as PeerSync, from);
+      return;
+    }
+
+    if (frame.type === "peer_join") {
+      const peer = (frame as AnyClusterFrame & { payload: PeerInfo }).payload;
+      if (peer.nodeId !== this.config.nodeId) {
+        this.router.addRelayPeer(peer, from.remoteNodeId!);
+      }
+      return;
+    }
+
+    if (frame.type === "peer_leave") {
+      const { nodeId } = (frame as AnyClusterFrame & { payload: { nodeId: string } }).payload;
+      const route = this.router.getRoute(nodeId);
+      if (route && route.reachableVia) {
+        this.router.removePeer(nodeId);
+      }
+      return;
+    }
+
+    if (frame.to && frame.to !== this.config.nodeId) {
+      this.router.tryRelay(frame);
+      return;
+    }
+
+    this.emit("frame", frame, from);
+  }
+
+  private sendPeerSync(conn: Connection) {
+    const peers = this.router.buildPeerSyncPayload();
+    conn.send({
+      type: "peer_sync",
+      from: this.config.nodeId,
+      timestamp: Date.now(),
+      payload: { peers },
+    } as AnyClusterFrame);
+  }
+
+  private handlePeerSync(frame: PeerSync, from: Connection) {
+    let changed = false;
+    for (const peer of frame.payload.peers) {
+      if (peer.nodeId === this.config.nodeId) continue;
+      if (peer.nodeId === from.remoteNodeId) {
+        const prev = this.router.getRoute(peer.nodeId);
+        const hadAgents = prev?.agents.length ?? 0;
+        this.router.updatePeerCapabilities(peer.nodeId, peer);
+        if (peer.agents.length !== hadAgents || peer.models.length !== (prev?.models.length ?? 0)) {
+          changed = true;
+        }
+      } else {
+        const existing = this.router.getRoute(peer.nodeId);
+        if (!existing) changed = true;
+        this.router.addRelayPeer(peer, from.remoteNodeId!);
+      }
+    }
+
+    // If we learned new info, re-sync with other peers
+    if (changed) {
+      for (const conn of this.router.getDirectConnections()) {
+        if (conn !== from && conn.isOpen) {
+          this.sendPeerSync(conn);
+        }
+      }
+    }
+  }
+
+  // ── Public API ─────────────────────────────────────────────────
+  sendTo(nodeId: string, frame: ClusterFrame | AnyClusterFrame): boolean {
+    return this.router.sendTo(nodeId, frame);
+  }
+
+  broadcast(frame: ClusterFrame | AnyClusterFrame) {
+    this.router.broadcast(frame);
+  }
+}

package/src/router.ts

@@ -0,0 +1,275 @@
+import type { ClusterFrame, AnyClusterFrame, PeerInfo, AgentInfo, ModelInfo } from "./types.ts";
+import type { Connection } from "./connection.ts";
+
+const DEFAULT_TTL = 3;
+const SEEN_FRAME_TTL = 60_000; // 60s dedup window
+const MAX_SEEN_FRAMES = 10_000;
+
+export interface RouteEntry {
+  nodeId: string;
+  agents: AgentInfo[];
+  models: ModelInfo[];
+  tags: string[];
+  connection: Connection | null; // null = reachable via relay
+  reachableVia: string | null; // relay nodeId, null = direct
+  lastSeen: number;
+  latencyMs: number;
+}
+
+export class Router {
+  private nodeId: string;
+  private localAgents: AgentInfo[];
+  private localModels: ModelInfo[];
+  private localTags: string[];
+  private routes = new Map<string, RouteEntry>();
+  private connections = new Map<string, Connection>(); // nodeId → direct connection
+  private seenFrames = new Map<string, number>(); // frameId → timestamp
+
+  constructor(
+    nodeId: string,
+    localCapabilities?: { agents: AgentInfo[]; models: ModelInfo[]; tags: string[] },
+  ) {
+    this.nodeId = nodeId;
+    this.localAgents = localCapabilities?.agents ?? [];
+    this.localModels = localCapabilities?.models ?? [];
+    this.localTags = localCapabilities?.tags ?? [];
+  }
+
+  // ── Route table management ─────────────────────────────────────
+  addDirectPeer(
+    nodeId: string,
+    connection: Connection,
+    capabilities: { agents: AgentInfo[]; models: ModelInfo[]; tags: string[] },
+  ) {
+    this.connections.set(nodeId, connection);
+    this.routes.set(nodeId, {
+      nodeId,
+      agents: capabilities.agents,
+      models: capabilities.models,
+      tags: capabilities.tags,
+      connection,
+      reachableVia: null,
+      lastSeen: Date.now(),
+      latencyMs: 0,
+    });
+  }
+
+  addRelayPeer(peer: PeerInfo, viaNodeId: string) {
+    // Don't add ourselves
+    if (peer.nodeId === this.nodeId) return;
+    // Don't overwrite direct connections with relay entries
+    const existing = this.routes.get(peer.nodeId);
+    if (existing?.connection) return;
+
+    this.routes.set(peer.nodeId, {
+      nodeId: peer.nodeId,
+      agents: peer.agents,
+      models: peer.models,
+      tags: peer.tags,
+      connection: null,
+      reachableVia: viaNodeId,
+      lastSeen: Date.now(),
+      latencyMs: 0,
+    });
+  }
+
+  removePeer(nodeId: string) {
+    this.connections.delete(nodeId);
+    this.routes.delete(nodeId);
+    // Also remove routes that relied on this node as relay
+    for (const [id, entry] of this.routes) {
+      if (entry.reachableVia === nodeId) {
+        this.routes.delete(id);
+      }
+    }
+  }
+
+  updatePeerCapabilities(
+    nodeId: string,
+    capabilities: { agents: AgentInfo[]; models: ModelInfo[]; tags: string[] },
+  ) {
+    const entry = this.routes.get(nodeId);
+    if (entry) {
+      entry.agents = capabilities.agents;
+      entry.models = capabilities.models;
+      entry.tags = capabilities.tags;
+      entry.lastSeen = Date.now();
+    }
+  }
+
+  // ── Routing ────────────────────────────────────────────────────
+  getRoute(nodeId: string): RouteEntry | undefined {
+    return this.routes.get(nodeId);
+  }
+
+  /** Resolve target agent to a specific nodeId. Supports agent ID or "tags:<tag>". */
+  resolveAgent(target: string): RouteEntry | undefined {
+    const isTagQuery = target.startsWith("tags:");
+    const tag = isTagQuery ? target.slice(5) : null;
+
+    let candidates: RouteEntry[] = [];
+    for (const entry of this.routes.values()) {
+      if (isTagQuery) {
+        if (entry.agents.some((a) => a.tags.includes(tag!)) || entry.tags.includes(tag!)) {
+          candidates.push(entry);
+        }
+      } else {
+        if (entry.agents.some((a) => a.id === target)) {
+          candidates.push(entry);
+        }
+      }
+    }
+
+    if (candidates.length === 0) return undefined;
+
+    // Sort: direct connections first, then by latency
+    candidates.sort((a, b) => {
+      const aDirect = a.connection ? 0 : 1;
+      const bDirect = b.connection ? 0 : 1;
+      if (aDirect !== bDirect) return aDirect - bDirect;
+      return a.latencyMs - b.latencyMs;
+    });
+
+    return candidates[0];
+  }
+
+  /** Find node that has a specific model. */
+  resolveModel(modelId: string): RouteEntry | undefined {
+    let candidates: RouteEntry[] = [];
+    for (const entry of this.routes.values()) {
+      if (entry.models.some((m) => m.id === modelId)) {
+        candidates.push(entry);
+      }
+    }
+    if (candidates.length === 0) return undefined;
+
+    candidates.sort((a, b) => {
+      const aDirect = a.connection ? 0 : 1;
+      const bDirect = b.connection ? 0 : 1;
+      if (aDirect !== bDirect) return aDirect - bDirect;
+      return a.latencyMs - b.latencyMs;
+    });
+    return candidates[0];
+  }
+
+  /** Resolve a node target (nodeId or "tags:<tag>"). */
+  resolveNode(target: string): RouteEntry | undefined {
+    if (target.startsWith("tags:")) {
+      const tag = target.slice(5);
+      let best: RouteEntry | undefined;
+      for (const entry of this.routes.values()) {
+        if (entry.tags.includes(tag)) {
+          if (!best || (entry.connection && !best.connection) || entry.latencyMs < best.latencyMs) {
+            best = entry;
+          }
+        }
+      }
+      return best;
+    }
+    return this.routes.get(target);
+  }
+
+  // ── Message sending and relay ──────────────────────────────────
+  /** Send a frame to a specific node, relaying if necessary. Returns true if sent. */
+  sendTo(targetNodeId: string, frame: ClusterFrame | AnyClusterFrame): boolean {
+    const route = this.routes.get(targetNodeId);
+    if (!route) return false;
+
+    if (route.connection?.isOpen) {
+      route.connection.send(frame);
+      return true;
+    }
+
+    // Relay through intermediate node
+    if (route.reachableVia) {
+      const relay = this.connections.get(route.reachableVia);
+      if (relay?.isOpen) {
+        relay.send(frame);
+        return true;
+      }
+    }
+
+    return false;
+  }
+
+  /** Broadcast a frame to all direct peers. */
+  broadcast(frame: ClusterFrame | AnyClusterFrame) {
+    for (const conn of this.connections.values()) {
+      if (conn.isOpen) {
+        conn.send(frame);
+      }
+    }
+  }
+
+  /** Handle relay: if message is not for us, forward it. Returns true if relayed. */
+  tryRelay(frame: ClusterFrame): boolean {
+    if (!frame.to || frame.to === this.nodeId) return false;
+
+    const ttl = (frame.ttl ?? DEFAULT_TTL) - 1;
+    if (ttl <= 0) return false;
+
+    const relayed = this.sendTo(frame.to, { ...frame, ttl });
+    return relayed;
+  }
+
+  // ── Deduplication ──────────────────────────────────────────────
+  /** Returns true if the frame has been seen before (duplicate). */
+  isDuplicate(frameId: string): boolean {
+    if (!frameId) return false;
+    if (this.seenFrames.has(frameId)) return true;
+
+    this.seenFrames.set(frameId, Date.now());
+    this.pruneSeenFrames();
+    return false;
+  }
+
+  /** Mark a request ID as failed so late responses are ignored. */
+  markFailed(requestId: string) {
+    this.seenFrames.set(`failed:${requestId}`, Date.now());
+  }
+
+  isFailed(requestId: string): boolean {
+    return this.seenFrames.has(`failed:${requestId}`);
+  }
+
+  private pruneSeenFrames() {
+    if (this.seenFrames.size <= MAX_SEEN_FRAMES) return;
+    const now = Date.now();
+    for (const [id, ts] of this.seenFrames) {
+      if (now - ts > SEEN_FRAME_TTL) {
+        this.seenFrames.delete(id);
+      }
+    }
+  }
+
+  // ── Accessors ──────────────────────────────────────────────────
+  getAllPeers(): RouteEntry[] {
+    return [...this.routes.values()];
+  }
+
+  getDirectConnections(): Connection[] {
+    return [...this.connections.values()];
+  }
+
+  /** Build PeerInfo list for peer_sync. */
+  buildPeerSyncPayload(): PeerInfo[] {
+    const peers: PeerInfo[] = [];
+    // Include ourselves
+    peers.push({
+      nodeId: this.nodeId,
+      agents: this.localAgents,
+      models: this.localModels,
+      tags: this.localTags,
+    });
+    for (const entry of this.routes.values()) {
+      peers.push({
+        nodeId: entry.nodeId,
+        agents: entry.agents,
+        models: entry.models,
+        tags: entry.tags,
+        reachableVia: entry.reachableVia ?? undefined,
+      });
+    }
+    return peers;
+  }
+}