@cortexkit/aft-opencode 0.42.0 → 0.43.0

package/src/shared/rpc-server.ts

@@ -1,4 +1,4 @@
-import { randomBytes } from "node:crypto";
+import { randomBytes, timingSafeEqual } from "node:crypto";
 import {
   existsSync,
   mkdirSync,
@@ -11,22 +11,92 @@ import {
 import { createServer, type IncomingMessage, type Server, type ServerResponse } from "node:http";
 import { dirname, join } from "node:path";
 import { log, warn } from "../logger";
+import {
+  drainNotifications,
+  type NotificationSink,
+  registerNotificationSink,
+  registerStatusChangeSink,
+  type StatusChangeSink,
+} from "./rpc-notifications";
 import { isPidAlive, parseRpcPortRecord, rpcPortFileDir } from "./rpc-utils";
 
 type RpcHandler = (params: Record<string, unknown>) => Promise<Record<string, unknown>>;
 
+type BunServe = <Data>(options: BunServeOptions<Data>) => BunServer<Data>;
+
+interface BunRuntime {
+  serve: BunServe;
+}
+
+interface BunServer<Data> {
+  port?: number;
+  stop(closeActiveConnections?: boolean): void;
+  upgrade(req: Request, options: { data: Data }): boolean;
+}
+
+interface BunServeOptions<Data> {
+  port: number;
+  hostname: string;
+  fetch(
+    req: Request,
+    server: BunServer<Data>,
+  ): Response | Promise<Response | undefined> | undefined;
+  websocket: {
+    open(ws: BunServerWebSocket<Data>): void;
+    message(ws: BunServerWebSocket<Data>, raw: string | Buffer): void;
+    close(ws: BunServerWebSocket<Data>): void;
+  };
+}
+
+interface BunServerWebSocket<Data> {
+  data: Data;
+  send(data: string): unknown;
+  close(code?: number, reason?: string): void;
+}
+
+interface WsData {
+  authed: boolean;
+  sessionId?: string;
+  unregisterNotification?: () => void;
+  unregisterStatus?: () => void;
+  authTimer?: ReturnType<typeof setTimeout>;
+}
+
 const PORT_FILE_HEARTBEAT_MS = 15_000;
+const MAX_BODY_BYTES = 1_048_576;
+const WS_AUTH_TIMEOUT_MS = 5_000;
+const WS_CLOSE_UNAUTHORIZED = 4401;
+
+function bunRuntime(): BunRuntime | undefined {
+  return (globalThis as typeof globalThis & { Bun?: BunRuntime }).Bun;
+}
+
+function tokensMatch(presented: string, expected: string): boolean {
+  const a = Buffer.from(presented, "utf8");
+  const b = Buffer.from(expected, "utf8");
+  if (a.length !== b.length) return false;
+  return timingSafeEqual(a, b);
+}
+
+function json(body: unknown, status = 200): Response {
+  return new Response(JSON.stringify(body), {
+    status,
+    headers: { "Content-Type": "application/json" },
+  });
+}
 
 export class AftRpcServer {
-  private server: Server | null = null;
+  private nodeServer: Server | null = null;
+  private bunServer: BunServer<WsData> | null = null;
   private port = 0;
-  private token: string | null = null;
+  private token = randomBytes(32).toString("hex");
   private handlers = new Map<string, RpcHandler>();
   private portFilePath: string;
   private portsDir: string;
   private heartbeatTimer: ReturnType<typeof setInterval> | null = null;
   /** Unique per-instance ID — distinguishes our entry from duplicate plugin loads. */
   private instanceId: string;
+  private sockets = new Set<BunServerWebSocket<WsData>>();
 
   constructor(storageDir: string, directory: string) {
     this.portsDir = rpcPortFileDir(storageDir, directory);
@@ -41,6 +111,36 @@ export class AftRpcServer {
 
   /** Start the server on a random port, write port to disk. */
   async start(): Promise<number> {
+    const bun = bunRuntime();
+    if (bun) {
+      return this.startBun(bun);
+    }
+    return this.startNode();
+  }
+
+  private async startBun(bun: BunRuntime): Promise<number> {
+    const server = bun.serve<WsData>({
+      port: 0,
+      hostname: "127.0.0.1",
+      fetch: (req, srv) => this.handleFetch(req, srv),
+      websocket: {
+        open: (ws) => {
+          ws.data.authTimer = setTimeout(() => {
+            if (!ws.data.authed) ws.close(WS_CLOSE_UNAUTHORIZED, "auth timeout");
+          }, WS_AUTH_TIMEOUT_MS);
+        },
+        message: (ws, raw) => this.handleWsMessage(ws, raw),
+        close: (ws) => this.closeWs(ws),
+      },
+    });
+
+    this.bunServer = server;
+    this.port = server.port ?? 0;
+    this.afterServerStarted();
+    return this.port;
+  }
+
+  private async startNode(): Promise<number> {
     return new Promise((resolve, reject) => {
       const server = createServer((req, res) => this.dispatch(req, res));
 
@@ -56,48 +156,8 @@ export class AftRpcServer {
           return;
         }
         this.port = addr.port;
-        this.token = randomBytes(32).toString("hex");
-        this.server = server;
-
-        // Write port file atomically
-        try {
-          const dir = dirname(this.portFilePath);
-          mkdirSync(dir, { recursive: true, mode: 0o700 });
-          const tmpPath = `${this.portFilePath}.tmp`;
-          // Record pid + started_at so the client can skip files whose owning
-          // process is dead (without a health-check round-trip) and prefer the
-          // freshest live server, instead of accumulating crash/restart leftovers.
-          writeFileSync(
-            tmpPath,
-            JSON.stringify({
-              port: this.port,
-              token: this.token,
-              pid: process.pid,
-              started_at: Date.now(),
-            }),
-            { encoding: "utf-8", mode: 0o600 },
-          );
-          renameSync(tmpPath, this.portFilePath);
-          log(`RPC server listening on 127.0.0.1:${this.port}`);
-          // Self-heal: the port file is this server's only discoverability
-          // record, and historically it was written exactly once — anything
-          // that deleted it (a client misjudging a load-induced health-check
-          // timeout as staleness, cache cleanup, manual sweeps) silently
-          // orphaned the server until host restart (issue #110: sidebar stuck
-          // on the lazy placeholder forever). Recreate it if it goes missing.
-          this.heartbeatTimer = setInterval(() => this.ensurePortFile(), PORT_FILE_HEARTBEAT_MS);
-          this.heartbeatTimer.unref?.();
-          // Hygiene: sweep dead siblings while we're here. The client only
-          // reclaims dead-pid files on a cold port scan, and it caches the
-          // first warm port afterwards — so crash/restart leftovers piled up
-          // (dozens of dead files per project hash were observed). Server
-          // startup is the natural sweep point: it runs once per instance and
-          // already owns this directory.
-          this.sweepDeadPortFiles();
-        } catch (err) {
-          warn(`Failed to write RPC port file: ${err}`);
-        }
-
+        this.nodeServer = server;
+        this.afterServerStarted();
         resolve(this.port);
       });
 
@@ -106,6 +166,40 @@ export class AftRpcServer {
     });
   }
 
+  private afterServerStarted(): void {
+    try {
+      this.writePortFile();
+      log(`RPC server listening on 127.0.0.1:${this.port}`);
+      // Self-heal: the port file is this server's only discoverability record.
+      // Anything that deletes it silently orphans the server until host restart.
+      // Recreate it if missing.
+      this.heartbeatTimer = setInterval(() => this.ensurePortFile(), PORT_FILE_HEARTBEAT_MS);
+      this.heartbeatTimer.unref?.();
+      // Hygiene: sweep dead port files from other instances while we're here.
+      // Server startup is the natural sweep point and already owns this directory.
+      this.sweepDeadPortFiles();
+    } catch (err) {
+      warn(`Failed to write RPC port file: ${err}`);
+    }
+  }
+
+  private writePortFile(): void {
+    const dir = dirname(this.portFilePath);
+    mkdirSync(dir, { recursive: true, mode: 0o700 });
+    const tmpPath = `${this.portFilePath}.tmp`;
+    writeFileSync(
+      tmpPath,
+      JSON.stringify({
+        port: this.port,
+        token: this.token,
+        pid: process.pid,
+        started_at: Date.now(),
+      }),
+      { encoding: "utf-8", mode: 0o600 },
+    );
+    renameSync(tmpPath, this.portFilePath);
+  }
+
   /**
    * Remove sibling port files whose owning process is provably dead.
    * Only files with a recorded pid that no longer maps to a live process are
@@ -143,21 +237,10 @@ export class AftRpcServer {
 
   /** Rewrite the port file if it disappeared (wrongful deletion recovery). */
   private ensurePortFile(): void {
-    if (!this.server || this.token === null) return;
+    if ((!this.nodeServer && !this.bunServer) || this.port <= 0) return;
     try {
       if (existsSync(this.portFilePath)) return;
-      const tmpPath = `${this.portFilePath}.tmp`;
-      writeFileSync(
-        tmpPath,
-        JSON.stringify({
-          port: this.port,
-          token: this.token,
-          pid: process.pid,
-          started_at: Date.now(),
-        }),
-        { encoding: "utf-8", mode: 0o600 },
-      );
-      renameSync(tmpPath, this.portFilePath);
+      this.writePortFile();
       log(`RPC port file was missing; rewrote ${this.portFilePath}`);
     } catch {
       // best-effort; retried on the next heartbeat
@@ -170,11 +253,23 @@ export class AftRpcServer {
       clearInterval(this.heartbeatTimer);
       this.heartbeatTimer = null;
     }
-    if (this.server) {
-      this.server.close();
-      this.server = null;
+    for (const ws of this.sockets) {
+      try {
+        this.closeWs(ws);
+        ws.close();
+      } catch {
+        // Ignore close errors during shutdown; remaining sockets are closed independently.
+      }
+    }
+    this.sockets.clear();
+    if (this.nodeServer) {
+      this.nodeServer.close();
+      this.nodeServer = null;
+    }
+    if (this.bunServer) {
+      this.bunServer.stop(true);
+      this.bunServer = null;
     }
-    this.token = null;
     try {
       unlinkSync(this.portFilePath);
     } catch {
@@ -182,6 +277,132 @@ export class AftRpcServer {
     }
   }
 
+  private async handleFetch(req: Request, srv: BunServer<WsData>): Promise<Response | undefined> {
+    const url = new URL(req.url);
+
+    if (url.pathname === "/ws") {
+      const upgraded = srv.upgrade(req, { data: { authed: false } });
+      if (upgraded) return undefined;
+      return new Response("upgrade failed", { status: 400 });
+    }
+
+    if (req.method === "GET" && url.pathname === "/health") {
+      return json({ ok: true, pid: process.pid });
+    }
+
+    if (req.method !== "POST" || !url.pathname.startsWith("/rpc/")) {
+      return new Response("Not Found", { status: 404 });
+    }
+
+    const method = url.pathname.slice(5);
+    const handler = this.handlers.get(method);
+    if (!handler) {
+      return json({ error: `Unknown method: ${method}` }, 404);
+    }
+
+    const bodyText = await req.text();
+    if (bodyText.length > MAX_BODY_BYTES) {
+      return new Response("Request too large", { status: 413 });
+    }
+
+    let params: Record<string, unknown> = {};
+    try {
+      if (bodyText.length > 0) {
+        params = JSON.parse(bodyText);
+      }
+    } catch {
+      return json({ error: "Invalid JSON" }, 400);
+    }
+
+    if (!tokensMatch(typeof params.token === "string" ? params.token : "", this.token)) {
+      return json({ error: "Forbidden" }, 403);
+    }
+
+    const { token: _token, ...handlerParams } = params;
+
+    try {
+      const result = await handler(handlerParams);
+      return json(result);
+    } catch (err) {
+      log(`RPC error: ${method} => ${err}`);
+      return json({ error: String(err) }, 500);
+    }
+  }
+
+  private handleWsMessage(ws: BunServerWebSocket<WsData>, raw: string | Buffer): void {
+    let msg: { type?: string; token?: string; sessionId?: string; lastReceivedId?: number };
+    try {
+      msg = JSON.parse(typeof raw === "string" ? raw : raw.toString("utf8"));
+    } catch {
+      return;
+    }
+
+    if (msg.type === "hello") {
+      if (!tokensMatch(typeof msg.token === "string" ? msg.token : "", this.token)) {
+        ws.send(JSON.stringify({ type: "error", error: "unauthorized" }));
+        ws.close(WS_CLOSE_UNAUTHORIZED, "bad token");
+        return;
+      }
+      if (ws.data.authTimer) {
+        clearTimeout(ws.data.authTimer);
+        ws.data.authTimer = undefined;
+      }
+
+      ws.data.unregisterNotification?.();
+      ws.data.unregisterStatus?.();
+      ws.data.authed = true;
+      ws.data.sessionId =
+        typeof msg.sessionId === "string" && msg.sessionId.length > 0 ? msg.sessionId : undefined;
+
+      const lastReceivedId = Number(msg.lastReceivedId ?? 0);
+      const backlog = drainNotifications(
+        Number.isFinite(lastReceivedId) ? lastReceivedId : 0,
+        ws.data.sessionId,
+      );
+
+      const notificationSink: NotificationSink = {
+        sessionId: ws.data.sessionId,
+        send: (notification) => {
+          ws.send(JSON.stringify({ type: "notification", notification }));
+        },
+      };
+      const statusSink: StatusChangeSink = {
+        sessionId: ws.data.sessionId,
+        send: (event) => {
+          ws.send(JSON.stringify({ type: "status-changed", ...event }));
+        },
+      };
+      ws.data.unregisterNotification = registerNotificationSink(notificationSink);
+      ws.data.unregisterStatus = registerStatusChangeSink(statusSink);
+      this.sockets.add(ws);
+
+      for (const notification of backlog) {
+        ws.send(JSON.stringify({ type: "notification", notification }));
+      }
+      ws.send(JSON.stringify({ type: "hello-ack" }));
+      return;
+    }
+
+    if (msg.type === "ack") {
+      const lastReceivedId = Number(msg.lastReceivedId ?? 0);
+      if (Number.isFinite(lastReceivedId) && lastReceivedId > 0) {
+        drainNotifications(lastReceivedId, ws.data.sessionId);
+      }
+    }
+  }
+
+  private closeWs(ws: BunServerWebSocket<WsData>): void {
+    if (ws.data.authTimer) {
+      clearTimeout(ws.data.authTimer);
+      ws.data.authTimer = undefined;
+    }
+    ws.data.unregisterNotification?.();
+    ws.data.unregisterNotification = undefined;
+    ws.data.unregisterStatus?.();
+    ws.data.unregisterStatus = undefined;
+    this.sockets.delete(ws);
+  }
+
   private dispatch(req: IncomingMessage, res: ServerResponse): void {
     const url = req.url ?? "";
 
@@ -208,7 +429,7 @@ export class AftRpcServer {
     let body = "";
     req.on("data", (chunk: Buffer) => {
       body += chunk.toString();
-      if (body.length > 1_048_576) {
+      if (body.length > MAX_BODY_BYTES) {
         res.writeHead(413);
         res.end("Request too large");
         req.destroy();
@@ -227,7 +448,7 @@ export class AftRpcServer {
         return;
       }
 
-      if (params.token !== this.token) {
+      if (!tokensMatch(typeof params.token === "string" ? params.token : "", this.token)) {
         res.writeHead(403, { "Content-Type": "application/json" });
         res.end(JSON.stringify({ error: "Forbidden" }));
         return;
@@ -235,9 +456,7 @@ export class AftRpcServer {
 
       const { token: _token, ...handlerParams } = params;
 
-      // Successful RPC calls are intentionally NOT logged. The /aft-status
-      // TUI dialog polls every 1.5s while open, which used to spam the log
-      // with RPC call/result pairs. Errors are still logged (real signal).
+      // Successful RPC calls are not logged to avoid noise; errors are still logged.
       handler(handlerParams)
         .then((result) => {
           res.writeHead(200, { "Content-Type": "application/json" });