coreci-chat/apps/control-plane/ws-server.ts

/**
 * control-plane ws-server — standalone WebSocket server for the Relay Agent
 * (REQ-011, REQ-012, REQ-013, Wave D Task 2).
 *
 * Next.js App Router route handlers are HTTP-only (no native WebSocket
 * upgrade), so for M1 the relay WebSocket server runs as a small standalone
 * HTTP server alongside Next.js. It binds 0.0.0.0:${CORECI_WS_PORT:-3001} and
 * upgrades requests to /api/relay/ws.
 *
 * Lifecycle per connection:
 *   1. On upgrade: read Authorization: Bearer <token>, verify the JWT relay
 *      token via @coreci/secrets verifyRelayToken (G-005). If invalid → close
 *      with code 4001 (REQ-012 Edge 12).
 *   2. On {type:"register"} message: INSERT into targets
 *      (tenant_id, hostname, os_name, os_version, ip_address, agent_version,
 *      last_seen_at=now()) under withTenant + audit append (provision event).
 *      Respond {type:"registered", targetId}.
 *   3. On {type:"ping"} message: UPDATE targets SET last_seen_at=now() WHERE
 *      id=targetId under withTenant. Respond {type:"pong", ts}.
 *
 * Connected agents are tracked in-memory for the dashboard (Wave E consumes
 * this via `getConnectedAgents()`).
 *
 * Run standalone: `tsx ws-server.ts` (next to `next dev`/`next start`).
 * The systemd unit / docker-compose starts both processes.
 */

import { createServer, type IncomingMessage } from "node:http";
import { WebSocketServer, type WebSocket } from "ws";
import { randomUUID } from "node:crypto";

import { verifyRelayToken, RelayTokenError } from "@coreci/secrets";
import { withTenant, appendAudit, type ScopedClient } from "@coreci/db";
import { getDb } from "./lib/db.js";

/** RELAY_TOKEN_SIGNING_KEY — infra bootstrap signing key (G-010 tier a). */
function signingKey(): string {
  const k = process.env.RELAY_TOKEN_SIGNING_KEY;
  if (!k) throw new Error("RELAY_TOKEN_SIGNING_KEY not configured");
  return k;
}

const WS_PATH = "/api/relay/ws";
const WS_PORT = Number(process.env.CORECI_WS_PORT ?? 3001);

// ─── In-memory connected-agent registry (Wave E consumes this) ───────────────
export interface ConnectedAgent {
  tenantId: string;
  targetId: string;
  hostname: string;
  osName: string;
  connectedAt: number;
  lastSeenAt: number;
}

const connectedAgents = new Map<WebSocket, ConnectedAgent>();

/** Snapshot of connected agents (Wave E dashboard reads this). */
export function getConnectedAgents(): ConnectedAgent[] {
  return Array.from(connectedAgents.values());
}

// ─── Reverse index + tool_call routing (Wave H, R-003) ───────────────────────
//
// `targetsByTenant: Map<tenantId, Map<targetId, WebSocket>>` is the reverse
// index built on `connectedAgents`. The SSH adapter's `RelayTransport` impl
// calls `getTargetWebSocket(tenantId, targetId)` to find the connected Relay
// Agent for a target; if offline → the broker returns HTTP 404 (no queueing).
//
// `pendingToolCalls: Map<callId, PendingToolCall>` tracks in-flight tool_calls
// awaiting tool_results. The 10s broker timeout (R-003) rejects the promise
// with a `timeout` error if no tool_result arrives.

const targetsByTenant = new Map<string, Map<string, WebSocket>>();
const pendingToolCalls = new Map<string, PendingToolCall>();

/** Broker-side timeout for the tool_call → tool_result round-trip (R-003). */
const BROKER_TOOL_CALL_TIMEOUT_MS = 10_000;

/**
 * Get the connected WebSocket for a (tenantId, targetId) tuple. Returns
 * undefined if the target is offline (no connected Relay Agent). The caller
 * (the SSH adapter's RelayTransport impl) surfaces HTTP 404 on offline.
 *
 * Tenant scoping is enforced: the reverse index is keyed by tenantId, so a
 * caller cannot resolve a target belonging to another tenant (INV-2 / RLS at
 * the routing layer; the `mcp_adapters` row was already tenant-scoped by the
 * router via `withTenant` + RLS before reaching the adapter).
 */
export function getTargetWebSocket(tenantId: string, targetId: string): WebSocket | undefined {
  const byTarget = targetsByTenant.get(tenantId);
  if (!byTarget) return undefined;
  const ws = byTarget.get(targetId);
  if (!ws || ws.readyState !== ws.OPEN) return undefined;
  return ws;
}

