@robodev-ai/runtime 0.2.0 → 0.4.0

package/src/local-storage.ts

@@ -0,0 +1,213 @@
+import { randomBytes } from "node:crypto";
+import { mkdir, readFile, rename, unlink, writeFile } from "node:fs/promises";
+import { dirname, join } from "node:path";
+import type { StorageClient, StorageObject, StorageUploadOptions } from "@robodev-ai/sdk";
+import {
+  assertObjectSize,
+  mintStorageServeToken,
+  parseLogicalKey,
+  storageListQuerySchema,
+} from "./storage-rules.js";
+
+export type StorageQueryable = {
+  query: <T extends Record<string, unknown> = Record<string, unknown>>(
+    sql: string,
+    values?: unknown[],
+  ) => Promise<{ rows: T[] }>;
+};
+
+export type LocalObjectRow = {
+  key: string;
+  public: boolean;
+  size_bytes: string | number;
+  content_type: string;
+  created_at: Date;
+};
+
+export type LocalStorageClientOptions = {
+  query: StorageQueryable["query"];
+  objectsDir: string;
+  publicBaseUrl: string;
+  projectId: string;
+  jwtSecret: string;
+};
+
+function notFound(key: string): never {
+  throw Object.assign(new Error(`Storage object not found: ${key}`), {
+    statusCode: 404,
+    code: "not_found",
+  });
+}
+
+function escapeLike(value: string): string {
+  return value.replace(/[\\%_]/g, "\\$&");
+}
+
+function objectPath(objectsDir: string, key: string): string {
+  return join(objectsDir, ...key.split("/"));
+}
+
+function stableObjectUrl(publicBaseUrl: string, key: string): string {
+  return `${publicBaseUrl.replace(/\/$/, "")}/storage/objects/${key}`;
+}
+
+function toStorageObject(publicBaseUrl: string, row: LocalObjectRow): StorageObject {
+  return {
+    key: row.key,
+    public: row.public,
+    size: Number(row.size_bytes),
+    contentType: row.content_type,
+    url: stableObjectUrl(publicBaseUrl, row.key),
+  };
+}
+
+async function loadRow(
+  query: StorageQueryable["query"],
+  key: string,
+): Promise<LocalObjectRow | null> {
+  const result = await query<LocalObjectRow>(
+    `SELECT key, public, size_bytes, content_type, created_at
+     FROM robodev_storage.objects WHERE key = $1`,
+    [key],
+  );
+  return result.rows[0] ?? null;
+}
+
+async function readObjectFile(objectsDir: string, key: string): Promise<Buffer | null> {
+  try {
+    return await readFile(objectPath(objectsDir, key));
+  } catch {
+    return null;
+  }
+}
+
+/** Creates the reserved `robodev_storage` schema. Idempotent, not drizzle, not public. */
+export async function ensureStorageSchema(db: StorageQueryable): Promise<void> {
+  await db.query(`CREATE SCHEMA IF NOT EXISTS robodev_storage`);
+  await db.query(`
+    CREATE TABLE IF NOT EXISTS robodev_storage.objects (
+      key TEXT PRIMARY KEY,
+      public BOOLEAN NOT NULL,
+      size_bytes INTEGER NOT NULL,
+      content_type TEXT NOT NULL,
+      created_at TIMESTAMPTZ NOT NULL DEFAULT now()
+    )
+  `);
+}
+
+export async function resolveLocalObject(input: {
+  query: StorageQueryable["query"];
+  objectsDir: string;
+  key: string;
+}): Promise<{ row: LocalObjectRow; body: Buffer } | null> {
+  const logicalKey = parseLogicalKey(input.key);
+  const row = await loadRow(input.query, logicalKey);
+  if (!row) return null;
+  const body = await readObjectFile(input.objectsDir, logicalKey);
+  if (!body) return null;
+  return { row, body };
+}
+
+export function createLocalStorageClient(options: LocalStorageClientOptions): StorageClient {
+  const query = options.query;
+  const objectsDir = options.objectsDir;
+  const publicBaseUrl = options.publicBaseUrl;
+
+  return {
+    async upload(key, data, upload?: StorageUploadOptions) {
+      const logicalKey = parseLogicalKey(key);
+      const contentType = upload?.contentType?.trim() || "application/octet-stream";
+      const isPublic = Boolean(upload?.public);
+      const buffer = typeof data === "string" ? Buffer.from(data) : Buffer.from(data);
+      assertObjectSize(buffer.byteLength);
+
+      await mkdir(objectsDir, { recursive: true });
+      const dest = objectPath(objectsDir, logicalKey);
+      await mkdir(dirname(dest), { recursive: true });
+      const tmp = join(objectsDir, `.tmp-${randomBytes(8).toString("hex")}`);
+      await writeFile(tmp, buffer);
+      await rename(tmp, dest);
+
+      const result = await query<LocalObjectRow>(
+        `INSERT INTO robodev_storage.objects (key, public, size_bytes, content_type)
+         VALUES ($1, $2, $3, $4)
+         ON CONFLICT (key) DO UPDATE SET
+           public = EXCLUDED.public,
+           size_bytes = EXCLUDED.size_bytes,
+           content_type = EXCLUDED.content_type
+         RETURNING key, public, size_bytes, content_type, created_at`,
+        [logicalKey, isPublic, buffer.byteLength, contentType],
+      );
+      const row = result.rows[0];
+      if (!row) throw new Error("Failed to upsert storage metadata");
+      return toStorageObject(publicBaseUrl, row);
+    },
+
+    async get(key) {
+      const logicalKey = parseLogicalKey(key);
+      const resolved = await resolveLocalObject({ query, objectsDir, key: logicalKey });
+      if (!resolved) notFound(logicalKey);
+      return {
+        body: resolved.body,
+        contentType: resolved.row.content_type,
+        public: resolved.row.public,
+      };
+    },
+
+    async getUrl(key, urlOptions) {
+      const logicalKey = parseLogicalKey(key);
+      const row = await loadRow(query, logicalKey);
+      if (!row) notFound(logicalKey);
+      const body = await readObjectFile(objectsDir, logicalKey);
+      if (!body) notFound(logicalKey);
+      if (row.public) return stableObjectUrl(publicBaseUrl, logicalKey);
+      const { exp, sig } = mintStorageServeToken(
+        options.jwtSecret,
+        options.projectId,
+        logicalKey,
+        urlOptions?.expiresIn,
+      );
+      return `${stableObjectUrl(publicBaseUrl, logicalKey)}?exp=${exp}&sig=${sig}`;
+    },
+
+    async delete(key) {
+      const logicalKey = parseLogicalKey(key);
+      const row = await loadRow(query, logicalKey);
+      if (!row) notFound(logicalKey);
+      try {
+        await unlink(objectPath(objectsDir, logicalKey));
+      } catch {
+        /* ignore missing file */
+      }
+      await query(`DELETE FROM robodev_storage.objects WHERE key = $1`, [logicalKey]);
+    },
+
+    async list(listOptions) {
+      const parsed = storageListQuerySchema.parse({
+        prefix: listOptions?.prefix,
+        limit: listOptions?.limit,
+        offset: listOptions?.offset,
+      });
+      const prefix = parsed.prefix ?? "";
+      if (prefix) parseLogicalKey(prefix.replace(/\/+$/, "") || prefix);
+      const like = prefix ? `${escapeLike(prefix)}%` : "%";
+      const count = await query<{ count: string }>(
+        `SELECT count(*)::text AS count FROM robodev_storage.objects
+         WHERE key LIKE $1 ESCAPE E'\\\\'`,
+        [like],
+      );
+      const result = await query<LocalObjectRow>(
+        `SELECT key, public, size_bytes, content_type, created_at
+         FROM robodev_storage.objects
+         WHERE key LIKE $1 ESCAPE E'\\\\'
+         ORDER BY created_at DESC
+         LIMIT $2 OFFSET $3`,
+        [like, parsed.limit, parsed.offset],
+      );
+      return {
+        objects: result.rows.map((row) => toStorageObject(publicBaseUrl, row)),
+        total: Number(count.rows[0]?.count ?? 0),
+      };
+    },
+  };
+}

