postgresai 0.16.0-dev.2 → 0.16.0-dev.4

package/test/joe.cli.test.ts

@@ -0,0 +1,298 @@
+import { describe, test, expect } from "bun:test";
+import { resolve } from "path";
+import { mkdtempSync } from "fs";
+import { tmpdir } from "os";
+
+// Sync runner for OFFLINE tests only (help / fail-fast). Never use it for a test
+// that round-trips the in-process Bun.serve: spawnSync blocks the JS event loop, so
+// the fake server can't answer and the two deadlock (see runCliAsync).
+function runCli(args: string[], env: Record<string, string> = {}) {
+  const cliPath = resolve(import.meta.dir, "..", "bin", "postgres-ai.ts");
+  const bunBin = typeof process.execPath === "string" && process.execPath.length > 0 ? process.execPath : "bun";
+  const result = Bun.spawnSync([bunBin, cliPath, ...args], {
+    env: { ...process.env, ...env },
+  });
+  return {
+    status: result.exitCode,
+    stdout: new TextDecoder().decode(result.stdout),
+    stderr: new TextDecoder().decode(result.stderr),
+  };
+}
+
+// Async runner for tests that hit the fake server — keeps the event loop free.
+async function runCliAsync(args: string[], env: Record<string, string> = {}) {
+  const cliPath = resolve(import.meta.dir, "..", "bin", "postgres-ai.ts");
+  const bunBin = typeof process.execPath === "string" && process.execPath.length > 0 ? process.execPath : "bun";
+  const proc = Bun.spawn([bunBin, cliPath, ...args], {
+    env: { ...process.env, ...env },
+    stdout: "pipe",
+    stderr: "pipe",
+  });
+  const [status, stdout, stderr] = await Promise.all([
+    proc.exited,
+    new Response(proc.stdout).text(),
+    new Response(proc.stderr).text(),
+  ]);
+  return { status, stdout, stderr };
+}
+
+/** Isolate config so tests never touch real user state; a shared dir threads sessions. */
+function isolatedEnv(extra: Record<string, string> = {}) {
+  const cfgHome = mkdtempSync(resolve(tmpdir(), "joe-cli-test-"));
+  return { XDG_CONFIG_HOME: cfgHome, HOME: cfgHome, ...extra };
+}
+
+interface RecordedRequest {
+  pathname: string;
+  body: Record<string, unknown>;
+}
+
+/**
+ * Fake PostgREST server for the Joe rpcs (backend not built yet). A `plan` on a
+ * project whose id is COLD_PROJECT stays `running` (to exercise the resume path);
+ * project id 999 returns the ai_enabled=false PT403; everything else completes.
+ */
+function startFakeApi() {
+  const requests: RecordedRequest[] = [];
+  const COLD_PROJECT = 77;
+  let nextId = 4711;
+  const commands = new Map<string, { command: string; projectId: number }>();
+
+  const server = Bun.serve({
+    hostname: "127.0.0.1",
+    port: 0,
+    async fetch(req) {
+      const url = new URL(req.url);
+      const bodyText = await req.text();
+      let body: Record<string, unknown> = {};
+      try {
+        body = bodyText ? (JSON.parse(bodyText) as Record<string, unknown>) : {};
+      } catch {
+        body = {};
+      }
+      requests.push({ pathname: url.pathname, body });
+
+      const respond = (obj: unknown, status = 200): Response =>
+        new Response(JSON.stringify(obj), { status, headers: { "Content-Type": "application/json" } });
+
+      if (req.method === "POST" && url.pathname.endsWith("/rpc/projects_list")) {
+        return respond([
+          { id: 12, alias: "main-db", name: "Main DB", label: "Main DB", joe_api_v2_enabled: true, tunnel_ready: true },
+          { id: 77, alias: "cold-db", name: "Cold DB", joe_ready: true, tunnel: true },
+        ]);
+      }
+
+      if (req.method === "POST" && url.pathname.endsWith("/rpc/joe_command_submit")) {
+        const projectId = Number(body.project_id);
+        if (projectId === 999) {
+          return respond({ message: "AI features disabled for this org — enable in AI Assistant Settings" }, 403);
+        }
+        const commandId = String(nextId++);
+        commands.set(commandId, { command: String(body.command), projectId });
+        // Echo the incoming session_id, or hand back a fresh one for a null session.
+        const sessionId = body.session_id != null ? String(body.session_id) : "88";
+        return respond({ command_id: commandId, session_id: sessionId, status: "queued" });
+      }
+
+      if (req.method === "POST" && url.pathname.endsWith("/rpc/joe_command_status")) {
+        const commandId = String(body.command_id);
+        const cmd = commands.get(commandId);
+        const status = cmd && cmd.projectId === COLD_PROJECT ? "running" : "done";
+        return respond({ command_id: commandId, status, error: null });
+      }
+
+      if (req.method === "POST" && url.pathname.endsWith("/rpc/joe_command_result")) {
+        const commandId = String(body.command_id);
+        const cmd = commands.get(commandId);
+        if (cmd?.command === "exec") {
+          return respond({ command_id: commandId, command: "exec", status: "done", row_count: 0, result_rows: [], notices: ["CREATE INDEX"], error: null });
+        }
+        return respond({
+          command_id: commandId,
+          status: "done",
+          queryid: "7712349901234567890",
+          plan_fingerprint: "a1b2c3d4",
+          plan_text: "Seq Scan on users  (cost=0.00..48210.00 rows=1 width=812)",
+          plan_json: { Plan: { "Node Type": "Seq Scan", "Relation Name": "users" } },
+          error: null,
+        });
+      }
+
+      return new Response("not found", { status: 404 });
+    },
+  });
+
+  const baseUrl = `http://${server.hostname}:${server.port}/api/general`;
+  return { baseUrl, requests, stop: () => server.stop(true) };
+}
+
+describe("CLI Joe command surface", () => {
+  test("top-level help lists the Joe verbs and projects/result", () => {
+    const r = runCli(["--help"], isolatedEnv());
+    const out = `${r.stdout}\n${r.stderr}`;
+    for (const verb of ["plan", "explain", "exec", "hypo", "activity", "terminate", "reset", "describe", "projects", "result", "status"]) {
+      expect(out).toContain(verb);
+    }
+  });
+
+  test("pgai projects prints the brief's table", async () => {
+    const api = startFakeApi();
+    try {
+      const r = await runCliAsync(["projects"], isolatedEnv({ PGAI_API_KEY: "k", PGAI_API_BASE_URL: api.baseUrl }));
+      expect(r.status).toBe(0);
+      expect(r.stdout).toContain("PROJECT_ID");
+      expect(r.stdout).toContain("ALIAS");
+      expect(r.stdout).toContain("JOE");
+      expect(r.stdout).toContain("TUNNEL");
+      expect(r.stdout).toContain("main-db");
+      expect(r.stdout).toContain("ready");
+      const listReq = api.requests.find((x) => x.pathname.endsWith("/rpc/projects_list"));
+      expect(listReq).toBeDefined();
+    } finally {
+      api.stop();
+    }
+  });
+
+  test("pgai plan --project 12 submits command='plan' and prints the plan", async () => {
+    const api = startFakeApi();
+    try {
+      const r = await runCliAsync(
+        ["plan", "select * from users where email = 'x@acme.io'", "--project", "12"],
+        isolatedEnv({ PGAI_API_KEY: "k", PGAI_API_BASE_URL: api.baseUrl })
+      );
+      expect(r.status).toBe(0);
+      expect(r.stdout).toContain("Seq Scan on users");
+      // client-side plan flag
+      expect(r.stdout).toContain("⚑");
+      const submit = api.requests.find((x) => x.pathname.endsWith("/rpc/joe_command_submit"));
+      expect(submit?.body.command).toBe("plan");
+      expect(submit?.body.project_id).toBe(12);
+      // numeric project => no projects_list lookup
+      expect(api.requests.some((x) => x.pathname.endsWith("/rpc/projects_list"))).toBe(false);
+    } finally {
+      api.stop();
+    }
+  });
+
+  test("pgai plan --project main-db resolves the alias via projects_list", async () => {
+    const api = startFakeApi();
+    try {
+      const r = await runCliAsync(
+        ["plan", "select 1", "--project", "main-db"],
+        isolatedEnv({ PGAI_API_KEY: "k", PGAI_API_BASE_URL: api.baseUrl })
+      );
+      expect(r.status).toBe(0);
+      expect(api.requests.some((x) => x.pathname.endsWith("/rpc/projects_list"))).toBe(true);
+      const submit = api.requests.find((x) => x.pathname.endsWith("/rpc/joe_command_submit"));
+      expect(submit?.body.project_id).toBe(12);
+    } finally {
+      api.stop();
+    }
+  });
+
+  test("optimize loop: plan then exec reuse the same session on one project", async () => {
+    const api = startFakeApi();
+    const env = isolatedEnv({ PGAI_API_KEY: "k", PGAI_API_BASE_URL: api.baseUrl });
+    try {
+      // First command starts a fresh session; the server hands back session 88.
+      const first = await runCliAsync(["plan", "select 1", "--project", "12", "--new-session"], env);
+      expect(first.status).toBe(0);
+      // Second command auto-reuses the stored session (88) on the same project.
+      const second = await runCliAsync(["exec", "create index on users (email)", "--project", "12"], env);
+      expect(second.status).toBe(0);
+      expect(second.stdout).toContain("CREATE INDEX");
+
+      const submits = api.requests.filter((x) => x.pathname.endsWith("/rpc/joe_command_submit"));
+      expect(submits.length).toBe(2);
+      // The re-plan/exec threads the session id returned by the first submit.
+      expect(submits[0].body.session_id).toBeNull(); // --new-session
+      expect(submits[1].body.session_id).toBe("88");
+      expect(submits[1].body.command).toBe("exec");
+      // exec mints a stable idempotency key.
+      expect(typeof submits[1].body.idempotency_key).toBe("string");
+    } finally {
+      api.stop();
+    }
+  });
+
+  test("cold clone exceeds the budget -> resume handle (exit 0), then pgai result", async () => {
+    const api = startFakeApi();
+    const env = isolatedEnv({ PGAI_API_KEY: "k", PGAI_API_BASE_URL: api.baseUrl });
+    try {
+      // project 77 (COLD) stays `running`; --budget 0 forces immediate resume.
+      const r = await runCliAsync(["plan", "select 1", "--project", "77", "--budget", "0"], env);
+      expect(r.status).toBe(0);
+      expect(r.stdout).toContain("resume:");
+      expect(r.stdout).toMatch(/pgai result \d+/);
+
+      const match = r.stdout.match(/pgai result (\d+)/);
+      const commandId = match ? match[1] : "";
+      expect(commandId).not.toBe("");
+
+      const resumed = await runCliAsync(["result", commandId], env);
+      expect(resumed.status).toBe(0);
+      expect(resumed.stdout).toContain("Seq Scan on users");
+    } finally {
+      api.stop();
+    }
+  });
+
+  test("ai_enabled=false submit is refused (PT403) with a clear message and exit 1", async () => {
+    const api = startFakeApi();
+    try {
+      const r = await runCliAsync(
+        ["plan", "select 1", "--project", "999"],
+        isolatedEnv({ PGAI_API_KEY: "k", PGAI_API_BASE_URL: api.baseUrl })
+      );
+      expect(r.status).toBe(1);
+      expect(`${r.stdout}\n${r.stderr}`).toContain("AI features disabled for this org");
+    } finally {
+      api.stop();
+    }
+  });
+
+  test("terminate requires a numeric pid and maps args.pid", async () => {
+    const api = startFakeApi();
+    try {
+      const r = await runCliAsync(
+        ["terminate", "4711", "--project", "12"],
+        isolatedEnv({ PGAI_API_KEY: "k", PGAI_API_BASE_URL: api.baseUrl })
+      );
+      expect(r.status).toBe(0);
+      const submit = api.requests.find((x) => x.pathname.endsWith("/rpc/joe_command_submit"));
+      expect(submit?.body.command).toBe("terminate");
+      expect(submit?.body.args).toEqual({ pid: 4711 });
+      expect(submit?.body.sql).toBeNull();
+    } finally {
+      api.stop();
+    }
+  });
+
+  test("describe maps args.object (and --variant)", async () => {
+    const api = startFakeApi();
+    try {
+      const r = await runCliAsync(
+        ["describe", "users", "--project", "12", "--variant", "\\d+"],
+        isolatedEnv({ PGAI_API_KEY: "k", PGAI_API_BASE_URL: api.baseUrl })
+      );
+      expect(r.status).toBe(0);
+      const submit = api.requests.find((x) => x.pathname.endsWith("/rpc/joe_command_submit"));
+      expect(submit?.body.command).toBe("describe");
+      expect(submit?.body.args).toEqual({ object: "users", variant: "\\d+" });
+    } finally {
+      api.stop();
+    }
+  });
+
+  test("missing --project fails fast", () => {
+    const r = runCli(["plan", "select 1"], isolatedEnv({ PGAI_API_KEY: "k" }));
+    expect(r.status).toBe(1);
+    expect(`${r.stdout}\n${r.stderr}`).toContain("Project is required");
+  });
+
+  test("missing API key fails fast", () => {
+    const r = runCli(["plan", "select 1", "--project", "12"], isolatedEnv());
+    expect(r.status).toBe(1);
+    expect(`${r.stdout}\n${r.stderr}`).toContain("API key is required");
+  });
+});