/**
 * Send a `tool_call` to the connected Relay Agent for (tenantId, targetId)
 * and await a `tool_result`. Returns the result envelope. Rejects with a
 * `ToolCallError`-like Error (`code` property) on:
 *   - `target_offline` — no connected Relay Agent for (tenantId, targetId).
 *   - `timeout` — no tool_result within BROKER_TOOL_CALL_TIMEOUT_MS (10s).
 *   - `send_failed` — the WebSocket write failed (connection lost mid-send).
 *
 * This is the `RelayTransport.sendToolCall` implementation the SSH adapter
 * consumes via the `apps/control-plane/lib/mcp.ts` wiring. The agent
 * independently enforces a 9.5s exec timeout so it returns a timeout
 * tool_result 0.5s before the broker gives up → the SSE stream closes
 * cleanly (R-003).
 */
export function sendToolCall(
  tenantId: string,
  targetId: string,
  call: { callId: string; command: string; timeoutMs: number },
): Promise<{ callId: string; stdout?: string; stderr?: string; exitCode: number; error?: string }> {
  const ws = getTargetWebSocket(tenantId, targetId);
  if (!ws) {
    const err = new Error(`target '${targetId}' is offline (no connected Relay Agent)`) as Error & { code: string };
    err.code = "target_offline";
    return Promise.reject(err);
  }

  return new Promise((resolve, reject) => {
    const message: ToolCallMessage = {
      type: "tool_call",
      callId: call.callId,
      command: call.command,
      timeoutMs: call.timeoutMs,
    };
    let wroteOk = false;
    try {
      ws.send(JSON.stringify(message));
      wroteOk = true;
    } catch (err) {
      const e = new Error(`tool_call write failed: ${err instanceof Error ? err.message : String(err)}`) as Error & { code: string };
      e.code = "send_failed";
      reject(e);
      return;
    }
    if (!wroteOk) return; // (defensive — reject already called)

    // Track the pending call; resolve on tool_result, reject on timeout.
    const timer = setTimeout(() => {
      pendingToolCalls.delete(call.callId);
      const e = new Error(`tool_call timeout after ${BROKER_TOOL_CALL_TIMEOUT_MS}ms`) as Error & { code: string };
      e.code = "timeout";
      reject(e);
    }, BROKER_TOOL_CALL_TIMEOUT_MS);

    pendingToolCalls.set(call.callId, {
      callId: call.callId,
      resolve: (result) => {
        clearTimeout(timer);
        pendingToolCalls.delete(call.callId);
        resolve(result);
      },
      reject: (err) => {
        clearTimeout(timer);
        pendingToolCalls.delete(call.callId);
        reject(err);
      },
      timer,
    });
  });
}

// ─── Wire message types ──────────────────────────────────────────────────────
interface RegisterMessage {
  type: "register";
  tenantId?: string; // ignored — resolved from the verified JWT
  hostname: string;
  os: string;
  osVersion: string;
  ip: string;
  agentVersion: string;
}

interface PingMessage {
  type: "ping";
  ts: number;
}

interface RegisteredResponse {
  type: "registered";
  targetId: string;
}

interface PongResponse {
  type: "pong";
  ts: number;
}

interface ErrorResponse {
  type: "error";
  error: string;
  code?: number;
  detail?: string;
}

/**
 * `tool_call` (server → agent, R-003 §4, Wave H). The broker sends this to the
 * connected Relay Agent for (tenantId, targetId) to run a whitelisted command.
 * The agent responds with a `tool_result` (ToolResultMessage, inbound).
 */
export interface ToolCallMessage {
  type: "tool_call";
  callId: string;
  command: string;
  timeoutMs: number;
}

/**
 * `tool_result` (agent → server, R-003 §4, Wave H). Inbound — handled by the
 * `handleMessage` switch (the agent sends this after running the command).
 * The broker resolves the pending call by `callId` and resolves the promise.
 */
interface ToolResultMessage {
  type: "tool_result";
  callId: string;
  stdout?: string;
  stderr?: string;
  exitCode: number;
  error?: string;
}

/**
 * A pending tool_call awaiting a tool_result. The broker's `RelayTransport`
 * impl creates this when sending a tool_call; the `tool_result` handler
 * resolves it (or rejects on timeout).
 */
interface PendingToolCall {
  callId: string;
  resolve: (result: ToolResultMessage) => void;
  reject: (err: Error) => void;
  timer: NodeJS.Timeout;
}