package/src/sockets.test.ts

@@ -0,0 +1,112 @@
+import assert from "node:assert/strict";
+import { test } from "node:test";
+import {
+  MAX_PROJECT_SOCKETS,
+  MAX_SOCKET_PAYLOAD_BYTES,
+  SOCKET_CLOSE_INVALID,
+  SOCKET_CLOSE_POLICY,
+  SOCKET_CLOSE_TOO_BIG,
+  attachProjectSocket,
+  closeProjectSockets,
+  dispatchSocketsSend,
+  joinSocketRoom,
+  parseSocketMessage,
+  projectSocketCount,
+  sendToSocketRoom,
+  setProjectSocketNames,
+  stringifySocketPayload,
+} from "./sockets.js";
+
+function mockSocket() {
+  const sent: string[] = [];
+  let readyState = 1;
+  return {
+    sent,
+    get readyState() {
+      return readyState;
+    },
+    send(data: string) {
+      sent.push(data);
+    },
+    close() {
+      readyState = 3;
+    },
+  };
+}
+
+test("51st connection is rejected with 1008", () => {
+  const projectId = "sock-limit";
+  closeProjectSockets(projectId);
+  for (let i = 0; i < MAX_PROJECT_SOCKETS; i++) {
+    const attached = attachProjectSocket({
+      projectId,
+      name: "chat",
+      socket: mockSocket(),
+    });
+    assert.equal(attached.ok, true);
+  }
+  assert.equal(projectSocketCount(projectId), MAX_PROJECT_SOCKETS);
+  const extra = attachProjectSocket({
+    projectId,
+    name: "chat",
+    socket: mockSocket(),
+  });
+  assert.equal(extra.ok, false);
+  if (!extra.ok) {
+    assert.equal(extra.code, SOCKET_CLOSE_POLICY);
+  }
+  closeProjectSockets(projectId);
+  assert.equal(projectSocketCount(projectId), 0);
+});
+
+test("unknown socket name throws unknown_socket", () => {
+  setProjectSocketNames("sock-known", ["chat"]);
+  assert.throws(
+    () => dispatchSocketsSend("sock-known", ["chat"], { name: "missing", payload: { ok: true } }),
+    (err: { message?: string; code?: string }) =>
+      err.message === "unknown_socket" && err.code === "unknown_socket",
+  );
+});
+
+test("room broadcast stays on the same socket name", () => {
+  const projectId = "sock-rooms";
+  closeProjectSockets(projectId);
+  setProjectSocketNames(projectId, ["chat", "alerts"]);
+  const chatA = mockSocket();
+  const chatB = mockSocket();
+  const alerts = mockSocket();
+  const a = attachProjectSocket({ projectId, name: "chat", socket: chatA });
+  const b = attachProjectSocket({ projectId, name: "chat", socket: chatB });
+  const c = attachProjectSocket({ projectId, name: "alerts", socket: alerts });
+  assert.ok(a.ok && b.ok && c.ok);
+  if (!a.ok || !b.ok || !c.ok) return;
+  joinSocketRoom(a.conn, "lobby");
+  joinSocketRoom(b.conn, "lobby");
+  joinSocketRoom(c.conn, "lobby");
+  sendToSocketRoom(projectId, "chat", "lobby", JSON.stringify({ hi: true }), a.conn);
+  assert.deepEqual(chatA.sent, []);
+  assert.deepEqual(chatB.sent, [JSON.stringify({ hi: true })]);
+  assert.deepEqual(alerts.sent, []);
+  dispatchSocketsSend(projectId, ["chat", "alerts"], {
+    name: "chat",
+    room: "missing",
+    payload: { nope: true },
+  });
+  assert.deepEqual(chatA.sent, []);
+  assert.deepEqual(chatB.sent, [JSON.stringify({ hi: true })]);
+  closeProjectSockets(projectId);
+});
+
+test("JSON payload size is capped at 100KB", () => {
+  const ok = stringifySocketPayload({ ok: true });
+  assert.equal(ok, JSON.stringify({ ok: true }));
+  assert.throws(() => stringifySocketPayload("a".repeat(MAX_SOCKET_PAYLOAD_BYTES + 1)), /100KB/);
+  const parsed = parseSocketMessage('{"ok":true}');
+  assert.equal(parsed.ok, true);
+  const invalid = parseSocketMessage("{nope");
+  assert.equal(invalid.ok, false);
+  if (!invalid.ok) assert.equal(invalid.code, SOCKET_CLOSE_INVALID);
+  const tooBig = parseSocketMessage("a".repeat(MAX_SOCKET_PAYLOAD_BYTES + 1));
+  assert.equal(tooBig.ok, false);
+  if (!tooBig.ok) assert.equal(tooBig.code, SOCKET_CLOSE_TOO_BIG);
+});

