postgresai 0.16.0-dev.0 → 0.16.0-dev.10

package/lib/mcp-server.ts

@@ -15,7 +15,23 @@ import {
   type ConfigChange,
 } from "./issues";
 import { fetchReports, fetchAllReports, fetchReportFiles, fetchReportFileData, parseFlexibleDate } from "./reports";
+import {
+  resolveDblabInstanceId,
+  createClone,
+  listClones,
+  getClone,
+  resetClone,
+  destroyClone,
+  listBranches,
+  createBranch,
+  deleteBranch,
+  branchLog,
+  listSnapshots,
+  createSnapshot,
+  destroySnapshot,
+} from "./dblab";
 import { uploadFile, downloadFile, buildMarkdownLink, uploadAttachments, appendAttachmentsToContent } from "./storage";
+import { executeJoeCommand, listProjects, type JoeCommand } from "./joe";
 import { resolveBaseUrls } from "./util";
 
 // MCP SDK imports - Bun handles these directly
@@ -50,6 +66,283 @@ export interface McpToolResponse {
   isError?: boolean;
 }
 
+/** MCP tool name → Joe command (SPEC §5 names). */
+export const JOE_TOOL_TO_COMMAND: Record<string, JoeCommand> = {
+  plan_query: "plan",
+  explain_query: "explain",
+  exec_sql: "exec",
+  hypo_index: "hypo",
+  activity: "activity",
+  terminate_backend: "terminate",
+  reset_clone: "reset",
+  describe: "describe",
+};
+
+interface McpToolDefinition {
+  name: string;
+  description: string;
+  inputSchema: Record<string, unknown>;
+}
+
+/**
+ * MCP tool definitions for the Joe command surface + org-level project discovery.
+ * Names match SPEC §5. Descriptions disclose the egress (finding B2): plan/explain
+ * feed the plan/stats to your LLM when AI features are enabled; exec_sql carries the
+ * stronger raw-data warning (real own-data table rows can enter the agent's LLM
+ * context). Every command tool carries an optional session_id (finding B-1).
+ */
+export function joeToolDefinitions(): McpToolDefinition[] {
+  const sessionProps = {
+    project_id: {
+      type: "string",
+      description: "Target project by numeric id OR alias/name (id-or-alias). --project 12 ≡ main-db.",
+    },
+    session_id: {
+      type: "string",
+      description: "Run on a specific session's clone (64-bit id as a string). Auto-reused per project when omitted.",
+    },
+    new_session: { type: "boolean", description: "Force a fresh session/clone (null session)." },
+    timeout_ms: { type: "number", description: "One-shot poll budget in ms (≤ 25000); on expiry returns a command_id to resume." },
+    debug: { type: "boolean", description: "Enable verbose debug logs." },
+  } as const;
+
+  return [
+    {
+      name: "plan_query",
+      description:
+        "Plan a query (EXPLAIN, plan-only — NEVER executes; the fast/safe default). Returns a sanitized plan tree (literals redacted). LLM-facing: feeds the plan/stats to your LLM when the org's AI features are enabled. Requires joe:plan.",
+      inputSchema: {
+        type: "object",
+        properties: { sql: { type: "string", description: "The query to plan (one explainable statement)." }, ...sessionProps },
+        required: ["sql", "project_id"],
+        additionalProperties: false,
+      },
+    },
+    {
+      name: "explain_query",
+      description:
+        "EXPLAIN ANALYZE a query — EXECUTES it on the hardened, ephemeral DBLab clone and returns the measured, annotated plan (node types, costs, actual-rows/timings, buffers — literals redacted), NOT raw SELECT result rows. LLM-facing: feeds the measured plan/stats to your LLM when AI features are enabled. Requires joe:exec; subject to the org execution policy (read_only allows SELECT-only).",
+      inputSchema: {
+        type: "object",
+        properties: { sql: { type: "string", description: "The query to EXPLAIN ANALYZE (executes on the clone)." }, ...sessionProps },
+        required: ["sql", "project_id"],
+        additionalProperties: false,
+      },
+    },
+    {
+      name: "exec_sql",
+      description:
+        "WARNING: executes arbitrary SQL (DDL/DML) on the disposable clone and CAN RETURN REAL (own-data) TABLE ROWS that will enter this agent's LLM context (subject to the agent row-cap / metadata-only minimization). Prefer plan_query (non-executing) when you only need a plan. Requires an explicit joe:exec grant and is subject to the org's joe_execution_policy.",
+      inputSchema: {
+        type: "object",
+        properties: { sql: { type: "string", description: "Arbitrary DDL/DML to run on the clone (e.g. create index, analyze)." }, ...sessionProps },
+        required: ["sql", "project_id"],
+        additionalProperties: false,
+      },
+    },
+    {
+      name: "hypo_index",
+      description:
+        "HypoPG hypothetical index + re-plan (no real index built, no data change) — would the planner switch to it? LLM-facing. Requires joe:plan.",
+      inputSchema: {
+        type: "object",
+        properties: {
+          sql: { type: "string", description: "Candidate index DDL (e.g. create index on orders (customer_id))." },
+          query: { type: "string", description: "Target query the hypothetical index is evaluated against." },
+          ...sessionProps,
+        },
+        required: ["sql", "query", "project_id"],
+        additionalProperties: false,
+      },
+    },
+    {
+      name: "activity",
+      description: "pg_stat_activity snapshot on the clone (query text redacted). LLM-facing. Requires joe:plan.",
+      inputSchema: { type: "object", properties: { ...sessionProps }, required: ["project_id"], additionalProperties: false },
+    },
+    {
+      name: "terminate_backend",
+      description: "pg_terminate_backend(pid) on the clone. Status-only (no data to the LLM). Requires joe:admin.",
+      inputSchema: {
+        type: "object",
+        properties: { pid: { type: "number", description: "Backend pid to terminate." }, ...sessionProps },
+        required: ["pid", "project_id"],
+        additionalProperties: false,
+      },
+    },
+    {
+      name: "reset_clone",
+      description: "Reset/recreate the session's thin clone. Status-only (no data to the LLM). Requires joe:admin.",
+      inputSchema: { type: "object", properties: { ...sessionProps }, required: ["project_id"], additionalProperties: false },
+    },
+    {
+      name: "describe",
+      description: "\\d-family schema/relation/index/db/view metadata introspection. LLM-facing (metadata). Requires joe:plan.",
+      inputSchema: {
+        type: "object",
+        properties: {
+          object: { type: "string", description: "Object to describe (e.g. users)." },
+          variant: { type: "string", description: "\\d-family variant (e.g. \\d+, \\di, \\dt)." },
+          ...sessionProps,
+        },
+        required: ["object", "project_id"],
+        additionalProperties: false,
+      },
+    },
+    {
+      name: "list_projects",
+      description:
+        "List the org's projects and which have Joe ready (org-level discovery — NOT a Joe endpoint; mirrors `pgai projects`). Use it to find the project_ids/aliases you may address with the Joe command tools.",
+      inputSchema: {
+        type: "object",
+        properties: {
+          org_id: { type: "number", description: "Organization ID (optional, falls back to config)." },
+          debug: { type: "boolean", description: "Enable verbose debug logs." },
+        },
+        additionalProperties: false,
+      },
+    },
+  ];
+}
+
+interface RunJoeToolParams {
+  toolName: string;
+  apiKey: string;
+  apiBaseUrl: string;
+  orgId?: number;
+  args: Record<string, unknown>;
+  debug: boolean;
+}
+
+/**
+ * Coerce the MCP `timeout_ms` argument to a usable poll budget, mirroring the
+ * CLI's `--budget` guard. `handleToolCall` performs no schema validation, so a
+ * non-numeric value would become NaN — which survives `?? DEFAULT_BUDGET_MS`
+ * downstream and turns the poll deadline into NaN (an UNBOUNDED loop). Anything
+ * non-finite (absent, null, NaN, Infinity, a non-numeric string) → undefined,
+ * so `runCommand` falls back to its default budget. Exported for tests.
+ */
+export function sanitizeBudgetMs(value: unknown): number | undefined {
+  if (value === undefined || value === null) {
+    return undefined;
+  }
+  const n = Number(value);
+  return Number.isFinite(n) ? n : undefined;
+}
+
+/**
+ * Execute a Joe command MCP tool: map the MCP args 1:1 to the rpc params, resolve
+ * `project_id` (id-or-alias), thread/persist the session, run the submit-then-poll
+ * one-shot, and return the outcome as JSON text. Mirrors the CLI verbs exactly.
+ */
+async function runJoeTool(params: RunJoeToolParams): Promise<McpToolResponse> {
+  const command = JOE_TOOL_TO_COMMAND[params.toolName];
+  const a = params.args;
+
+  const project = String(a.project_id ?? "").trim();
+  if (!project) {
+    return { content: [{ type: "text", text: "project_id is required" }], isError: true };
+  }
+
+  let sql: string | null = null;
+  let cmdArgs: Record<string, unknown> | null = null;
+
+  if (command === "plan" || command === "explain" || command === "exec" || command === "hypo") {
+    sql = String(a.sql ?? "").trim();
+    if (!sql) {
+      return { content: [{ type: "text", text: "sql is required" }], isError: true };
+    }
+  }
+  if (command === "hypo") {
+    const query = String(a.query ?? "").trim();
+    if (!query) {
+      return { content: [{ type: "text", text: "query is required for hypo_index" }], isError: true };
+    }
+    cmdArgs = { query };
+  }
+  if (command === "terminate") {
+    // A pid must be a strict positive integer — `Number.isFinite` alone lets
+    // -5 / 1.5 / 0 through to a pg_terminate_backend against the wrong (or no)
+    // backend. Mirrors the CLI's bare-digits (^[0-9]+$) guard.
+    const raw = a.pid;
+    const pid =
+      typeof raw === "number"
+        ? raw
+        : typeof raw === "string" && /^[0-9]+$/.test(raw.trim())
+          ? Number(raw.trim())
+          : Number.NaN;
+    if (!Number.isInteger(pid) || pid <= 0) {
+      return {
+        content: [{ type: "text", text: "pid must be a strict positive integer (backend pid) for terminate_backend" }],
+        isError: true,
+      };
+    }
+    cmdArgs = { pid };
+  }
+  if (command === "describe") {
+    const object = String(a.object ?? "").trim();
+    if (!object) {
+      return { content: [{ type: "text", text: "object is required for describe" }], isError: true };
+    }
+    cmdArgs = { object };
+    if (a.variant !== undefined) cmdArgs.variant = String(a.variant);
+  }
+
+  const session = a.session_id !== undefined && a.session_id !== null ? String(a.session_id) : null;
+  const newSession = a.new_session === true;
+  const budgetMs = sanitizeBudgetMs(a.timeout_ms);
+
+  const outcome = await executeJoeCommand({
+    apiKey: params.apiKey,
+    apiBaseUrl: params.apiBaseUrl,
+    command,
+    project,
+    sql,
+    args: cmdArgs,
+    session,
+    newSession,
+    orgId: params.orgId,
+    budgetMs,
+    debug: params.debug,
+  });
+
+  const payload = {
+    command,
+    command_id: outcome.commandId,
+    session_id: outcome.sessionId,
+    status: outcome.status,
+    budget_expired: outcome.budgetExpired,
+    resume: outcome.budgetExpired ? `pgai joe result ${outcome.commandId}` : undefined,
+    result: outcome.result,
+  };
+  const isError = outcome.status === "error" || outcome.status === "timed_out";
+  return { content: [{ type: "text", text: JSON.stringify(payload, null, 2) }], isError };
+}
+
+/**
+ * Resolve a DBLab-companion tool's `project` arg to its single DBLab instance id
+ * (SPEC §8). Mirrors the CLI's `resolveDblabTarget`. Throws on a missing/unknown
+ * project — the caller's try/catch turns that into an `isError` tool response.
+ */
+async function resolveDblabInstanceForTool(
+  apiKey: string,
+  apiBaseUrl: string,
+  args: Record<string, unknown>,
+  cfgOrgId: number | null,
+  debug: boolean
+): Promise<string> {
+  const project = String(args.project || "").trim();
+  if (!project) {
+    throw new Error("project is required (project id or alias)");
+  }
+  const orgId = args.org_id !== undefined ? Number(args.org_id) : cfgOrgId ?? undefined;
+  return resolveDblabInstanceId({ apiKey, apiBaseUrl, project, orgId, debug });
+}
+
+function optStr(v: unknown): string | undefined {
+  return v !== undefined && v !== null && String(v).length > 0 ? String(v) : undefined;
+}
+
 /** Handle MCP tool calls - exported for testing */
 export async function handleToolCall(
   req: McpToolRequest,
@@ -369,6 +662,144 @@ export async function handleToolCall(
       return { content: [{ type: "text", text: JSON.stringify(files, null, 2) }] };
     }
 
+    // Org-level project discovery (NOT a Joe endpoint) — mirrors `pgai projects`.
+    if (toolName === "list_projects") {
+      const orgId = args.org_id !== undefined ? Number(args.org_id) : cfg.orgId ?? undefined;
+      const projects = await listProjects({ apiKey, apiBaseUrl, orgId, debug });
+      return { content: [{ type: "text", text: JSON.stringify(projects, null, 2) }] };
+    }
+
+    // Joe command tools (SPEC §5) — each a thin wrapper over the same joe.ts pipeline.
+    if (toolName in JOE_TOOL_TO_COMMAND) {
+      const orgId = cfg.orgId ?? undefined;
+      return await runJoeTool({ toolName, apiKey, apiBaseUrl, orgId, args, debug });
+    }
+
+    // ---- DBLab companion tools (SPEC §8) --------------------------------
+    // Each mirrors the CLI `pgai dblab clone|branch|snapshot …` verb 1:1: resolve the
+    // project's single DBLab instance, then proxy the existing Platform DBLab
+    // API (v1.dblab_api_call). Destructive verbs are `joe:admin`-gated backend
+    // side (a PT403 surfaces as an isError response via the outer catch).
+    if (toolName === "clone_create") {
+      const instanceId = await resolveDblabInstanceForTool(apiKey, apiBaseUrl, args, cfg.orgId, debug);
+      const result = await createClone({
+        apiKey, apiBaseUrl, instanceId,
+        cloneId: optStr(args.clone_id),
+        branch: optStr(args.branch),
+        snapshotId: optStr(args.snapshot_id),
+        dbUser: optStr(args.db_user),
+        dbPassword: optStr(args.db_password),
+        isProtected: args.protected === true,
+        debug,
+      });
+      return { content: [{ type: "text", text: JSON.stringify(result, null, 2) }] };
+    }
+
+    if (toolName === "clone_list") {
+      const instanceId = await resolveDblabInstanceForTool(apiKey, apiBaseUrl, args, cfg.orgId, debug);
+      const result = await listClones({ apiKey, apiBaseUrl, instanceId, debug });
+      return { content: [{ type: "text", text: JSON.stringify(result, null, 2) }] };
+    }
+
+    if (toolName === "clone_status") {
+      const cloneId = String(args.clone_id || "").trim();
+      if (!cloneId) return { content: [{ type: "text", text: "clone_id is required" }], isError: true };
+      const instanceId = await resolveDblabInstanceForTool(apiKey, apiBaseUrl, args, cfg.orgId, debug);
+      const result = await getClone({ apiKey, apiBaseUrl, instanceId, cloneId, debug });
+      return { content: [{ type: "text", text: JSON.stringify(result, null, 2) }] };
+    }
+
+    if (toolName === "clone_reset") {
+      const cloneId = String(args.clone_id || "").trim();
+      if (!cloneId) return { content: [{ type: "text", text: "clone_id is required" }], isError: true };
+      const instanceId = await resolveDblabInstanceForTool(apiKey, apiBaseUrl, args, cfg.orgId, debug);
+      const result = await resetClone({
+        apiKey, apiBaseUrl, instanceId, cloneId,
+        snapshotId: optStr(args.snapshot_id),
+        latest: args.latest === true ? true : undefined,
+        debug,
+      });
+      return { content: [{ type: "text", text: JSON.stringify(result ?? { reset: true, cloneId }, null, 2) }] };
+    }
+
+    if (toolName === "clone_destroy") {
+      const cloneId = String(args.clone_id || "").trim();
+      if (!cloneId) return { content: [{ type: "text", text: "clone_id is required" }], isError: true };
+      const instanceId = await resolveDblabInstanceForTool(apiKey, apiBaseUrl, args, cfg.orgId, debug);
+      const result = await destroyClone({ apiKey, apiBaseUrl, instanceId, cloneId, debug });
+      return { content: [{ type: "text", text: JSON.stringify(result ?? { destroyed: true, cloneId }, null, 2) }] };
+    }
+
+    if (toolName === "branch_list") {
+      const instanceId = await resolveDblabInstanceForTool(apiKey, apiBaseUrl, args, cfg.orgId, debug);
+      const result = await listBranches({ apiKey, apiBaseUrl, instanceId, debug });
+      return { content: [{ type: "text", text: JSON.stringify(result, null, 2) }] };
+    }
+
+    if (toolName === "branch_create") {
+      const branchName = String(args.name || "").trim();
+      if (!branchName) return { content: [{ type: "text", text: "name is required" }], isError: true };
+      const instanceId = await resolveDblabInstanceForTool(apiKey, apiBaseUrl, args, cfg.orgId, debug);
+      const result = await createBranch({
+        apiKey, apiBaseUrl, instanceId, branchName,
+        baseBranch: optStr(args.base_branch),
+        snapshotId: optStr(args.snapshot_id),
+        debug,
+      });
+      return { content: [{ type: "text", text: JSON.stringify(result, null, 2) }] };
+    }
+
+    if (toolName === "branch_delete") {
+      const branchName = String(args.name || "").trim();
+      if (!branchName) return { content: [{ type: "text", text: "name is required" }], isError: true };
+      const instanceId = await resolveDblabInstanceForTool(apiKey, apiBaseUrl, args, cfg.orgId, debug);
+      const result = await deleteBranch({ apiKey, apiBaseUrl, instanceId, branchName, debug });
+      return { content: [{ type: "text", text: JSON.stringify(result ?? { deleted: true, branch: branchName }, null, 2) }] };
+    }
+
+    if (toolName === "branch_log") {
+      const branchName = String(args.name || "").trim();
+      if (!branchName) return { content: [{ type: "text", text: "name is required" }], isError: true };
+      const instanceId = await resolveDblabInstanceForTool(apiKey, apiBaseUrl, args, cfg.orgId, debug);
+      const result = await branchLog({ apiKey, apiBaseUrl, instanceId, branchName, debug });
+      return { content: [{ type: "text", text: JSON.stringify(result, null, 2) }] };
+    }
+
+    if (toolName === "snapshot_list") {
+      const instanceId = await resolveDblabInstanceForTool(apiKey, apiBaseUrl, args, cfg.orgId, debug);
+      const result = await listSnapshots({
+        apiKey, apiBaseUrl, instanceId,
+        branchName: optStr(args.branch),
+        dataset: optStr(args.dataset),
+        debug,
+      });
+      return { content: [{ type: "text", text: JSON.stringify(result, null, 2) }] };
+    }
+
+    if (toolName === "snapshot_create") {
+      const cloneId = String(args.clone_id || "").trim();
+      if (!cloneId) return { content: [{ type: "text", text: "clone_id is required" }], isError: true };
+      const instanceId = await resolveDblabInstanceForTool(apiKey, apiBaseUrl, args, cfg.orgId, debug);
+      const result = await createSnapshot({
+        apiKey, apiBaseUrl, instanceId, cloneId,
+        message: optStr(args.message),
+        debug,
+      });
+      return { content: [{ type: "text", text: JSON.stringify(result, null, 2) }] };
+    }
+
+    if (toolName === "snapshot_destroy") {
+      const snapshotId = String(args.snapshot_id || "").trim();
+      if (!snapshotId) return { content: [{ type: "text", text: "snapshot_id is required" }], isError: true };
+      const instanceId = await resolveDblabInstanceForTool(apiKey, apiBaseUrl, args, cfg.orgId, debug);
+      const result = await destroySnapshot({
+        apiKey, apiBaseUrl, instanceId, snapshotId,
+        force: args.force === true,
+        debug,
+      });
+      return { content: [{ type: "text", text: JSON.stringify(result ?? { destroyed: true, snapshot: snapshotId }, null, 2) }] };
+    }
+
     throw new Error(`Unknown tool: ${toolName}`);
   } catch (err) {
     const message = err instanceof Error ? err.message : String(err);
@@ -653,6 +1084,200 @@ export async function startMcpServer(rootOpts?: RootOptsLike, extra?: { debug?:
             additionalProperties: false,
           },
         },
+        ...joeToolDefinitions(),
+        // DBLab companion tools (SPEC §8) — proxy the existing Platform DBLab API
+        // (v1.dblab_api_call) the Console drives; mirror `pgai dblab clone|branch|snapshot`.
+        // `project` (id or alias) resolves the project's single DBLab instance.
+        {
+          name: "clone_create",
+          description: "Create a DBLab thin clone of the project's database (same as CLI 'pgai dblab clone create'). Requires the joe:plan scope.",
+          inputSchema: {
+            type: "object",
+            properties: {
+              project: { type: "string", description: "Project id or alias" },
+              branch: { type: "string", description: "Branch to clone from" },
+              snapshot_id: { type: "string", description: "Snapshot id to clone from" },
+              clone_id: { type: "string", description: "Clone id (DBLab generates one when omitted)" },
+              db_user: { type: "string", description: "Clone DB user" },
+              db_password: { type: "string", description: "Clone DB password" },
+              protected: { type: "boolean", description: "Protect the clone from auto-deletion" },
+              org_id: { type: "number", description: "Organization id (optional, falls back to config)" },
+              debug: { type: "boolean", description: "Enable verbose debug logs" },
+            },
+            required: ["project"],
+            additionalProperties: false,
+          },
+        },
+        {
+          name: "clone_list",
+          description: "List the project's DBLab thin clones (same as CLI 'pgai dblab clone list').",
+          inputSchema: {
+            type: "object",
+            properties: {
+              project: { type: "string", description: "Project id or alias" },
+              org_id: { type: "number", description: "Organization id (optional, falls back to config)" },
+              debug: { type: "boolean", description: "Enable verbose debug logs" },
+            },
+            required: ["project"],
+            additionalProperties: false,
+          },
+        },
+        {
+          name: "clone_status",
+          description: "Show a DBLab clone's status (same as CLI 'pgai dblab clone status').",
+          inputSchema: {
+            type: "object",
+            properties: {
+              project: { type: "string", description: "Project id or alias" },
+              clone_id: { type: "string", description: "Clone id" },
+              org_id: { type: "number", description: "Organization id (optional, falls back to config)" },
+              debug: { type: "boolean", description: "Enable verbose debug logs" },
+            },
+            required: ["project", "clone_id"],
+            additionalProperties: false,
+          },
+        },
+        {
+          name: "clone_reset",
+          description: "Reset a DBLab clone to a pristine snapshot (same as CLI 'pgai dblab clone reset'). Requires the joe:admin scope.",
+          inputSchema: {
+            type: "object",
+            properties: {
+              project: { type: "string", description: "Project id or alias" },
+              clone_id: { type: "string", description: "Clone id" },
+              snapshot_id: { type: "string", description: "Snapshot id to reset to (default: latest)" },
+              latest: { type: "boolean", description: "Reset to the latest snapshot" },
+              org_id: { type: "number", description: "Organization id (optional, falls back to config)" },
+              debug: { type: "boolean", description: "Enable verbose debug logs" },
+            },
+            required: ["project", "clone_id"],
+            additionalProperties: false,
+          },
+        },
+        {
+          name: "clone_destroy",
+          description: "Destroy a DBLab clone (same as CLI 'pgai dblab clone destroy'). Requires the joe:admin scope.",
+          inputSchema: {
+            type: "object",
+            properties: {
+              project: { type: "string", description: "Project id or alias" },
+              clone_id: { type: "string", description: "Clone id" },
+              org_id: { type: "number", description: "Organization id (optional, falls back to config)" },
+              debug: { type: "boolean", description: "Enable verbose debug logs" },
+            },
+            required: ["project", "clone_id"],
+            additionalProperties: false,
+          },
+        },
+        {
+          name: "branch_list",
+          description: "List the project's DBLab branches (same as CLI 'pgai dblab branch list').",
+          inputSchema: {
+            type: "object",
+            properties: {
+              project: { type: "string", description: "Project id or alias" },
+              org_id: { type: "number", description: "Organization id (optional, falls back to config)" },
+              debug: { type: "boolean", description: "Enable verbose debug logs" },
+            },
+            required: ["project"],
+            additionalProperties: false,
+          },
+        },
+        {
+          name: "branch_create",
+          description: "Create a DBLab branch (same as CLI 'pgai dblab branch create'). Requires the joe:plan scope.",
+          inputSchema: {
+            type: "object",
+            properties: {
+              project: { type: "string", description: "Project id or alias" },
+              name: { type: "string", description: "Branch name" },
+              snapshot_id: { type: "string", description: "Snapshot id to base the branch on" },
+              base_branch: { type: "string", description: "Parent branch to fork from" },
+              org_id: { type: "number", description: "Organization id (optional, falls back to config)" },
+              debug: { type: "boolean", description: "Enable verbose debug logs" },
+            },
+            required: ["project", "name"],
+            additionalProperties: false,
+          },
+        },
+        {
+          name: "branch_delete",
+          description: "Delete a DBLab branch (same as CLI 'pgai dblab branch delete'). Requires the joe:admin scope.",
+          inputSchema: {
+            type: "object",
+            properties: {
+              project: { type: "string", description: "Project id or alias" },
+              name: { type: "string", description: "Branch name" },
+              org_id: { type: "number", description: "Organization id (optional, falls back to config)" },
+              debug: { type: "boolean", description: "Enable verbose debug logs" },
+            },
+            required: ["project", "name"],
+            additionalProperties: false,
+          },
+        },
+        {
+          name: "branch_log",
+          description: "Show a DBLab branch's snapshot log (same as CLI 'pgai dblab branch log').",
+          inputSchema: {
+            type: "object",
+            properties: {
+              project: { type: "string", description: "Project id or alias" },
+              name: { type: "string", description: "Branch name" },
+              org_id: { type: "number", description: "Organization id (optional, falls back to config)" },
+              debug: { type: "boolean", description: "Enable verbose debug logs" },
+            },
+            required: ["project", "name"],
+            additionalProperties: false,
+          },
+        },
+        {
+          name: "snapshot_list",
+          description: "List the project's DBLab snapshots (same as CLI 'pgai dblab snapshot list').",
+          inputSchema: {
+            type: "object",
+            properties: {
+              project: { type: "string", description: "Project id or alias" },
+              branch: { type: "string", description: "Filter by branch" },
+              dataset: { type: "string", description: "Filter by dataset" },
+              org_id: { type: "number", description: "Organization id (optional, falls back to config)" },
+              debug: { type: "boolean", description: "Enable verbose debug logs" },
+            },
+            required: ["project"],
+            additionalProperties: false,
+          },
+        },
+        {
+          name: "snapshot_create",
+          description: "Create a DBLab snapshot from a clone (same as CLI 'pgai dblab snapshot create'). Requires the joe:plan scope.",
+          inputSchema: {
+            type: "object",
+            properties: {
+              project: { type: "string", description: "Project id or alias" },
+              clone_id: { type: "string", description: "Clone id to snapshot" },
+              message: { type: "string", description: "Snapshot message" },
+              org_id: { type: "number", description: "Organization id (optional, falls back to config)" },
+              debug: { type: "boolean", description: "Enable verbose debug logs" },
+            },
+            required: ["project", "clone_id"],
+            additionalProperties: false,
+          },
+        },
+        {
+          name: "snapshot_destroy",
+          description: "Destroy a DBLab snapshot (same as CLI 'pgai dblab snapshot destroy'). Requires the joe:admin scope.",
+          inputSchema: {
+            type: "object",
+            properties: {
+              project: { type: "string", description: "Project id or alias" },
+              snapshot_id: { type: "string", description: "Snapshot id" },
+              force: { type: "boolean", description: "Force-delete even when dependent clones exist" },
+              org_id: { type: "number", description: "Organization id (optional, falls back to config)" },
+              debug: { type: "boolean", description: "Enable verbose debug logs" },
+            },
+            required: ["project", "snapshot_id"],
+            additionalProperties: false,
+          },
+        },
       ],
     };
   });