// ─── WS server bootstrap ─────────────────────────────────────────────────────
export async function startWsServer(port = WS_PORT): Promise<{ server: ReturnType<typeof createServer>; wss: WebSocketServer }> {
  // Ensure the DB is bootstrapped (loads migrations) before accepting connections.
  await getDb();

  const server = createServer();
  const wss = new WebSocketServer({ noServer: true });

  server.on("upgrade", (req, socket, head) => {
    const url = new URL(req.url ?? "", `http://${req.headers.host ?? "localhost"}`);
    if (url.pathname !== WS_PATH) {
      socket.destroy();
      return;
    }

    const token = bearerToken(req);
    if (!token) {
      socket.destroy();
      return;
    }

    let tenantId: string;
    try {
      const claims = verifyRelayToken(token, signingKey());
      tenantId = claims.tenantId;
    } catch (err) {
      // REQ-012 Edge 12: registration/token failure → close with 4001.
      logWarn(`relay ws auth failed: ${err instanceof Error ? err.message : String(err)}`);
      socket.destroy();
      return;
    }

    wss.handleUpgrade(req, socket, head, (ws) => {
      // Stash the resolved tenantId on the socket for the message handlers.
      (ws as WebSocket & { tenantId: string }).tenantId = tenantId;
      wss.emit("connection", ws, req);
    });
  });

  wss.on("connection", (ws) => {
    const tenantId = (ws as WebSocket & { tenantId: string }).tenantId;
    logInfo(`relay ws connected (tenant=${tenantId})`);

    ws.on("message", (data) => {
      handleMessage(ws, tenantId, data)
        .catch((err) => {
          logWarn(`relay ws message error: ${err instanceof Error ? err.message : String(err)}`);
          safeSend(ws, { type: "error", error: "internal_error", detail: err instanceof Error ? err.message : String(err) } satisfies ErrorResponse);
        });
    });

    ws.on("close", () => {
      const agent = connectedAgents.get(ws);
      if (agent) {
        logInfo(`relay ws closed (tenant=${agent.tenantId} target=${agent.targetId})`);
        connectedAgents.delete(ws);
        // Wave H: remove from the reverse index. If the agent reconnects, the
        // new WebSocket replaces this entry on register. Also reject any
        // in-flight tool_calls for this target (the connection is gone).
        const byTarget = targetsByTenant.get(agent.tenantId);
        if (byTarget) {
          if (byTarget.get(agent.targetId) === ws) byTarget.delete(agent.targetId);
          if (byTarget.size === 0) targetsByTenant.delete(agent.tenantId);
        }
      }
      // Reject any pending tool_calls whose WebSocket was this one. We don't
      // have a per-call ws ref, so we reject ALL pending calls that can no
      // longer be delivered (their ws is closed). A future enhancement could
      // track the ws per pending call; M2's volume is low so a sweep is fine.
      for (const [callId, pending] of pendingToolCalls) {
        // Best-effort: reject pending calls (the broker timeout will also
        // fire if we miss one here). This is defensive cleanup.
        pending.reject(new Error("connection_lost"));
        pendingToolCalls.delete(callId);
      }
    });

    ws.on("error", (err) => {
      logWarn(`relay ws socket error: ${err.message}`);
    });
  });

  return new Promise((resolve, reject) => {
    server.on("error", reject);
    server.listen(port, "0.0.0.0", () => {
      logInfo(`relay ws server listening on ws://0.0.0.0:${port}${WS_PATH}`);
      resolve({ server, wss });
    });
  });
}

// ─── Message handling ────────────────────────────────────────────────────────
async function handleMessage(ws: WebSocket, tenantId: string, data: unknown): Promise<void> {
  let msg: Record<string, unknown>;
  try {
    if (typeof data !== "string" && !(data instanceof Buffer) && !Array.isArray(data)) {
      throw new Error("non-text message");
    }
    const text = data instanceof Buffer ? data.toString("utf8") : Array.isArray(data) ? Buffer.concat(data as Buffer[]).toString("utf8") : String(data);
    msg = JSON.parse(text) as Record<string, unknown>;
  } catch (err) {
    safeSend(ws, { type: "error", error: "invalid_json", detail: err instanceof Error ? err.message : String(err) } satisfies ErrorResponse);
    return;
  }

  switch (msg.type) {
    case "register":
      await handleRegister(ws, tenantId, msg as unknown as RegisterMessage);
      return;
    case "ping":
      await handlePing(ws, tenantId, msg as unknown as PingMessage);
      return;
    case "tool_result":
      handleToolResult(msg as unknown as ToolResultMessage);
      return;
    default:
      safeSend(ws, { type: "error", error: `unknown message type: ${String(msg.type)}` } satisfies ErrorResponse);
  }
}

/**
 * `handleToolResult` — resolve the pending tool_call by `callId` (Wave H,
 * R-003). The agent sends this after running the whitelisted command. If no
 * pending call exists (e.g. the broker already timed out and deleted the
 * entry), the result is dropped (the agent's work is wasted — acceptable;
 * the broker's 10s timeout is the gate).
 */