package/src/sockets.ts

@@ -0,0 +1,293 @@
+import type {
+  ApiAuth,
+  SocketHandlerContext,
+  SocketRoomClient,
+  SocketsClient,
+} from "@robodev-ai/sdk";
+import { clampTimeoutMs, withHandlerTimeout } from "./handler-timeout.js";
+import { id } from "./ids.js";
+
+export const MAX_PROJECT_SOCKETS = 50;
+export const MAX_SOCKET_PAYLOAD_BYTES = 100 * 1024;
+export const SOCKET_IDLE_PING_MS = 60_000;
+export const SOCKET_CLOSE_POLICY = 1008;
+export const SOCKET_CLOSE_INVALID = 1007;
+export const SOCKET_CLOSE_TOO_BIG = 1009;
+export const SOCKET_CLOSE_RESTART = 1012;
+
+export type SocketWire = {
+  send: (data: string) => void;
+  close: (code?: number, reason?: string) => void;
+  readyState: number;
+};
+
+export type LiveConn = {
+  id: string;
+  projectId: string;
+  name: string;
+  socket: SocketWire;
+  rooms: Set<string>;
+};
+
+export type SocketSessionHelpers = {
+  send: (payload: unknown) => void;
+  close: (code?: number, reason?: string) => void;
+  room: SocketRoomClient;
+};
+
+export type BindSocketSessionInput = {
+  projectId: string;
+  name: string;
+  socket: SocketWire;
+  timeoutMs?: number;
+  createContext: (helpers: SocketSessionHelpers) => SocketHandlerContext;
+  onConnect?: (ctx: SocketHandlerContext) => Promise<void> | void;
+  onMessage: (ctx: SocketHandlerContext, value: unknown) => Promise<void> | void;
+  onClose?: (ctx: SocketHandlerContext) => Promise<void> | void;
+};
+
+const byProject = new Map<string, Set<LiveConn>>();
+const knownNamesByProject = new Map<string, string[]>();
+
+export function setProjectSocketNames(projectId: string, names: string[]): void {
+  knownNamesByProject.set(projectId, names);
+}
+
+export function stringifySocketPayload(payload: unknown): string {
+  const json = JSON.stringify(payload);
+  if (Buffer.byteLength(json, "utf8") > MAX_SOCKET_PAYLOAD_BYTES) {
+    throw new Error("socket payload must be at most 100KB");
+  }
+  return json;
+}
+
+export function parseSocketMessage(
+  raw: string,
+): { ok: true; value: unknown } | { ok: false; code: number } {
+  if (Buffer.byteLength(raw, "utf8") > MAX_SOCKET_PAYLOAD_BYTES) {
+    return { ok: false, code: SOCKET_CLOSE_TOO_BIG };
+  }
+  try {
+    return { ok: true, value: JSON.parse(raw) as unknown };
+  } catch {
+    return { ok: false, code: SOCKET_CLOSE_INVALID };
+  }
+}
+
+export function socketAuthDisplay(auth: ApiAuth): "required" | "public" {
+  return auth === "required" || typeof auth === "function" ? "required" : "public";
+}
+
+export function assertKnownSocket(name: string, known: Iterable<string>): void {
+  if (!new Set(known).has(name)) {
+    throw Object.assign(new Error("unknown_socket"), { code: "unknown_socket" });
+  }
+}
+
+export function projectSocketCount(projectId: string): number {
+  return byProject.get(projectId)?.size ?? 0;
+}
+
+export function closeProjectSockets(
+  projectId: string,
+  code = SOCKET_CLOSE_RESTART,
+  reason = "deployed",
+): void {
+  const set = byProject.get(projectId);
+  if (!set) return;
+  for (const conn of [...set]) {
+    try {
+      conn.socket.close(code, reason);
+    } catch {
+      /* already closed */
+    }
+  }
+  byProject.delete(projectId);
+  knownNamesByProject.delete(projectId);
+}
+
+export function attachProjectSocket(input: {
+  projectId: string;
+  name: string;
+  socket: SocketWire;
+}): { ok: true; conn: LiveConn } | { ok: false; code: number; reason: string } {
+  const current = byProject.get(input.projectId) ?? new Set<LiveConn>();
+  if (current.size >= MAX_PROJECT_SOCKETS) {
+    return { ok: false, code: SOCKET_CLOSE_POLICY, reason: "too_many_connections" };
+  }
+  const conn: LiveConn = {
+    id: id("sock"),
+    projectId: input.projectId,
+    name: input.name,
+    socket: input.socket,
+    rooms: new Set(),
+  };
+  current.add(conn);
+  byProject.set(input.projectId, current);
+  return { ok: true, conn };
+}
+
+export function detachProjectSocket(conn: LiveConn): void {
+  const set = byProject.get(conn.projectId);
+  if (!set) return;
+  set.delete(conn);
+  if (set.size === 0) byProject.delete(conn.projectId);
+}
+
+export function joinSocketRoom(conn: LiveConn, roomId: string): void {
+  conn.rooms.add(roomId);
+}
+
+export function leaveSocketRoom(conn: LiveConn, roomId: string): void {
+  conn.rooms.delete(roomId);
+}
+
+function sendRaw(socket: SocketWire, json: string): void {
+  if (socket.readyState !== 1) return;
+  socket.send(json);
+}
+
+export function sendToSocketName(projectId: string, name: string, json: string): void {
+  const set = byProject.get(projectId);
+  if (!set) return;
+  for (const conn of set) {
+    if (conn.name === name) sendRaw(conn.socket, json);
+  }
+}
+
+export function sendToSocketRoom(
+  projectId: string,
+  name: string,
+  roomId: string,
+  json: string,
+  except?: LiveConn,
+): void {
+  const set = byProject.get(projectId);
+  if (!set) return;
+  for (const conn of set) {
+    if (conn === except) continue;
+    if (conn.name === name && conn.rooms.has(roomId)) sendRaw(conn.socket, json);
+  }
+}
+
+export function dispatchSocketsSend(
+  projectId: string,
+  knownNames: Iterable<string>,
+  input: { name: string; room?: string; payload: unknown },
+): void {
+  assertKnownSocket(input.name, knownNames);
+  const json = stringifySocketPayload(input.payload);
+  if (input.room !== undefined) {
+    sendToSocketRoom(projectId, input.name, input.room, json);
+    return;
+  }
+  sendToSocketName(projectId, input.name, json);
+}
+
+export function createProjectSocketsClient(projectId: string): SocketsClient {
+  function knownNames(): string[] {
+    return knownNamesByProject.get(projectId) ?? [];
+  }
+  return {
+    send(input) {
+      dispatchSocketsSend(projectId, knownNames(), input);
+    },
+    broadcast(input) {
+      dispatchSocketsSend(projectId, knownNames(), input);
+    },
+  };
+}
+
+async function runSocketCallback(
+  timeoutMs: number | undefined,
+  run: () => Promise<void> | void,
+): Promise<void> {
+  await withHandlerTimeout(() => Promise.resolve(run()), clampTimeoutMs(timeoutMs));
+}
+
+export async function bindSocketSession(
+  input: BindSocketSessionInput,
+): Promise<
+  | { ok: false; code: number; reason: string }
+  | { ok: true; onText: (raw: string) => void; onBinary: () => void; onClose: () => void }
+> {
+  const attached = attachProjectSocket({
+    projectId: input.projectId,
+    name: input.name,
+    socket: input.socket,
+  });
+  if (!attached.ok) return attached;
+
+  const helpers: SocketSessionHelpers = {
+    send(payload) {
+      sendRaw(input.socket, stringifySocketPayload(payload));
+    },
+    close(code, reason) {
+      input.socket.close(code, reason);
+    },
+    room: {
+      join: (roomId) => joinSocketRoom(attached.conn, roomId),
+      leave: (roomId) => leaveSocketRoom(attached.conn, roomId),
+      broadcast: (roomId, payload) => {
+        sendToSocketRoom(
+          input.projectId,
+          input.name,
+          roomId,
+          stringifySocketPayload(payload),
+          attached.conn,
+        );
+      },
+    },
+  };
+
+  const ctx = input.createContext(helpers);
+
+  let closed = false;
+  const finish = async () => {
+    if (closed) return;
+    closed = true;
+    detachProjectSocket(attached.conn);
+    if (input.onClose) {
+      try {
+        await runSocketCallback(input.timeoutMs, () => input.onClose?.(ctx));
+      } catch {
+        /* onClose must not throw to the wire */
+      }
+    }
+  };
+
+  if (input.onConnect) {
+    try {
+      await runSocketCallback(input.timeoutMs, () => input.onConnect?.(ctx));
+    } catch {
+      input.socket.close(1011, "handler_failed");
+      await finish();
+      return { ok: false, code: 1011, reason: "handler_failed" };
+    }
+  }
+
+  return {
+    ok: true,
+    onText(raw) {
+      if (closed) return;
+      const parsed = parseSocketMessage(raw);
+      if (!parsed.ok) {
+        input.socket.close(
+          parsed.code,
+          parsed.code === SOCKET_CLOSE_TOO_BIG ? "too_big" : "invalid_json",
+        );
+        return;
+      }
+      void runSocketCallback(input.timeoutMs, () => input.onMessage(ctx, parsed.value)).catch(
+        () => undefined,
+      );
+    },
+    onBinary() {
+      if (closed) return;
+      input.socket.close(1003, "binary not supported");
+    },
+    onClose() {
+      void finish();
+    },
+  };
+}