package/lib/supabase.ts

@@ -733,24 +733,6 @@ export async function verifyInitSetupViaSupabase(params: {
   }
 
   // Check helper functions - first verify they exist to avoid has_function_privilege errors
-  const explainFnExistsRes = await params.client.query(
-    "SELECT oid FROM pg_proc WHERE proname = 'explain_generic' AND pronamespace = (SELECT oid FROM pg_namespace WHERE nspname = 'postgres_ai')",
-    true
-  );
-  if (explainFnExistsRes.rowCount === 0) {
-    missingRequired.push("function postgres_ai.explain_generic exists");
-  } else {
-    const explainFnRes = await params.client.query(
-      `SELECT has_function_privilege('${escapeLiteral(role)}', 'postgres_ai.explain_generic(text, text, text)', 'EXECUTE') as ok`,
-      true
-    );
-    if (!explainFnRes.rows?.[0]?.ok) {
-      missingRequired.push(
-        "EXECUTE on postgres_ai.explain_generic(text, text, text)"
-      );
-    }
-  }
-
   const tableDescribeFnExistsRes = await params.client.query(
     "SELECT oid FROM pg_proc WHERE proname = 'table_describe' AND pronamespace = (SELECT oid FROM pg_namespace WHERE nspname = 'postgres_ai')",
     true