function handleToolResult(msg: ToolResultMessage): void {
  if (!msg.callId) {
    logWarn(`tool_result: missing callId — dropping`);
    return;
  }
  const pending = pendingToolCalls.get(msg.callId);
  if (!pending) {
    // Late result (broker already timed out) or duplicate — drop.
    logInfo(`tool_result: no pending call for ${msg.callId} (timed out or duplicate) — dropping`);
    return;
  }
  pending.resolve(msg);
}

interface TargetRow {
  id: string;
}

async function handleRegister(ws: WebSocket, tenantId: string, msg: RegisterMessage): Promise<void> {
  if (!msg.hostname || !msg.os || !msg.agentVersion) {
    safeSend(ws, { type: "error", error: "register: missing required fields (hostname, os, agentVersion)" } satisfies ErrorResponse);
    return;
  }

  // INSERT the target under withTenant (RLS enforces tenant scoping) and
  // append a provision audit event in the same transaction (Edge 7: audit
  // failure halts the operation, so the target INSERT rolls back too).
  const targetId = await withTenant(tenantId, async (c: ScopedClient) => {
    const insertRes = await c.query<TargetRow>(
      `INSERT INTO targets (tenant_id, hostname, os_name, os_version, ip_address, agent_version, last_seen_at)
       VALUES ($1, $2, $3, $4, $5, $6, now())
       RETURNING id`,
      [tenantId, msg.hostname, msg.os, msg.osVersion || "", msg.ip || null, msg.agentVersion],
    );
    const id = insertRes.rows[0]?.id;
    if (!id) throw new Error("register: target INSERT returned no id");

    await appendAudit(c, {
      tenantId,
      eventType: "provision",
      payload: {
        action: "relay_target_registered",
        hostname: msg.hostname,
        os: msg.os,
        osVersion: msg.osVersion,
        ip: msg.ip,
        agentVersion: msg.agentVersion,
        targetId: id,
      },
      targetId: id,
    });
    return id;
  });

  // Track in-memory for the dashboard (Wave E).
  connectedAgents.set(ws, {
    tenantId,
    targetId,
    hostname: msg.hostname,
    osName: msg.os,
    connectedAt: Date.now(),
    lastSeenAt: Date.now(),
  });
  // Wave H: also track in the reverse index (tenantId → targetId → ws) so
  // the SSH adapter's RelayTransport can route tool_calls by target.
  let byTarget = targetsByTenant.get(tenantId);
  if (!byTarget) {
    byTarget = new Map();
    targetsByTenant.set(tenantId, byTarget);
  }
  byTarget.set(targetId, ws);

  const response: RegisteredResponse = { type: "registered", targetId };
  safeSend(ws, response);
  logInfo(`relay target registered (tenant=${tenantId} target=${targetId} hostname=${msg.hostname})`);
}

async function handlePing(ws: WebSocket, tenantId: string, msg: PingMessage): Promise<void> {
  const agent = connectedAgents.get(ws);
  if (!agent) {
    safeSend(ws, { type: "error", error: "ping before register" } satisfies ErrorResponse);
    return;
  }

  // UPDATE last_seen under withTenant (RLS enforces scoping).
  await withTenant(tenantId, async (c: ScopedClient) => {
    await c.query(
      `UPDATE targets SET last_seen_at = now() WHERE id = $1`,
      [agent.targetId],
    );
  });
  agent.lastSeenAt = Date.now();

  const response: PongResponse = { type: "pong", ts: msg.ts ?? Date.now() };
  safeSend(ws, response);
}

// ─── Helpers ─────────────────────────────────────────────────────────────────
function bearerToken(req: IncomingMessage): string | null {
  const header = req.headers["authorization"];
  if (!header || typeof header !== "string") return null;
  const m = header.match(/^Bearer\s+(.+)$/i);
  return m ? m[1]!.trim() : null;
}

function safeSend(ws: WebSocket, msg: RegisteredResponse | PongResponse | ErrorResponse): void {
  if (ws.readyState !== ws.OPEN) return;
  try {
    ws.send(JSON.stringify(msg));
  } catch (err) {
    logWarn(`relay ws send failed: ${err instanceof Error ? err.message : String(err)}`);
  }
}

function logInfo(msg: string): void {
  console.log(`[relay-ws] ${msg}`);
}
function logWarn(msg: string): void {
  console.warn(`[relay-ws] ${msg}`);
}

// silence unused import in type-only contexts
void RelayTokenError;
void randomUUID;

// ─── Entrypoint (run standalone: `tsx ws-server.ts`) ──────────────────────────
if (import.meta.url === `file://${process.argv[1]}`) {
  startWsServer().catch((err) => {
    console.error(`[relay-ws] fatal: ${err instanceof Error ? err.message : String(err)}`);
    process.exit(1);
  });
}