postgresai 0.16.0-dev.1 → 0.16.0-dev.3

package/test/joe.test.ts

@@ -0,0 +1,726 @@
+import { describe, test, expect, mock, afterEach } from "bun:test";
+import { mkdtempSync, existsSync, readFileSync } from "fs";
+import { tmpdir } from "os";
+import { resolve } from "path";
+import {
+  submitCommand,
+  getCommandStatus,
+  getCommandResult,
+  listProjects,
+  resolveProjectId,
+  isNumericProjectRef,
+  computeIdempotencyKey,
+  runCommand,
+  executeJoeCommand,
+  readStoredSessionId,
+  writeStoredSessionId,
+  clearStoredSessionId,
+  clientSidePlanFlags,
+  formatProjectsTable,
+  formatJoeResult,
+  JOE_COMMANDS,
+  type JoeCommand,
+  type JoeStatus,
+} from "../lib/joe";
+import { handleToolCall, joeToolDefinitions, JOE_TOOL_TO_COMMAND, type McpToolRequest } from "../lib/mcp-server";
+
+const BASE = "https://api.example.com";
+const originalFetch = globalThis.fetch;
+
+interface Captured {
+  url: string;
+  body: Record<string, unknown>;
+  headers: Record<string, string>;
+}
+
+/** Install a fetch mock that routes on the /rpc/<fn> suffix and records requests. */
+function installFetch(routes: Record<string, (body: Record<string, unknown>) => Response>): Captured[] {
+  const captured: Captured[] = [];
+  globalThis.fetch = mock((url: string, options: RequestInit) => {
+    const body = options.body ? (JSON.parse(options.body as string) as Record<string, unknown>) : {};
+    const headers = (options.headers as Record<string, string>) || {};
+    captured.push({ url, body, headers });
+    const fn = new URL(url).pathname.split("/rpc/")[1] ?? "";
+    const handler = routes[fn];
+    if (!handler) {
+      return Promise.resolve(new Response("not found", { status: 404 }));
+    }
+    return Promise.resolve(handler(body));
+  }) as unknown as typeof fetch;
+  return captured;
+}
+
+function json(obj: unknown, status = 200): Response {
+  return new Response(JSON.stringify(obj), { status, headers: { "Content-Type": "application/json" } });
+}
+
+function tempDir(): string {
+  return mkdtempSync(resolve(tmpdir(), "joe-test-"));
+}
+
+afterEach(() => {
+  globalThis.fetch = originalFetch;
+});
+
+describe("submitCommand", () => {
+  test("throws when apiKey missing", async () => {
+    await expect(
+      submitCommand({ apiKey: "", apiBaseUrl: BASE, command: "plan", projectId: 12 })
+    ).rejects.toThrow("API key is required");
+  });
+
+  test("maps to /rpc/joe_command_submit with the correct payload", async () => {
+    const captured = installFetch({
+      joe_command_submit: () => json({ command_id: "4711", session_id: "88", status: "queued" }),
+    });
+    const res = await submitCommand({
+      apiKey: "k",
+      apiBaseUrl: BASE,
+      command: "plan",
+      projectId: 12,
+      sql: "select 1",
+      sessionId: "88",
+      idempotencyKey: "idem-1",
+    });
+    expect(captured[0].url).toBe(`${BASE}/rpc/joe_command_submit`);
+    expect(captured[0].headers["access-token"]).toBe("k");
+    expect(captured[0].body).toEqual({
+      project_id: 12,
+      command: "plan",
+      sql: "select 1",
+      args: null,
+      session_id: "88",
+      idempotency_key: "idem-1",
+    });
+    expect(res).toEqual({ command_id: "4711", session_id: "88", status: "queued" });
+  });
+
+  test("sends null sql/session_id and args for structured commands", async () => {
+    const captured = installFetch({
+      joe_command_submit: () => json({ command_id: "1", session_id: null, status: "queued" }),
+    });
+    await submitCommand({
+      apiKey: "k",
+      apiBaseUrl: BASE,
+      command: "terminate",
+      projectId: 5,
+      args: { pid: 4711 },
+      idempotencyKey: "x",
+    });
+    expect(captured[0].body.sql).toBeNull();
+    expect(captured[0].body.session_id).toBeNull();
+    expect(captured[0].body.args).toEqual({ pid: 4711 });
+  });
+
+  test("surfaces the ai_enabled=false PT403 disabled-AI message", async () => {
+    installFetch({
+      joe_command_submit: () =>
+        json({ message: "AI features disabled for this org — enable in AI Assistant Settings" }, 403),
+    });
+    await expect(
+      submitCommand({ apiKey: "k", apiBaseUrl: BASE, command: "plan", projectId: 12, sql: "select 1" })
+    ).rejects.toThrow(/AI features disabled for this org/);
+  });
+
+  test("surfaces a scope refusal (PT403)", async () => {
+    installFetch({
+      joe_command_submit: () => json({ message: "missing scope joe:exec" }, 403),
+    });
+    await expect(
+      submitCommand({ apiKey: "k", apiBaseUrl: BASE, command: "exec", projectId: 12, sql: "create index" })
+    ).rejects.toThrow(/missing scope joe:exec/);
+  });
+});
+
+describe("getCommandStatus / getCommandResult", () => {
+  test("status maps to /rpc/joe_command_status with command_id", async () => {
+    const captured = installFetch({
+      joe_command_status: () => json({ command_id: "4711", status: "running", error: null }),
+    });
+    const res = await getCommandStatus({ apiKey: "k", apiBaseUrl: BASE, commandId: "4711" });
+    expect(captured[0].url).toBe(`${BASE}/rpc/joe_command_status`);
+    expect(captured[0].body).toEqual({ command_id: "4711" });
+    expect(res.status).toBe("running");
+  });
+
+  test("result maps to /rpc/joe_command_result and returns the plan body", async () => {
+    const captured = installFetch({
+      joe_command_result: () =>
+        json({
+          command_id: "4711",
+          status: "done",
+          queryid: "77",
+          plan_fingerprint: "a1b2",
+          plan_text: "Seq Scan on users",
+          plan_json: { Plan: { "Node Type": "Seq Scan", "Relation Name": "users" } },
+          error: null,
+        }),
+    });
+    const res = await getCommandResult({ apiKey: "k", apiBaseUrl: BASE, commandId: "4711" });
+    expect(captured[0].url).toBe(`${BASE}/rpc/joe_command_result`);
+    expect(res.plan_text).toBe("Seq Scan on users");
+  });
+
+  test("PT404 (not found / out-of-allowlist) surfaces as an error", async () => {
+    installFetch({
+      joe_command_status: () => json({ message: "command not found" }, 404),
+    });
+    await expect(
+      getCommandStatus({ apiKey: "k", apiBaseUrl: BASE, commandId: "9999" })
+    ).rejects.toThrow(/Failed to fetch command status/);
+  });
+});
+
+describe("listProjects", () => {
+  test("normalizes id/joe_api_v2_enabled/tunnel variants", async () => {
+    installFetch({
+      projects_list: () =>
+        json([
+          { id: 12, alias: "main-db", name: "Main DB", label: "Main DB", joe_api_v2_enabled: true, tunnel_ready: true },
+          { project_id: 13, alias: "dw", name: "Warehouse", joe_ready: false, tunnel: false },
+        ]),
+    });
+    const projects = await listProjects({ apiKey: "k", apiBaseUrl: BASE });
+    expect(projects).toEqual([
+      { project_id: 12, alias: "main-db", name: "Main DB", label: "Main DB", joe_ready: true, tunnel: true },
+      { project_id: 13, alias: "dw", name: "Warehouse", label: null, joe_ready: false, tunnel: false },
+    ]);
+  });
+
+  test("passes org_id when provided", async () => {
+    const captured = installFetch({ projects_list: () => json([]) });
+    await listProjects({ apiKey: "k", apiBaseUrl: BASE, orgId: 7 });
+    expect(captured[0].body).toEqual({ org_id: 7 });
+  });
+});
+
+describe("resolveProjectId (id-or-alias)", () => {
+  test("isNumericProjectRef distinguishes ids from aliases", () => {
+    expect(isNumericProjectRef("12")).toBe(true);
+    expect(isNumericProjectRef(" 12 ")).toBe(true);
+    expect(isNumericProjectRef("main-db")).toBe(false);
+    expect(isNumericProjectRef("12a")).toBe(false);
+  });
+
+  test("numeric ref resolves WITHOUT a projects lookup", async () => {
+    const captured = installFetch({ projects_list: () => json([]) });
+    const id = await resolveProjectId({ apiKey: "k", apiBaseUrl: BASE, project: "12" });
+    expect(id).toBe(12);
+    expect(captured.length).toBe(0);
+  });
+
+  test("alias resolves via the projects list (case-insensitive)", async () => {
+    installFetch({
+      projects_list: () => json([{ id: 12, alias: "main-db", name: "Main DB", joe_ready: true }]),
+    });
+    const id = await resolveProjectId({ apiKey: "k", apiBaseUrl: BASE, project: "MAIN-DB" });
+    expect(id).toBe(12);
+  });
+
+  test("--project 12 ≡ --project main-db", async () => {
+    installFetch({
+      projects_list: () => json([{ id: 12, alias: "main-db", name: "Main DB", joe_ready: true }]),
+    });
+    const byId = await resolveProjectId({ apiKey: "k", apiBaseUrl: BASE, project: "12" });
+    const byAlias = await resolveProjectId({ apiKey: "k", apiBaseUrl: BASE, project: "main-db" });
+    expect(byId).toBe(byAlias);
+  });
+
+  test("unknown alias throws", async () => {
+    installFetch({ projects_list: () => json([{ id: 12, alias: "main-db", name: "Main DB" }]) });
+    await expect(
+      resolveProjectId({ apiKey: "k", apiBaseUrl: BASE, project: "nope" })
+    ).rejects.toThrow(/Project not found for alias\/name 'nope'/);
+  });
+
+  test("empty project ref throws", async () => {
+    await expect(
+      resolveProjectId({ apiKey: "k", apiBaseUrl: BASE, project: "" })
+    ).rejects.toThrow(/project is required/);
+  });
+});
+
+describe("computeIdempotencyKey", () => {
+  test("exec/explain keys are STABLE (deterministic) across calls", () => {
+    const a = computeIdempotencyKey("exec", 12, "create index on users (email)", null);
+    const b = computeIdempotencyKey("exec", 12, "create index on users (email)", null);
+    expect(a).toBe(b);
+    // A different statement diverges.
+    const c = computeIdempotencyKey("exec", 12, "create index on orders (id)", null);
+    expect(a).not.toBe(c);
+  });
+
+  test("plan keys are fresh per invocation", () => {
+    const a = computeIdempotencyKey("plan", 12, "select 1", null);
+    const b = computeIdempotencyKey("plan", 12, "select 1", null);
+    expect(a).not.toBe(b);
+  });
+});
+
+describe("runCommand — submit-then-poll", () => {
+  test("polls until terminal then fetches the result", async () => {
+    let statusCalls = 0;
+    installFetch({
+      joe_command_submit: () => json({ command_id: "4711", session_id: "88", status: "queued" }),
+      joe_command_status: () => {
+        statusCalls += 1;
+        return json({ command_id: "4711", status: statusCalls < 2 ? "running" : "done", error: null });
+      },
+      joe_command_result: () =>
+        json({ command_id: "4711", status: "done", plan_text: "Index Scan", plan_json: {}, error: null }),
+    });
+    const outcome = await runCommand({
+      apiKey: "k",
+      apiBaseUrl: BASE,
+      command: "plan",
+      projectId: 12,
+      sql: "select 1",
+      pollIntervalMs: 0,
+      sleep: async () => {},
+    });
+    expect(outcome.status).toBe("done");
+    expect(outcome.budgetExpired).toBe(false);
+    expect(outcome.commandId).toBe("4711");
+    expect(outcome.sessionId).toBe("88");
+    expect(outcome.result?.plan_text).toBe("Index Scan");
+    expect(statusCalls).toBeGreaterThanOrEqual(2);
+  });
+
+  test("budget expiry returns a resume handle without a result", async () => {
+    installFetch({
+      joe_command_submit: () => json({ command_id: "4712", session_id: "88", status: "queued" }),
+      joe_command_status: () => json({ command_id: "4712", status: "running", error: null }),
+      joe_command_result: () => json({ command_id: "4712", status: "done", error: null }),
+    });
+    let clock = 1000;
+    const outcome = await runCommand({
+      apiKey: "k",
+      apiBaseUrl: BASE,
+      command: "plan",
+      projectId: 12,
+      sql: "select 1",
+      budgetMs: 5,
+      pollIntervalMs: 0,
+      now: () => (clock += 10),
+      sleep: async () => {},
+    });
+    expect(outcome.budgetExpired).toBe(true);
+    expect(outcome.status).toBe("running");
+    expect(outcome.result).toBeNull();
+    expect(outcome.commandId).toBe("4712");
+  });
+
+  test("terminal error state fetches the result (error body)", async () => {
+    installFetch({
+      joe_command_submit: () => json({ command_id: "5", session_id: null, status: "queued" }),
+      joe_command_status: () => json({ command_id: "5", status: "error", error: "PT500" }),
+      joe_command_result: () => json({ command_id: "5", status: "error", error: "boom", plan_json: {} }),
+    });
+    const outcome = await runCommand({
+      apiKey: "k",
+      apiBaseUrl: BASE,
+      command: "plan",
+      projectId: 12,
+      sql: "select 1",
+      pollIntervalMs: 0,
+      sleep: async () => {},
+    });
+    expect(outcome.status).toBe("error");
+    expect(outcome.result?.error).toBe("boom");
+  });
+
+  test("exec submits a STABLE idempotency key", async () => {
+    const captured = installFetch({
+      joe_command_submit: () => json({ command_id: "6", session_id: "88", status: "queued" }),
+      joe_command_status: () => json({ command_id: "6", status: "done", error: null }),
+      joe_command_result: () => json({ command_id: "6", status: "done", row_count: 0, notices: ["CREATE INDEX"], error: null }),
+    });
+    await runCommand({
+      apiKey: "k",
+      apiBaseUrl: BASE,
+      command: "exec",
+      projectId: 12,
+      sql: "create index on users (email)",
+      pollIntervalMs: 0,
+      sleep: async () => {},
+    });
+    const expected = computeIdempotencyKey("exec", 12, "create index on users (email)", null);
+    expect(captured[0].body.idempotency_key).toBe(expected);
+  });
+});
+
+describe("executeJoeCommand — session threading + persistence", () => {
+  test("auto-reuses the stored session id for the project", async () => {
+    const dir = tempDir();
+    writeStoredSessionId(12, "88", dir);
+    const captured = installFetch({
+      joe_command_submit: () => json({ command_id: "1", session_id: "88", status: "queued" }),
+      joe_command_status: () => json({ command_id: "1", status: "done", error: null }),
+      joe_command_result: () => json({ command_id: "1", status: "done", plan_text: "ok", error: null }),
+    });
+    await executeJoeCommand({
+      apiKey: "k",
+      apiBaseUrl: BASE,
+      command: "plan",
+      project: "12",
+      sql: "select 1",
+      sessionDir: dir,
+      pollIntervalMs: 0,
+      sleep: async () => {},
+    });
+    expect(captured[0].body.session_id).toBe("88");
+  });
+
+  test("persists the session id returned by submit", async () => {
+    const dir = tempDir();
+    installFetch({
+      joe_command_submit: () => json({ command_id: "1", session_id: "99", status: "queued" }),
+      joe_command_status: () => json({ command_id: "1", status: "done", error: null }),
+      joe_command_result: () => json({ command_id: "1", status: "done", plan_text: "ok", error: null }),
+    });
+    await executeJoeCommand({
+      apiKey: "k",
+      apiBaseUrl: BASE,
+      command: "plan",
+      project: "12",
+      sql: "select 1",
+      sessionDir: dir,
+      pollIntervalMs: 0,
+      sleep: async () => {},
+    });
+    expect(readStoredSessionId(12, dir)).toBe("99");
+  });
+
+  test("--new-session clears the stored session and submits null", async () => {
+    const dir = tempDir();
+    writeStoredSessionId(12, "88", dir);
+    const captured = installFetch({
+      joe_command_submit: () => json({ command_id: "1", session_id: "100", status: "queued" }),
+      joe_command_status: () => json({ command_id: "1", status: "done", error: null }),
+      joe_command_result: () => json({ command_id: "1", status: "done", plan_text: "ok", error: null }),
+    });
+    await executeJoeCommand({
+      apiKey: "k",
+      apiBaseUrl: BASE,
+      command: "plan",
+      project: "12",
+      sql: "select 1",
+      newSession: true,
+      sessionDir: dir,
+      pollIntervalMs: 0,
+      sleep: async () => {},
+    });
+    expect(captured[0].body.session_id).toBeNull();
+    // The freshly returned session id is persisted for the next command.
+    expect(readStoredSessionId(12, dir)).toBe("100");
+  });
+
+  test("explicit --session overrides the stored session", async () => {
+    const dir = tempDir();
+    writeStoredSessionId(12, "88", dir);
+    const captured = installFetch({
+      joe_command_submit: () => json({ command_id: "1", session_id: "77", status: "queued" }),
+      joe_command_status: () => json({ command_id: "1", status: "done", error: null }),
+      joe_command_result: () => json({ command_id: "1", status: "done", plan_text: "ok", error: null }),
+    });
+    await executeJoeCommand({
+      apiKey: "k",
+      apiBaseUrl: BASE,
+      command: "plan",
+      project: "12",
+      sql: "select 1",
+      session: "77",
+      sessionDir: dir,
+      pollIntervalMs: 0,
+      sleep: async () => {},
+    });
+    expect(captured[0].body.session_id).toBe("77");
+  });
+
+  test("resolves an alias to its project id before submitting", async () => {
+    const dir = tempDir();
+    const captured = installFetch({
+      projects_list: () => json([{ id: 12, alias: "main-db", name: "Main DB", joe_ready: true }]),
+      joe_command_submit: () => json({ command_id: "1", session_id: "88", status: "queued" }),
+      joe_command_status: () => json({ command_id: "1", status: "done", error: null }),
+      joe_command_result: () => json({ command_id: "1", status: "done", plan_text: "ok", error: null }),
+    });
+    const outcome = await executeJoeCommand({
+      apiKey: "k",
+      apiBaseUrl: BASE,
+      command: "plan",
+      project: "main-db",
+      sql: "select 1",
+      sessionDir: dir,
+      pollIntervalMs: 0,
+      sleep: async () => {},
+    });
+    expect(outcome.projectId).toBe(12);
+    const submit = captured.find((c) => c.url.endsWith("/rpc/joe_command_submit"));
+    expect(submit?.body.project_id).toBe(12);
+  });
+});
+
+describe("session store helpers", () => {
+  test("write / read / clear round-trip", () => {
+    const dir = tempDir();
+    expect(readStoredSessionId(42, dir)).toBeNull();
+    writeStoredSessionId(42, "500", dir);
+    expect(readStoredSessionId(42, dir)).toBe("500");
+    expect(existsSync(resolve(dir, "joe-sessions.json"))).toBe(true);
+    clearStoredSessionId(42, dir);
+    expect(readStoredSessionId(42, dir)).toBeNull();
+  });
+
+  test("does not clobber other projects' sessions", () => {
+    const dir = tempDir();
+    writeStoredSessionId(1, "a", dir);
+    writeStoredSessionId(2, "b", dir);
+    clearStoredSessionId(1, dir);
+    expect(readStoredSessionId(1, dir)).toBeNull();
+    expect(readStoredSessionId(2, dir)).toBe("b");
+    const raw = JSON.parse(readFileSync(resolve(dir, "joe-sessions.json"), "utf8"));
+    expect(raw).toEqual({ "2": "b" });
+  });
+});
+
+describe("presentation helpers", () => {
+  test("clientSidePlanFlags flags Seq Scans in nested plans", () => {
+    const flags = clientSidePlanFlags({
+      Plan: {
+        "Node Type": "Nested Loop",
+        Plans: [{ "Node Type": "Seq Scan", "Relation Name": "users" }],
+      },
+    });
+    expect(flags.length).toBe(1);
+    expect(flags[0]).toContain("Seq Scan on users");
+  });
+
+  test("formatProjectsTable renders the brief's column header", () => {
+    const table = formatProjectsTable([
+      { project_id: 12, alias: "main-db", name: "Main DB", label: null, joe_ready: true, tunnel: true },
+    ]);
+    const [header, row] = table.split("\n");
+    expect(header).toContain("PROJECT_ID");
+    expect(header).toContain("ALIAS");
+    expect(header).toContain("JOE");
+    expect(header).toContain("TUNNEL");
+    expect(row).toContain("12");
+    expect(row).toContain("main-db");
+    expect(row).toContain("ready");
+    expect(row).toContain("yes");
+  });
+
+  test("formatJoeResult renders per-command bodies", () => {
+    const plan = formatJoeResult("plan", {
+      command_id: "1",
+      status: "done",
+      error: null,
+      plan_text: "Seq Scan on users",
+      plan_json: { Plan: { "Node Type": "Seq Scan", "Relation Name": "users" } },
+      queryid: "77",
+      plan_fingerprint: "a1",
+    });
+    expect(plan).toContain("Seq Scan on users");
+    expect(plan).toContain("⚑");
+
+    const exec = formatJoeResult("exec", {
+      command_id: "2",
+      status: "done",
+      error: null,
+      row_count: 0,
+      notices: ["CREATE INDEX"],
+    });
+    expect(exec).toContain("CREATE INDEX");
+
+    const terminate = formatJoeResult("terminate", {
+      command_id: "3",
+      status: "done",
+      error: null,
+      terminated: true,
+      pid: 4711,
+    });
+    expect(terminate).toContain("pid 4711");
+  });
+});
+
+describe("MCP tools", () => {
+  const ROOT = { apiKey: "k", apiBaseUrl: BASE };
+  const originalXdg = process.env.XDG_CONFIG_HOME;
+
+  function makeReq(name: string, args?: Record<string, unknown>): McpToolRequest {
+    return { params: { name, arguments: args } };
+  }
+
+  afterEach(() => {
+    globalThis.fetch = originalFetch;
+    if (originalXdg === undefined) delete process.env.XDG_CONFIG_HOME;
+    else process.env.XDG_CONFIG_HOME = originalXdg;
+  });
+
+  test("joeToolDefinitions exposes every SPEC §5 name + list_projects", () => {
+    const names = joeToolDefinitions().map((t) => t.name);
+    expect(names).toEqual([
+      "plan_query",
+      "explain_query",
+      "exec_sql",
+      "hypo_index",
+      "activity",
+      "terminate_backend",
+      "reset_clone",
+      "describe",
+      "list_projects",
+    ]);
+    // Tool-name -> command map covers the 8 command tools.
+    expect(Object.keys(JOE_TOOL_TO_COMMAND).length).toBe(8);
+  });
+
+  test("exec_sql carries the stronger raw-data warning; plan/explain note the plan egress", () => {
+    const defs = joeToolDefinitions();
+    const exec = defs.find((t) => t.name === "exec_sql");
+    expect(exec?.description).toMatch(/REAL .*TABLE ROWS/i);
+    expect(exec?.description).toContain("LLM context");
+    const plan = defs.find((t) => t.name === "plan_query");
+    expect(plan?.description).toMatch(/plan-only/i);
+    expect(plan?.description).toMatch(/NEVER executes/i);
+    const explain = defs.find((t) => t.name === "explain_query");
+    expect(explain?.description).toMatch(/EXECUTES/);
+    expect(explain?.description).toMatch(/NOT raw SELECT result rows/i);
+  });
+
+  test("plan_query calls executeJoeCommand and returns the plan result", () => {
+    process.env.XDG_CONFIG_HOME = tempDir();
+    const captured = installFetch({
+      joe_command_submit: () => json({ command_id: "4711", session_id: "88", status: "queued" }),
+      joe_command_status: () => json({ command_id: "4711", status: "done", error: null }),
+      joe_command_result: () =>
+        json({ command_id: "4711", status: "done", plan_text: "Seq Scan on users", plan_json: {}, error: null }),
+    });
+    return handleToolCall(makeReq("plan_query", { sql: "select 1", project_id: "12" }), ROOT).then((res) => {
+      expect(res.isError).toBeFalsy();
+      const submit = captured.find((c) => c.url.endsWith("/rpc/joe_command_submit"));
+      expect(submit?.body.command).toBe("plan");
+      expect(submit?.body.project_id).toBe(12);
+      const payload = JSON.parse(res.content[0].text) as { status: string; result: { plan_text: string } };
+      expect(payload.status).toBe("done");
+      expect(payload.result.plan_text).toBe("Seq Scan on users");
+    });
+  });
+
+  test("exec_sql maps sql; requires it", async () => {
+    process.env.XDG_CONFIG_HOME = tempDir();
+    const captured = installFetch({
+      joe_command_submit: () => json({ command_id: "1", session_id: "88", status: "queued" }),
+      joe_command_status: () => json({ command_id: "1", status: "done", error: null }),
+      joe_command_result: () => json({ command_id: "1", command: "exec", status: "done", row_count: 0, notices: ["CREATE INDEX"], error: null }),
+    });
+    const ok = await handleToolCall(makeReq("exec_sql", { sql: "create index on users (email)", project_id: "12" }), ROOT);
+    expect(ok.isError).toBeFalsy();
+    const submit = captured.find((c) => c.url.endsWith("/rpc/joe_command_submit"));
+    expect(submit?.body.command).toBe("exec");
+
+    const missing = await handleToolCall(makeReq("exec_sql", { project_id: "12" }), ROOT);
+    expect(missing.isError).toBe(true);
+    expect(missing.content[0].text).toContain("sql is required");
+  });
+
+  test("hypo_index requires query and maps args.query", async () => {
+    process.env.XDG_CONFIG_HOME = tempDir();
+    const captured = installFetch({
+      joe_command_submit: () => json({ command_id: "1", session_id: "88", status: "queued" }),
+      joe_command_status: () => json({ command_id: "1", status: "done", error: null }),
+      joe_command_result: () => json({ command_id: "1", status: "done", hypo_used: true, hypo_plan: {}, error: null }),
+    });
+    const missing = await handleToolCall(makeReq("hypo_index", { sql: "create index on orders (customer_id)", project_id: "12" }), ROOT);
+    expect(missing.isError).toBe(true);
+    expect(missing.content[0].text).toContain("query is required");
+
+    const ok = await handleToolCall(
+      makeReq("hypo_index", { sql: "create index on orders (customer_id)", query: "select 1", project_id: "12" }),
+      ROOT
+    );
+    expect(ok.isError).toBeFalsy();
+    const submit = captured.find((c) => c.url.endsWith("/rpc/joe_command_submit"));
+    expect(submit?.body.command).toBe("hypo");
+    expect(submit?.body.args).toEqual({ query: "select 1" });
+  });
+
+  test("terminate_backend maps args.pid; describe maps args.object", async () => {
+    process.env.XDG_CONFIG_HOME = tempDir();
+    const captured = installFetch({
+      joe_command_submit: () => json({ command_id: "1", session_id: null, status: "queued" }),
+      joe_command_status: () => json({ command_id: "1", status: "done", error: null }),
+      joe_command_result: () => json({ command_id: "1", status: "done", terminated: true, pid: 4711, snapshot: {}, error: null }),
+    });
+    await handleToolCall(makeReq("terminate_backend", { pid: 4711, project_id: "12" }), ROOT);
+    let submit = captured.find((c) => c.url.endsWith("/rpc/joe_command_submit"));
+    expect(submit?.body.command).toBe("terminate");
+    expect(submit?.body.args).toEqual({ pid: 4711 });
+
+    captured.length = 0;
+    await handleToolCall(makeReq("describe", { object: "users", variant: "\\d+", project_id: "12" }), ROOT);
+    submit = captured.find((c) => c.url.endsWith("/rpc/joe_command_submit"));
+    expect(submit?.body.command).toBe("describe");
+    expect(submit?.body.args).toEqual({ object: "users", variant: "\\d+" });
+  });
+
+  test("list_projects tool returns the normalized projects", async () => {
+    process.env.XDG_CONFIG_HOME = tempDir();
+    installFetch({
+      projects_list: () => json([{ id: 12, alias: "main-db", name: "Main DB", joe_ready: true, tunnel: true }]),
+    });
+    const res = await handleToolCall(makeReq("list_projects", {}), ROOT);
+    expect(res.isError).toBeFalsy();
+    const projects = JSON.parse(res.content[0].text) as Array<{ project_id: number; alias: string }>;
+    expect(projects[0].project_id).toBe(12);
+    expect(projects[0].alias).toBe("main-db");
+  });
+
+  test("resolves an alias project_id before submitting", async () => {
+    process.env.XDG_CONFIG_HOME = tempDir();
+    const captured = installFetch({
+      projects_list: () => json([{ id: 12, alias: "main-db", name: "Main DB", joe_ready: true }]),
+      joe_command_submit: () => json({ command_id: "1", session_id: "88", status: "queued" }),
+      joe_command_status: () => json({ command_id: "1", status: "done", error: null }),
+      joe_command_result: () => json({ command_id: "1", status: "done", plan_text: "ok", error: null }),
+    });
+    await handleToolCall(makeReq("plan_query", { sql: "select 1", project_id: "main-db" }), ROOT);
+    const submit = captured.find((c) => c.url.endsWith("/rpc/joe_command_submit"));
+    expect(submit?.body.project_id).toBe(12);
+  });
+
+  test("ai_enabled=false PT403 surfaces as an MCP error", async () => {
+    process.env.XDG_CONFIG_HOME = tempDir();
+    installFetch({
+      joe_command_submit: () =>
+        json({ message: "AI features disabled for this org — enable in AI Assistant Settings" }, 403),
+    });
+    const res = await handleToolCall(makeReq("plan_query", { sql: "select 1", project_id: "12" }), ROOT);
+    expect(res.isError).toBe(true);
+    expect(res.content[0].text).toContain("AI features disabled for this org");
+  });
+
+  test("missing project_id is an error", async () => {
+    process.env.XDG_CONFIG_HOME = tempDir();
+    const res = await handleToolCall(makeReq("plan_query", { sql: "select 1" }), ROOT);
+    expect(res.isError).toBe(true);
+    expect(res.content[0].text).toContain("project_id is required");
+  });
+});
+
+describe("command set", () => {
+  test("JOE_COMMANDS is the SPEC command set", () => {
+    const expected: JoeCommand[] = ["plan", "explain", "exec", "hypo", "activity", "terminate", "reset", "describe"];
+    expect([...JOE_COMMANDS]).toEqual(expected);
+  });
+
+  test("each command submits with its fixed command name", async () => {
+    for (const command of JOE_COMMANDS) {
+      const captured = installFetch({
+        joe_command_submit: () => json({ command_id: "1", session_id: null, status: "queued" as JoeStatus }),
+      });
+      await submitCommand({ apiKey: "k", apiBaseUrl: BASE, command, projectId: 12, idempotencyKey: "x" });
+      expect(captured[0].body.command).toBe(command);
+    }
+  });
+});