@askdb/studio 0.2.0-beta.3 → 0.2.0-beta.30

package/dist/server.js

@@ -1,26 +1,41 @@
-import { existsSync, readFileSync, rmSync } from "node:fs";
+import { randomUUID } from "node:crypto";
+import { existsSync, readFileSync, rmSync, writeFileSync } from "node:fs";
 import { createServer } from "node:http";
-import { extname, relative, join, resolve } from "node:path";
+import { extname, relative, join, resolve, dirname, sep } from "node:path";
+import { spawn } from "node:child_process";
 import { fileURLToPath } from "node:url";
 import { generateText as defaultGenerateText } from "ai";
-import { getAskDbRuntimeConfig } from "@askdb/config";
-import { ask, askDbAiKeyMissingMessage, createAskDbEmbeddingModel, createAskDbLanguageModelFromEnv, isBuiltInDialectId, loadSchema, resolveAskDbAiConfig, resolveAskDbEmbeddingConfig, suggestEnrichment, } from "@askdb/core";
-import { buildSchemaIndex, chunkContentHash, chunkSchema, createRetriever, loadChunkerSourcesFromDir, } from "@askdb/rag";
-import { createAiSdkEmbedder } from "@askdb/rag/embedders/ai-sdk";
-import { createFileStore } from "@askdb/rag/stores/file";
-import { buildDefaultTableBody, buildFrontmatter, buildSuggestionContext, buildSuggestionTarget, buildTableDraft, loadWorkspace, replaceH2Section, replaceTableDescription, saveTable, } from "@askdb/enrich";
+import { bootstrapAskDbEnv, getAskDbRuntimeConfig } from "@askdb/config";
+import { createAiRegistry, } from "@askdb/ai";
+import { anthropicProvider } from "@askdb/ai-anthropic";
+import { azureProvider } from "@askdb/ai-azure";
+import { googleProvider } from "@askdb/ai-google";
+import { openaiProvider } from "@askdb/ai-openai";
+import { ask, isBuiltInDialectId, loadSchema, suggestEnrichment, tenantScopeSchema, writeTenantPolicyMarkdown, tenantPolicyFrontmatterSchema, } from "@askdb/core";
+import { buildSchemaIndex, chunkContentHash, chunkSchema, createAiSdkEmbedder, createFileStore, createMemoryStore, createPgvectorStore, createRetriever, loadChunkerSourcesFromDir, } from "@askdb/rag";
+import { buildDefaultTableBody, buildFrontmatter, buildSuggestionContext, buildSuggestionTarget, buildTableDraft, loadWorkspace, replaceH2Section, replaceTableDescription, saveConcepts, saveTable, } from "@askdb/enrich";
+import { EXECUTE_DRIVER_REGISTRY, isDriverInstalled } from "./execute-registry.js";
+import { resolveStudioIntrospectionPlan, runStudioIntrospection } from "./introspection.js";
+import { probeSetupState, writeSetupConfig, SetupError } from "./setup.js";
+const ai = createAiRegistry([openaiProvider, azureProvider, googleProvider, anthropicProvider]);
 const CLIENT_DIR = fileURLToPath(new URL("./client/", import.meta.url));
 const STUDIO_RAG_MOCK_DIMENSIONS = 64;
 const STUDIO_RAG_MOCK_EMBEDDER_ID = `studio:mock-lexical-${STUDIO_RAG_MOCK_DIMENSIONS}`;
 const DEFAULT_EMBEDDING_MODEL = "text-embedding-3-small";
+let studioPgvectorStoreFactoryForTests;
+export function setStudioPgvectorStoreFactoryForTests(factory) {
+    studioPgvectorStoreFactoryForTests = factory;
+}
 export function createStudioServer(options) {
     const schemaDir = resolve(options.schema);
-    if (!existsSync(schemaDir)) {
+    const setupReason = options.setupReason ?? null;
+    if (!setupReason && !existsSync(schemaDir)) {
         throw new Error(`schema directory not found: ${schemaDir}`);
     }
     const state = {
         schemaDir,
-        workspace: loadWorkspace(schemaDir),
+        workspace: setupReason ? null : loadWorkspace(schemaDir),
+        setupReason,
     };
     return createServer(async (req, res) => {
         try {
@@ -31,39 +46,122 @@ export function createStudioServer(options) {
             if (req.method === "GET" && url.pathname.startsWith("/assets/")) {
                 return serveClientFile(res, decodeURIComponent(url.pathname.slice(1)));
             }
+            if (req.method === "GET" && url.pathname === "/api/setup/status") {
+                return writeJson(res, 200, buildSetupStatus(state));
+            }
+            if (req.method === "POST" && url.pathname === "/api/setup/config") {
+                if (!isLoopbackRequest(req)) {
+                    return writeJson(res, 403, { error: { message: "Setup is only available from loopback clients." } });
+                }
+                const body = await readJson(req);
+                const result = handleSetupConfig(state, body);
+                return writeJson(res, 200, result);
+            }
+            if (req.method === "POST" && url.pathname === "/api/setup/introspect") {
+                if (!isLoopbackRequest(req)) {
+                    return writeJson(res, 403, { error: { message: "Setup is only available from loopback clients." } });
+                }
+                const result = await handleSetupIntrospect(state);
+                return writeJson(res, 200, result);
+            }
+            // Everything below needs a loaded workspace — route setup-mode clients to the wizard.
+            if (state.setupReason && url.pathname.startsWith("/api/")) {
+                return writeJson(res, 409, {
+                    error: { message: "Studio needs setup. Open Studio in the browser to finish the guided setup." },
+                    setup: buildSetupStatus(state),
+                });
+            }
             if (req.method === "GET" && url.pathname === "/api/workspace") {
-                return writeJson(res, 200, serializeWorkspace(state.workspace));
+                return writeJson(res, 200, serializeWorkspace(requireWorkspace(state)));
+            }
+            if (req.method === "GET" && url.pathname === "/api/introspect/status") {
+                requireWorkspace(state);
+                return writeJson(res, 200, buildIntrospectionPlanDto());
+            }
+            if (req.method === "POST" && url.pathname === "/api/introspect") {
+                if (!isLoopbackRequest(req)) {
+                    return writeJson(res, 403, { error: { message: "Resync is only available from loopback clients." } });
+                }
+                const result = await handleResync(state);
+                return writeJson(res, 200, result);
             }
             if (req.method === "POST" && url.pathname.startsWith("/api/tables/")) {
                 const tableId = decodeURIComponent(url.pathname.slice("/api/tables/".length));
                 const body = await readJson(req);
                 const draft = parseTableDraftBody(body);
                 saveDraft(state, tableId, draft);
-                return writeJson(res, 200, serializeWorkspace(state.workspace));
+                return writeJson(res, 200, serializeWorkspace(requireWorkspace(state)));
+            }
+            if (req.method === "POST" && url.pathname === "/api/concepts") {
+                const body = await readJson(req);
+                const concepts = parseConceptsBody(body);
+                saveConceptsDraft(state, concepts);
+                return writeJson(res, 200, serializeWorkspace(requireWorkspace(state)));
+            }
+            if (req.method === "POST" && url.pathname === "/api/tenant-policy") {
+                const body = await readJson(req);
+                const { frontmatter, bodyText } = parseTenantPolicyBody(body);
+                saveTenantPolicy(state, frontmatter, bodyText);
+                return writeJson(res, 200, serializeWorkspace(requireWorkspace(state)));
+            }
+            if (req.method === "POST" && url.pathname === "/api/suggest-tenant-policy") {
+                const draft = await suggestTenantPolicyDraft(state);
+                return writeJson(res, 200, draft);
             }
             if (req.method === "POST" && url.pathname === "/api/suggest") {
                 const body = await readJson(req);
                 const source = parseSuggestSource(body);
-                const candidates = await suggestForSource(state.workspace, source);
+                const candidates = await suggestForSource(requireWorkspace(state), source);
                 return writeJson(res, 200, { candidates });
             }
             if (req.method === "GET" && url.pathname === "/api/rag/status") {
-                return writeJson(res, 200, getRagStatus(state.schemaDir));
+                return writeJson(res, 200, await getRagStatus(state));
             }
             if (req.method === "POST" && url.pathname === "/api/rag/index") {
-                const result = await indexRag(state.schemaDir);
+                const result = await indexRag(state);
                 return writeJson(res, 200, result);
             }
             if (req.method === "POST" && url.pathname === "/api/rag/query") {
                 const body = await readJson(req);
                 const query = parseRagQuery(body);
-                const result = await queryRag(state.schemaDir, query);
+                const result = await queryRag(state, query);
                 return writeJson(res, 200, result);
             }
             if (req.method === "POST" && url.pathname === "/api/ask") {
                 const body = await readJson(req);
                 const options = parseAskBody(body);
-                const result = await askSampleQuestion(state.schemaDir, options);
+                const result = await askSampleQuestion(state, options);
+                return writeJson(res, 200, result);
+            }
+            if (req.method === "GET" && url.pathname === "/api/history") {
+                return writeJson(res, 200, readPlaygroundHistory(state.schemaDir));
+            }
+            if (req.method === "POST" && url.pathname === "/api/history") {
+                const body = await readJson(req);
+                const entry = parsePlaygroundHistoryEntry(body);
+                appendPlaygroundHistory(state.schemaDir, entry);
+                return writeJson(res, 200, { ok: true });
+            }
+            if (req.method === "DELETE" && url.pathname.startsWith("/api/history/")) {
+                const id = decodeURIComponent(url.pathname.slice("/api/history/".length));
+                const result = deletePlaygroundHistoryEntry(state.schemaDir, id);
+                return writeJson(res, 200, result);
+            }
+            if (req.method === "GET" && url.pathname === "/api/execute/status") {
+                const result = await getExecuteStatus(state.schemaDir);
+                return writeJson(res, 200, result);
+            }
+            if (req.method === "POST" && url.pathname === "/api/execute/install-driver") {
+                if (!isLoopbackRequest(req)) {
+                    return writeJson(res, 403, { error: { message: "Driver install is only available from loopback clients." } });
+                }
+                const body = await readJson(req);
+                const result = await installExecuteDriver(body, state.schemaDir);
+                return writeJson(res, 200, result);
+            }
+            if (req.method === "POST" && url.pathname === "/api/execute") {
+                const body = await readJson(req);
+                const result = await executeQuery(body, state.schemaDir);
                 return writeJson(res, 200, result);
             }
             if (req.method === "GET" && !url.pathname.startsWith("/api/")) {
@@ -72,7 +170,7 @@ export function createStudioServer(options) {
             writeJson(res, 404, { error: { message: "Not found" } });
         }
         catch (error) {
-            const status = error instanceof StudioHttpError ? error.status : 500;
+            const status = error instanceof StudioHttpError || error instanceof SetupError ? error.status : 500;
             writeJson(res, status, {
                 error: { message: error instanceof Error ? error.message : String(error) },
             });
@@ -83,9 +181,7 @@ export function serializeWorkspace(workspace) {
     const rt = getAskDbRuntimeConfig();
     const aiConfig = (() => {
         try {
-            return resolveAskDbAiConfig(rt.ai.aiEnv, {
-                modelEnvVar: rt.ai.studioModel ? "ASKDB_STUDIO_MODEL" : undefined,
-            });
+            return ai.resolveAiConfig(rt.ai.aiEnv);
         }
         catch {
             // A misconfigured AI env (e.g. azure without resourceName) shouldn't crash the workspace
@@ -93,9 +189,15 @@ export function serializeWorkspace(workspace) {
             return undefined;
         }
     })();
+    const schema = loadSchema(workspace.schemaDir);
+    const schemaProvider = "provider" in schema && typeof schema.provider === "string" ? schema.provider : undefined;
+    const dialect = rt.nlToSql.dialect ??
+        (schemaProvider && isBuiltInDialectId(schemaProvider) ? schemaProvider : "postgres");
     return {
         schemaDir: workspace.schemaDir,
+        schemaPathRelative: toRelativeSchemaPath(workspace.schemaDir),
         schemaId: workspace.physical.schemaId,
+        dialect,
         warnings: workspace.warnings,
         aiConfigured: Boolean(aiConfig),
         model: aiConfig?.model ?? "gpt-4o-mini",
@@ -113,13 +215,192 @@ export function serializeWorkspace(workspace) {
             };
         }),
         concepts: workspace.concepts?.frontmatter.concepts ?? [],
+        tenantPolicy: schema.tenantPolicy ?? null,
+    };
+}
+function toRelativeSchemaPath(schemaDir) {
+    const rel = relative(process.cwd(), schemaDir);
+    if (rel === "")
+        return ".";
+    if (rel.startsWith("..") || rel.includes(`..${sep}`))
+        return schemaDir;
+    return rel.startsWith(".") ? rel : `./${rel}`;
+}
+// ---------------------------------------------------------------------------
+// Setup wizard + server-side introspection (resync)
+// ---------------------------------------------------------------------------
+function requireWorkspace(state) {
+    if (!state.workspace) {
+        throw new StudioHttpError(409, "Studio is in setup mode — finish the guided setup first.");
+    }
+    return state.workspace;
+}
+function buildSetupStatus(state) {
+    const projectDir = process.cwd();
+    if (!state.setupReason) {
+        return {
+            needed: false,
+            reason: null,
+            projectDir,
+            configPath: null,
+            outputDir: state.schemaDir,
+            configLoadError: null,
+        };
+    }
+    const probe = probeSetupState(projectDir);
+    return {
+        needed: true,
+        reason: probe.hasConfig ? "no-artifact" : "no-config",
+        projectDir,
+        configPath: probe.configPath,
+        outputDir: probe.outputDir,
+        configLoadError: probe.loadError,
+    };
+}
+function handleSetupConfig(state, body) {
+    if (!state.setupReason) {
+        throw new StudioHttpError(409, "Studio is already set up — edit askdb.config.ts directly.");
+    }
+    const input = parseSetupConfigBody(body);
+    const result = writeSetupConfig(process.cwd(), input);
+    // Config now exists; the artifact usually doesn't yet.
+    const probe = probeSetupState(process.cwd());
+    if (probe.hasArtifact && probe.outputDir) {
+        state.schemaDir = probe.outputDir;
+        state.workspace = loadWorkspace(probe.outputDir);
+        state.setupReason = null;
+    }
+    else {
+        state.setupReason = "no-artifact";
+    }
+    return {
+        configPath: result.configPath,
+        envExamplePath: result.envExamplePath,
+        envVars: result.envVars,
+        installed: result.installed,
+        manualInstallCommand: result.manualInstallCommand,
+        packageJsonCreated: result.packageJsonCreated,
+        status: buildSetupStatus(state),
+    };
+}
+async function handleSetupIntrospect(state) {
+    if (!state.setupReason) {
+        throw new StudioHttpError(409, "Studio is already set up — use Resync schema instead.");
+    }
+    // probeSetupState re-bootstraps, picking up any .env created since server start.
+    const probe = probeSetupState(process.cwd());
+    if (!probe.hasConfig) {
+        throw new StudioHttpError(409, "Write askdb.config.ts first (setup step 1).");
+    }
+    if (probe.loadError || !probe.outputDir) {
+        throw new StudioHttpError(409, `askdb.config.ts exists but could not be loaded: ${probe.loadError ?? "unknown error"}. ` +
+            "Install the project dependencies shown in setup step 1, then retry.");
+    }
+    const outDir = probe.outputDir;
+    const run = await runStudioIntrospection({
+        outDir,
+        hasExistingArtifact: existsSync(join(outDir, "schema.json")),
+    });
+    state.schemaDir = outDir;
+    state.workspace = loadWorkspace(outDir);
+    state.setupReason = null;
+    return {
+        ok: true,
+        engine: run.engine,
+        schemaId: run.schemaId,
+        tables: run.tables,
+        warnings: run.warnings,
+        workspace: serializeWorkspace(state.workspace),
+    };
+}
+async function handleResync(state) {
+    const workspace = requireWorkspace(state);
+    // Refresh the runtime snapshot so config edits made while Studio runs are honored.
+    try {
+        bootstrapAskDbEnv({ cwd: process.cwd() });
+    }
+    catch {
+        // Keep the snapshot the server started with when the config disappeared mid-session.
+    }
+    const run = await runStudioIntrospection({
+        outDir: state.schemaDir,
+        schemaId: workspace.physical.schemaId,
+        hasExistingArtifact: true,
+    });
+    state.workspace = loadWorkspace(state.schemaDir);
+    return {
+        ok: true,
+        engine: run.engine,
+        schemaId: run.schemaId,
+        tables: run.tables,
+        warnings: run.warnings,
+        workspace: serializeWorkspace(state.workspace),
     };
 }
+function buildIntrospectionPlanDto() {
+    const plan = resolveStudioIntrospectionPlan();
+    return plan.ok
+        ? { ok: true, engine: plan.engine, sourceLabel: plan.sourceLabel }
+        : { ok: false, engine: plan.engine, error: plan.error };
+}
+function parseSetupConfigBody(body) {
+    if (!isRecord(body)) {
+        throw new StudioHttpError(400, "Request body must be a JSON object.");
+    }
+    const databases = ["postgres", "mysql", "sqlite", "sqlserver", "prisma"];
+    const aiProviders = ["openai", "anthropic", "google", "azure"];
+    if (typeof body.database !== "string" || !databases.includes(body.database)) {
+        throw new StudioHttpError(400, `\`database\` must be one of: ${databases.join(", ")}.`);
+    }
+    if (typeof body.aiProvider !== "string" ||
+        !aiProviders.includes(body.aiProvider)) {
+        throw new StudioHttpError(400, `\`aiProvider\` must be one of: ${aiProviders.join(", ")}.`);
+    }
+    const optionalString = (key) => {
+        const value = body[key];
+        if (value === undefined || value === null || value === "")
+            return undefined;
+        if (typeof value !== "string")
+            throw new StudioHttpError(400, `\`${key}\` must be a string.`);
+        return value;
+    };
+    return {
+        database: body.database,
+        aiProvider: body.aiProvider,
+        connectionEnv: optionalString("connectionEnv"),
+        sqliteFile: optionalString("sqliteFile"),
+        prismaSchema: optionalString("prismaSchema"),
+        aiKeyEnv: optionalString("aiKeyEnv"),
+        schemaOut: optionalString("schemaOut"),
+    };
+}
+function saveConceptsDraft(state, concepts) {
+    saveConcepts(requireWorkspace(state), { concepts });
+    state.workspace = loadWorkspace(state.schemaDir);
+}
+function saveTenantPolicy(state, frontmatter, body) {
+    const filePath = join(state.schemaDir, "tenant-policy.md");
+    const md = writeTenantPolicyMarkdown(frontmatter, body);
+    writeFileSync(filePath, md, "utf8");
+    state.workspace = loadWorkspace(state.schemaDir);
+}
+function parseTenantPolicyBody(body) {
+    if (!isRecord(body) || !isRecord(body.frontmatter)) {
+        throw new StudioHttpError(400, "`frontmatter` object is required.");
+    }
+    const parsed = tenantPolicyFrontmatterSchema.safeParse(body.frontmatter);
+    if (!parsed.success) {
+        throw new StudioHttpError(400, `Invalid tenant policy: ${parsed.error.message}`);
+    }
+    const bodyText = typeof body.body === "string" ? body.body : "";
+    return { frontmatter: parsed.data, bodyText };
+}
 function saveDraft(state, tableId, draft) {
-    const table = state.workspace.tables.find((candidate) => candidate.physical.id === tableId);
+    const workspace = requireWorkspace(state);
+    const table = workspace.tables.find((candidate) => candidate.physical.id === tableId);
     if (!table)
         throw new StudioHttpError(404, `No such table: ${tableId}`);
-    const frontmatter = buildFrontmatter(table.physical, state.workspace.physical.schemaId, draft);
+    const frontmatter = buildFrontmatter(table.physical, workspace.physical.schemaId, draft);
     let body = table.parsed
         ? replaceTableDescription(table.parsed.body, draft.description)
         : buildDefaultTableBody(table.physical.name, draft.description);
@@ -129,43 +410,123 @@ function saveDraft(state, tableId, draft) {
     if (draft.exampleQuestions !== undefined) {
         body = replaceH2Section(body, "Example questions", draft.exampleQuestions);
     }
-    saveTable(state.workspace, tableId, frontmatter, body);
+    saveTable(workspace, tableId, frontmatter, body);
     state.workspace = loadWorkspace(state.schemaDir);
 }
 async function suggestForSource(workspace, source) {
     const rt = getAskDbRuntimeConfig();
-    const modelEnvVar = rt.ai.studioModel
-        ? "ASKDB_STUDIO_MODEL"
-        : rt.ai.tuiModel
-            ? "ASKDB_TUI_MODEL"
-            : undefined;
-    const model = await createAskDbLanguageModelFromEnv(rt.ai.aiEnv, {
-        modelEnvVar,
-    });
+    const model = await ai.createLanguageModelFromEnv(rt.ai.aiEnv);
     if (!model) {
-        throw new StudioHttpError(400, askDbAiKeyMissingMessage("AI enrichment suggestions"));
+        throw new StudioHttpError(400, ai.keyMissingMessage("AI enrichment suggestions"));
     }
     const candidates = await suggestEnrichment(buildSuggestionTarget(workspace, source), buildSuggestionContext(workspace, source.tableId), model);
     return candidates.map((candidate) => ({ text: candidate.text }));
 }
-async function askSampleQuestion(schemaDir, options) {
+async function suggestTenantPolicyDraft(state) {
+    const rt = getAskDbRuntimeConfig();
+    const model = await ai.createLanguageModelFromEnv(rt.ai.aiEnv);
+    if (!model) {
+        throw new StudioHttpError(400, ai.keyMissingMessage("AI tenant policy suggestion"));
+    }
+    // Build a compact DDL-like representation of the schema for the prompt
+    const workspace = requireWorkspace(state);
+    const schemaId = workspace.physical.schemaId;
+    const schemaDdl = workspace.tables
+        .map((t) => {
+        const cols = t.physical.columns
+            .map((c) => {
+            const parts = [c.name, c.type];
+            if (c.primaryKey)
+                parts.push("PRIMARY KEY");
+            if (!c.nullable)
+                parts.push("NOT NULL");
+            return `  ${parts.join(" ")}`;
+        })
+            .join(",\n");
+        const rels = (t.physical.relationships ?? [])
+            .map((r) => `  -- FK: ${r.from} -> ${r.to}`)
+            .join("\n");
+        return `CREATE TABLE ${t.physical.schema}.${t.physical.name} (\n${cols}\n);\n${rels}`.trim();
+    })
+        .join("\n\n");
+    const tableIds = workspace.tables.map((t) => t.physical.id);
+    const systemPrompt = `You are an expert database architect specializing in multi-tenant SaaS systems.
+You analyze database schemas and produce a tenant-policy configuration in a specific JSON format.
+
+The tenant policy identifies:
+1. **roots** — Tables whose rows represent tenants (organizations, companies, accounts, etc).
+   Each root has: id (table ID like "table:public.orgs"), tenantIdColumn (column ref like "table:public.orgs#id"), label (human-readable name).
+   If a root has a parent root (e.g. sub_org belongs to org), include parent: { root: "<parent-root-id>", foreignKey: "<fk-column-ref>" }.
+2. **hierarchy** — Edges showing parent→child relationships between roots or tenant-scoped tables.
+   Each edge: { parent: "<table-id>", child: "<table-id>", foreignKey: "<child-fk-column-ref>" }.
+3. **scopedTables** — Tables that belong to a tenant via a direct column or a JOIN chain.
+   Each entry: { id: "<table-id>", scopeThrough: [{ root: "<root-id>", column: "<column-ref>" }] } for direct FK,
+   or { id: "<table-id>", scopeThrough: [{ root: "<root-id>", join: [{ from: "<col-ref>", to: "<col-ref>" }] }] } for indirect.
+4. **polymorphicTables** — Tables with a type/id polymorphic association pattern.
+   Each entry: { id: "<table-id>", typeColumn: "<col-ref>", idColumn: "<col-ref>", mapping: { "<type_value>": "<target-table-id>" } }.
+5. **globalTables** — Tables shared across all tenants (lookup tables, reference data). Array of table IDs.
+6. **enforcement** — "strict" (reject queries on unknown tables) or "warn" (allow with warnings).
+
+Column references use the format "table:<schema>.<table>#<column>".
+Table IDs use the format "table:<schema>.<table>".
+
+Available table IDs: ${JSON.stringify(tableIds)}
+
+Return a JSON object with two keys:
+- "frontmatter": the TenantPolicyFrontmatter object (schemaId, enforcement, roots, hierarchy?, scopedTables?, polymorphicTables?, globalTables?)
+- "body": a markdown body string with sections "## Hierarchy", "## Scope rules", "## Sensitive interactions" explaining the policy in plain English.
+
+Return ONLY valid JSON, no markdown fences or explanations.`;
+    const userPrompt = `Analyze this database schema and draft a complete multi-tenant policy.
+
+Schema ID: ${schemaId}
+
+${schemaDdl}
+
+Look for:
+- Tables that represent organizations, companies, accounts, tenants, or workspaces — these are roots.
+- Foreign key columns like org_id, company_id, tenant_id, account_id that scope data to a tenant.
+- Tables with no tenant FK that are clearly shared/global (lookup tables, enums, config).
+- Polymorphic association patterns (type_column + id_column pointing to multiple parent tables).
+- Multi-level hierarchies (org → team → user, company → department → employee).
+
+If the schema has no clear multi-tenant pattern, still make your best guess at what would be the tenant root.`;
+    const result = await defaultGenerateText({
+        model,
+        system: systemPrompt,
+        prompt: userPrompt,
+    });
+    let parsed;
+    try {
+        parsed = JSON.parse(result.text);
+    }
+    catch {
+        throw new StudioHttpError(500, `AI returned invalid JSON for tenant policy suggestion. Raw text: ${result.text.slice(0, 500)}`);
+    }
+    // Validate frontmatter shape
+    const validation = tenantPolicyFrontmatterSchema.safeParse(parsed.frontmatter);
+    if (!validation.success) {
+        throw new StudioHttpError(500, `AI-generated tenant policy frontmatter is invalid: ${validation.error.issues.map((i) => i.message).join("; ")}`);
+    }
+    parsed.frontmatter = validation.data;
+    return parsed;
+}
+async function askSampleQuestion(state, options) {
     const rt = getAskDbRuntimeConfig();
     const mockSql = rt.dev.mockSql;
-    const aiConfig = mockSql ? undefined : resolveAskDbAiConfig(rt.ai.aiEnv);
+    const aiConfig = mockSql ? undefined : ai.resolveAiConfig(rt.ai.aiEnv);
     if (!mockSql && !aiConfig) {
-        throw new StudioHttpError(400, `${askDbAiKeyMissingMessage("Sample NL-to-SQL generation")} Set ASKDB_MOCK_SQL to bypass the live model.`);
+        throw new StudioHttpError(400, `${ai.keyMissingMessage("Sample NL-to-SQL generation")} Set ASKDB_MOCK_SQL to bypass the live model.`);
     }
-    const schema = loadSchema(schemaDir);
+    const schema = loadSchema(state.schemaDir);
     const retrievedChunks = [];
     const usage = createRequestUsageCollector();
     const ragIndex = options.useRag
-        ? await createCurrentStudioRagIndex(schemaDir, "using RAG for sample generation", usage)
+        ? await createCurrentStudioRagIndex(state, "using RAG for sample generation", usage)
         : undefined;
     const model = mockSql
         ? undefined
-        : (await createAskDbLanguageModelFromEnv(rt.ai.aiEnv, {
-            modelEnvVar: rt.ai.studioModel ? "ASKDB_STUDIO_MODEL" : undefined,
-        }));
+        : (await ai.createLanguageModelFromEnv(rt.ai.aiEnv));
     const schemaProvider = "provider" in schema && typeof schema.provider === "string"
         ? schema.provider
         : undefined;
@@ -177,6 +538,8 @@ async function askSampleQuestion(schemaDir, options) {
         model,
         dialect,
         explain: true,
+        ...(options.tenantScope ? { tenantScope: options.tenantScope } : {}),
+        ...(options.tenantSqlMode ? { tenantSqlMode: options.tenantSqlMode } : {}),
         ...(ragIndex
             ? {
                 retriever: async (params) => {
@@ -185,7 +548,7 @@ async function askSampleQuestion(schemaDir, options) {
                         results = await ragIndex.retriever(params);
                     }
                     catch (error) {
-                        throw formatStudioRagEmbeddingError(error, ragIndex.config);
+                        throw formatStudioRagOperationError(error, ragIndex.config);
                     }
                     retrievedChunks.splice(0, retrievedChunks.length, ...results);
                     return results;
@@ -202,6 +565,8 @@ async function askSampleQuestion(schemaDir, options) {
             : {
                 generateText: createTrackedGenerateText(usage),
             },
+    }).finally(async () => {
+        await ragIndex?.dispose();
     });
     return {
         sql: result.sql,
@@ -211,90 +576,115 @@ async function askSampleQuestion(schemaDir, options) {
             enabled: options.useRag,
             chunks: options.useRag ? retrievedChunks.map(serializeRagResult) : [],
         },
+        tenant: options.tenantScope
+            ? {
+                enabled: true,
+                sqlMode: options.tenantSqlMode ?? "sql-only",
+                bindings: result.tenantBindings ?? [],
+                params: result.tenantParams ?? [],
+            }
+            : null,
         usage: usage.toDto(),
     };
 }
-function getRagStatus(schemaDir) {
+async function getRagStatus(state) {
     const config = resolveStudioRagEmbedderConfig();
-    const sources = loadChunkerSourcesFromDir(schemaDir);
+    const sources = loadChunkerSourcesFromDir(state.schemaDir);
     const chunkResult = chunkSchema(sources);
-    const lockPath = join(schemaDir, "schema.lock.json");
-    const embeddingsJsonPath = join(schemaDir, "schema.embeddings.json");
-    const embeddingsBinPath = join(schemaDir, "schema.embeddings.bin");
+    const lockPath = join(state.schemaDir, "schema.lock.json");
     const lock = readOptionalJson(lockPath);
-    const embeddings = readOptionalJson(embeddingsJsonPath);
     const currentHashes = Object.fromEntries(chunkResult.chunks.map((chunk) => [chunk.id, chunkContentHash(chunk.text)]));
     const lockHashes = lock?.hashes ?? {};
     const hashIds = Object.keys(currentHashes);
-    const stale = !lock ||
-        !existsSync(embeddingsJsonPath) ||
-        !existsSync(embeddingsBinPath) ||
-        lock.embedderId !== config.embedderId ||
-        embeddings?.dimensions !== config.dimensions ||
-        Object.keys(lockHashes).length !== hashIds.length ||
-        hashIds.some((id) => lockHashes[id] !== currentHashes[id]);
-    return {
-        schemaId: sources.schema.schemaId,
-        embedder: {
-            kind: config.kind,
-            label: config.label,
-            configured: config.configured,
-            expectedId: config.embedderId,
-            indexedId: lock?.embedderId ?? null,
-            provider: config.kind === "ai-sdk" ? config.provider : null,
-            model: config.kind === "ai-sdk" ? config.model : null,
-            baseUrl: config.kind === "ai-sdk" ? config.baseUrl ?? null : null,
-        },
-        embedderId: lock?.embedderId ?? config.embedderId,
-        expectedEmbedderId: config.embedderId,
-        hasIndex: Boolean(lock && existsSync(embeddingsJsonPath) && existsSync(embeddingsBinPath)),
-        stale,
-        updatedAt: lock?.updatedAt ?? null,
-        chunksTotal: chunkResult.chunks.length,
-        chunksIndexed: Array.isArray(embeddings?.records) ? embeddings.records.length : 0,
-        dimensions: typeof embeddings?.dimensions === "number" ? embeddings.dimensions : config.dimensions,
-        expectedDimensions: config.dimensions,
-        sensitiveExcluded: chunkResult.stats.sensitiveExcluded,
-        sensitiveIncluded: chunkResult.stats.sensitiveIncluded,
-        files: {
-            lock: existsSync(lockPath),
-            embeddingsJson: existsSync(embeddingsJsonPath),
-            embeddingsBin: existsSync(embeddingsBinPath),
-        },
-    };
+    const store = await openStudioRagStore(state, config.dimensions);
+    try {
+        const chunksIndexed = await countStudioRagStoreChunks(store, sources.schema.schemaId);
+        const stale = !lock ||
+            lock.embedderId !== config.embedderId ||
+            Object.keys(lockHashes).length !== hashIds.length ||
+            hashIds.some((id) => lockHashes[id] !== currentHashes[id]) ||
+            chunksIndexed !== hashIds.length;
+        const fileArtifacts = store.kind === "file" && store.basePath
+            ? {
+                lock: existsSync(lockPath),
+                embeddingsJson: existsSync(`${store.basePath}.embeddings.json`),
+                embeddingsBin: existsSync(`${store.basePath}.embeddings.bin`),
+            }
+            : {
+                lock: existsSync(lockPath),
+                embeddingsJson: false,
+                embeddingsBin: false,
+            };
+        return {
+            schemaId: sources.schema.schemaId,
+            store: {
+                kind: store.kind,
+                basePath: store.basePath ?? null,
+                table: store.table ?? null,
+                indexStrategy: store.indexStrategy ?? null,
+            },
+            embedder: {
+                kind: config.kind,
+                label: config.label,
+                configured: config.configured,
+                expectedId: config.embedderId,
+                indexedId: lock?.embedderId ?? null,
+                provider: config.kind === "ai-sdk" ? config.provider : null,
+                model: config.kind === "ai-sdk" ? config.model : null,
+                baseUrl: config.kind === "ai-sdk" ? config.baseUrl ?? null : null,
+            },
+            embedderId: lock?.embedderId ?? config.embedderId,
+            expectedEmbedderId: config.embedderId,
+            hasIndex: Boolean(lock && chunksIndexed > 0),
+            stale,
+            updatedAt: lock?.updatedAt ?? null,
+            chunksTotal: chunkResult.chunks.length,
+            chunksIndexed,
+            dimensions: config.dimensions,
+            expectedDimensions: config.dimensions,
+            sensitiveExcluded: chunkResult.stats.sensitiveExcluded,
+            sensitiveIncluded: chunkResult.stats.sensitiveIncluded,
+            files: fileArtifacts,
+        };
+    }
+    finally {
+        await store.dispose();
+    }
 }
-async function indexRag(schemaDir) {
+async function indexRag(state) {
     const config = resolveStudioRagEmbedderConfig();
     if (!config.configured) {
         throw new StudioHttpError(400, studioRagAiSdkKeyMissingMessage());
     }
     const usage = createRequestUsageCollector();
-    clearIncompatibleRagStore(schemaDir, config);
-    const sources = loadChunkerSourcesFromDir(schemaDir);
-    const store = createFileStore({ basePath: join(schemaDir, "schema") });
+    clearIncompatibleRagStore(state, config);
+    const sources = loadChunkerSourcesFromDir(state.schemaDir);
+    const store = await openStudioRagStore(state, config.dimensions);
     let result;
     try {
         result = await buildSchemaIndex({
             schema: sources,
             embedder: await createStudioRagEmbedder(config, usage),
-            store,
+            store: store.store,
             embedderId: config.embedderId,
-            lockFilePath: join(schemaDir, "schema.lock.json"),
+            lockFilePath: join(state.schemaDir, "schema.lock.json"),
         });
     }
     catch (error) {
-        throw formatStudioRagEmbeddingError(error, config);
+        throw formatStudioRagOperationError(error, config);
+    }
+    finally {
+        await store.dispose();
     }
-    store.flush();
     return {
-        status: getRagStatus(schemaDir),
+        status: await getRagStatus(state),
         stats: result.stats,
         usage: usage.toDto(),
     };
 }
-async function queryRag(schemaDir, query) {
+async function queryRag(state, query) {
     const usage = createRequestUsageCollector();
-    const ragIndex = await createCurrentStudioRagIndex(schemaDir, "querying chunks", usage);
+    const ragIndex = await createCurrentStudioRagIndex(state, "querying chunks", usage);
     let results;
     try {
         results = await ragIndex.retriever({
@@ -307,7 +697,10 @@ async function queryRag(schemaDir, query) {
         });
     }
     catch (error) {
-        throw formatStudioRagEmbeddingError(error, ragIndex.config);
+        throw formatStudioRagOperationError(error, ragIndex.config);
+    }
+    finally {
+        await ragIndex.dispose();
     }
     return {
         question: query.question,
@@ -316,12 +709,12 @@ async function queryRag(schemaDir, query) {
         usage: usage.toDto(),
     };
 }
-async function createCurrentStudioRagIndex(schemaDir, action, usage) {
+async function createCurrentStudioRagIndex(state, action, usage) {
     const config = resolveStudioRagEmbedderConfig();
     if (!config.configured) {
         throw new StudioHttpError(400, studioRagAiSdkKeyMissingMessage());
     }
-    const status = getRagStatus(schemaDir);
+    const status = (await getRagStatus(state));
     if (!status.hasIndex) {
         throw new StudioHttpError(400, `Build the RAG index before ${action}.`);
     }
@@ -331,6 +724,7 @@ async function createCurrentStudioRagIndex(schemaDir, action, usage) {
     if (!status.schemaId || typeof status.chunksTotal !== "number") {
         throw new StudioHttpError(500, "Studio RAG status is missing schema metadata.");
     }
+    const store = await openStudioRagStore(state, config.dimensions);
     return {
         config,
         status: {
@@ -339,10 +733,85 @@ async function createCurrentStudioRagIndex(schemaDir, action, usage) {
         },
         retriever: createRetriever({
             embedder: await createStudioRagEmbedder(config, usage),
-            store: createFileStore({ basePath: join(schemaDir, "schema") }),
+            store: store.store,
         }),
+        dispose: async () => {
+            await store.dispose();
+        },
     };
 }
+function resolveStudioRagStoreConfig(state) {
+    const rt = getAskDbRuntimeConfig();
+    const kind = rt.structured.rag.store;
+    if (kind === "memory")
+        return { kind };
+    if (kind === "file") {
+        const basePath = rt.structured.rag.storeConfig.file?.basePath?.trim();
+        return { kind, basePath: basePath ? resolve(basePath) : join(state.schemaDir, "schema") };
+    }
+    const connectionString = pickFlat(rt.flat, "ASKDB_PGVECTOR_URL");
+    return {
+        kind,
+        connectionString,
+        table: rt.structured.rag.storeConfig.pgvector?.table?.trim() || undefined,
+        indexStrategy: pickFlat(rt.flat, "ASKDB_PGVECTOR_INDEX_STRATEGY"),
+    };
+}
+async function openStudioRagStore(state, dimensions) {
+    const config = resolveStudioRagStoreConfig(state);
+    if (config.kind === "memory") {
+        state.ragMemoryStore ??= createMemoryStore();
+        return {
+            kind: "memory",
+            store: state.ragMemoryStore,
+            dispose: async () => { },
+        };
+    }
+    if (config.kind === "file") {
+        const store = createFileStore({ basePath: config.basePath });
+        return {
+            kind: "file",
+            store,
+            basePath: config.basePath,
+            dispose: async () => {
+                store.flush();
+            },
+        };
+    }
+    if (!config.connectionString) {
+        throw new StudioHttpError(400, 'Studio pgvector RAG requires `ASKDB_PGVECTOR_URL` via `askdb.config.ts`.');
+    }
+    const pgvectorFactory = studioPgvectorStoreFactoryForTests ?? createPgvectorStore;
+    const store = pgvectorFactory({
+        connectionString: config.connectionString,
+        table: config.table,
+        dimensions,
+        ...(config.indexStrategy ? { indexStrategy: config.indexStrategy } : {}),
+    });
+    await store.ensureSchema();
+    return {
+        kind: "pgvector",
+        store,
+        table: config.table,
+        indexStrategy: config.indexStrategy,
+        dispose: async () => {
+            await store.close();
+        },
+    };
+}
+async function countStudioRagStoreChunks(store, schemaId) {
+    if (typeof store.store.count === "function") {
+        return store.store.count({ schemaId });
+    }
+    if (typeof store.store.size === "function") {
+        return store.store.size();
+    }
+    return 0;
+}
+function pickFlat(flat, key) {
+    const value = flat[key];
+    return typeof value === "string" && value.trim() !== "" ? value.trim() : undefined;
+}
 function pickEnv(env, key) {
     const v = env[key];
     return typeof v === "string" && v.trim() !== "" ? v.trim() : undefined;
@@ -350,7 +819,7 @@ function pickEnv(env, key) {
 function resolveStudioRagEmbedderConfig() {
     const rt = getAskDbRuntimeConfig();
     const base = rt.ai.aiEnv;
-    const explicitKind = pickEnv(base, "ASKDB_RAG_EMBEDDER") ?? pickEnv(base, "ASKDB_STUDIO_RAG_EMBEDDER");
+    const explicitKind = pickEnv(base, "ASKDB_RAG_EMBEDDER");
     const kind = explicitKind?.toLowerCase();
     if (kind === "mock") {
         return {
@@ -365,8 +834,8 @@ function resolveStudioRagEmbedderConfig() {
         throw new StudioHttpError(400, `Unsupported Studio RAG embedder: ${kind}`);
     }
     const env = buildStudioRagEmbeddingEnv(kind, base);
-    const aiConfig = resolveAskDbEmbeddingConfig(env, {
-        modelEnvVar: "ASKDB_STUDIO_RAG_EMBEDDER_MODEL",
+    const aiConfig = ai.resolveEmbeddingConfig(env, {
+        modelEnvVar: "ASKDB_RAG_EMBEDDER_MODEL",
         modelDefault: DEFAULT_EMBEDDING_MODEL,
     });
     if (!aiConfig && kind === undefined) {
@@ -380,8 +849,7 @@ function resolveStudioRagEmbedderConfig() {
     }
     const provider = aiConfig?.provider ?? fallbackStudioRagProvider(kind, base);
     const model = aiConfig?.model ?? DEFAULT_EMBEDDING_MODEL;
-    const dimensionOverride = readPositiveIntegerEnv(pickEnv(base, "ASKDB_RAG_EMBEDDER_DIMENSIONS") ??
-        pickEnv(base, "ASKDB_STUDIO_RAG_EMBEDDER_DIMENSIONS"));
+    const dimensionOverride = readPositiveIntegerEnv(pickEnv(base, "ASKDB_RAG_EMBEDDER_DIMENSIONS"));
     const dimensions = dimensionOverride ?? defaultEmbeddingDimensions(model);
     return {
         kind: "ai-sdk",
@@ -397,14 +865,12 @@ function resolveStudioRagEmbedderConfig() {
     };
 }
 function buildStudioRagEmbeddingEnv(kind, base) {
-    const apiKeyOverride = pickEnv(base, "ASKDB_RAG_EMBEDDER_API_KEY") ?? pickEnv(base, "ASKDB_STUDIO_RAG_EMBEDDER_API_KEY");
-    const modelOverride = pickEnv(base, "ASKDB_RAG_EMBEDDER_MODEL") ?? pickEnv(base, "ASKDB_STUDIO_RAG_EMBEDDER_MODEL");
-    const baseUrlOverride = pickEnv(base, "ASKDB_RAG_EMBEDDER_BASE_URL") ?? pickEnv(base, "ASKDB_STUDIO_RAG_EMBEDDER_BASE_URL");
+    const apiKeyOverride = pickEnv(base, "ASKDB_RAG_EMBEDDER_API_KEY");
+    const baseUrlOverride = pickEnv(base, "ASKDB_RAG_EMBEDDER_BASE_URL");
     return {
         ...base,
         ...(kind === "openai" ? { ASKDB_AI_PROVIDER: "openai" } : {}),
         ...(apiKeyOverride ? { ASKDB_AI_API_KEY: apiKeyOverride } : {}),
-        ...(modelOverride ? { ASKDB_STUDIO_RAG_EMBEDDER_MODEL: modelOverride } : {}),
         ...(baseUrlOverride ? { ASKDB_AI_BASE_URL: baseUrlOverride } : {}),
     };
 }
@@ -426,7 +892,7 @@ async function createStudioRagEmbedder(config, usage) {
     if (!config.aiConfig) {
         throw new StudioHttpError(400, studioRagAiSdkKeyMissingMessage());
     }
-    const model = await createAskDbEmbeddingModel(config.aiConfig, {
+    const model = await ai.createEmbeddingModel(config.aiConfig, {
         dimensions: config.requestDimensions,
     });
     return createAiSdkEmbedder({
@@ -441,11 +907,16 @@ async function createStudioRagEmbedder(config, usage) {
         },
     });
 }
-function formatStudioRagEmbeddingError(error, config) {
+function formatStudioRagOperationError(error, config) {
+    if (error instanceof StudioHttpError)
+        return error;
     if (config.kind === "mock") {
         return new StudioHttpError(500, error instanceof Error ? error.message : String(error));
     }
     const apiError = findApiCallError(error);
+    if (!apiError) {
+        return new StudioHttpError(500, error instanceof Error ? error.message : String(error));
+    }
     const parts = [
         `Studio RAG embedding request failed for provider ${config.provider}, model ${config.model}.`,
     ];
@@ -488,15 +959,18 @@ function truncateForMessage(value) {
         return undefined;
     return trimmed.length > 500 ? `${trimmed.slice(0, 500)}...` : trimmed;
 }
-function clearIncompatibleRagStore(schemaDir, config) {
-    const embeddingsJsonPath = join(schemaDir, "schema.embeddings.json");
+function clearIncompatibleRagStore(state, config) {
+    const store = resolveStudioRagStoreConfig(state);
+    if (store.kind !== "file")
+        return;
+    const embeddingsJsonPath = `${store.basePath}.embeddings.json`;
     const embeddings = readOptionalJson(embeddingsJsonPath);
     if (embeddings?.dimensions === undefined || embeddings.dimensions === config.dimensions)
         return;
     for (const path of [
-        join(schemaDir, "schema.embeddings.json"),
-        join(schemaDir, "schema.embeddings.bin"),
-        join(schemaDir, "schema.lock.json"),
+        `${store.basePath}.embeddings.json`,
+        `${store.basePath}.embeddings.bin`,
+        join(state.schemaDir, "schema.lock.json"),
     ]) {
         rmSync(path, { force: true });
     }
@@ -546,6 +1020,234 @@ function stableTokenHash(token) {
     }
     return hash >>> 0;
 }
+// ---------------------------------------------------------------------------
+// Playground history helpers
+// ---------------------------------------------------------------------------
+const HISTORY_MAX_STORED = 200;
+const HISTORY_MAX_RETURNED = 50;
+function playgroundHistoryPath(schemaDir) {
+    return join(schemaDir, "playground-history.json");
+}
+function readPlaygroundHistory(schemaDir) {
+    const raw = readOptionalJson(playgroundHistoryPath(schemaDir));
+    const entries = Array.isArray(raw) ? raw : [];
+    return { entries: entries.slice(0, HISTORY_MAX_RETURNED) };
+}
+function appendPlaygroundHistory(schemaDir, entry) {
+    const histPath = playgroundHistoryPath(schemaDir);
+    const raw = readOptionalJson(histPath);
+    const existing = Array.isArray(raw)
+        ? raw
+        : [];
+    const newEntry = {
+        ...entry,
+        id: randomUUID(),
+        timestamp: new Date().toISOString(),
+    };
+    const updated = [newEntry, ...existing].slice(0, HISTORY_MAX_STORED);
+    writeFileSync(histPath, JSON.stringify(updated, null, 2), "utf8");
+}
+function deletePlaygroundHistoryEntry(schemaDir, id) {
+    const histPath = playgroundHistoryPath(schemaDir);
+    const raw = readOptionalJson(histPath);
+    const existing = Array.isArray(raw)
+        ? raw
+        : [];
+    const filtered = existing.filter((e) => e.id !== id);
+    if (filtered.length === existing.length) {
+        return { ok: false, error: "not found" };
+    }
+    writeFileSync(histPath, JSON.stringify(filtered, null, 2), "utf8");
+    return { ok: true };
+}
+function parsePlaygroundHistoryEntry(body) {
+    if (!isRecord(body)) {
+        throw new StudioHttpError(400, "Request body must be a JSON object.");
+    }
+    if (typeof body.question !== "string" || body.question.trim() === "") {
+        throw new StudioHttpError(400, "`question` is required.");
+    }
+    if (body.mode !== "full" && body.mode !== "rag") {
+        throw new StudioHttpError(400, "`mode` must be `full` or `rag`.");
+    }
+    if (typeof body.sqlMode !== "string") {
+        throw new StudioHttpError(400, "`sqlMode` is required.");
+    }
+    if (typeof body.sql !== "string") {
+        throw new StudioHttpError(400, "`sql` is required.");
+    }
+    return body;
+}
+// ---------------------------------------------------------------------------
+// Execute endpoint helpers
+// ---------------------------------------------------------------------------
+async function getExecuteStatus(schemaDir) {
+    const rt = getAskDbRuntimeConfig();
+    const { provider, databaseUrl, file } = rt.studio.execute;
+    const def = EXECUTE_DRIVER_REGISTRY[provider];
+    const projectRoot = findProjectRoot(schemaDir) ?? schemaDir;
+    const installed = isDriverInstalled(def.packageName, projectRoot);
+    const connectionKind = provider === "sqlite" ? "file" : "url";
+    const configured = connectionKind === "file" ? Boolean(file) : Boolean(databaseUrl);
+    return {
+        provider,
+        label: def.label,
+        configured,
+        connectionKind,
+        packageName: def.packageName,
+        installed,
+        installCommand: def.installCommand,
+        canInstallFromStudio: true,
+        manualInstallReason: null,
+    };
+}
+async function installExecuteDriver(body, schemaDir) {
+    const rt = getAskDbRuntimeConfig();
+    const configuredProvider = rt.studio.execute.provider;
+    // Accept the configured provider or an explicit allowlisted provider.
+    let provider = configuredProvider;
+    if (isRecord(body) && typeof body.provider === "string") {
+        const requested = body.provider;
+        if (!(requested in EXECUTE_DRIVER_REGISTRY)) {
+            return {
+                ok: false,
+                provider: configuredProvider,
+                packageName: EXECUTE_DRIVER_REGISTRY[configuredProvider].packageName,
+                command: [],
+                error: `Unknown provider: ${body.provider}`,
+                installed: false,
+            };
+        }
+        provider = requested;
+    }
+    const def = EXECUTE_DRIVER_REGISTRY[provider];
+    const { packageManager, command, args, manualReason } = detectPackageManager(schemaDir, def.packageName);
+    if (manualReason) {
+        return {
+            ok: false,
+            provider,
+            packageName: def.packageName,
+            command: [],
+            error: manualReason,
+            installed: false,
+        };
+    }
+    // Run the install.
+    const { stdout, stderr, code } = await spawnCommand(command, args, schemaDir);
+    const projectRoot = findProjectRoot(schemaDir) ?? schemaDir;
+    const installed = isDriverInstalled(def.packageName, projectRoot);
+    return {
+        ok: code === 0,
+        provider,
+        packageName: def.packageName,
+        command: [command, ...args],
+        stdout,
+        stderr,
+        error: code !== 0 ? `${packageManager} exited with code ${code}` : undefined,
+        installed,
+    };
+}
+function detectPackageManager(cwd, packageName) {
+    // Walk up from schemaDir to find the project root's lockfile.
+    const projectRoot = findProjectRoot(cwd);
+    if (!projectRoot) {
+        return {
+            packageManager: "",
+            command: "",
+            args: [],
+            manualReason: `Could not locate a package.json above ${cwd}. Run the install manually.`,
+        };
+    }
+    if (existsSync(join(projectRoot, "pnpm-lock.yaml"))) {
+        return {
+            packageManager: "pnpm",
+            command: "pnpm",
+            args: ["add", packageName],
+            manualReason: null,
+        };
+    }
+    if (existsSync(join(projectRoot, "package-lock.json"))) {
+        return {
+            packageManager: "npm",
+            command: "npm",
+            args: ["install", packageName],
+            manualReason: null,
+        };
+    }
+    if (existsSync(join(projectRoot, "yarn.lock"))) {
+        return {
+            packageManager: "yarn",
+            command: "yarn",
+            args: ["add", packageName],
+            manualReason: null,
+        };
+    }
+    if (existsSync(join(projectRoot, "bun.lockb")) || existsSync(join(projectRoot, "bun.lock"))) {
+        return {
+            packageManager: "bun",
+            command: "bun",
+            args: ["add", packageName],
+            manualReason: null,
+        };
+    }
+    return {
+        packageManager: "",
+        command: "",
+        args: [],
+        manualReason: `No recognized lockfile found in ${projectRoot}. Run: pnpm add ${packageName}`,
+    };
+}
+function findProjectRoot(startDir) {
+    let current = resolve(startDir);
+    for (let i = 0; i < 20; i++) {
+        if (existsSync(join(current, "package.json")))
+            return current;
+        const parent = dirname(current);
+        if (parent === current)
+            return null;
+        current = parent;
+    }
+    return null;
+}
+function spawnCommand(command, args, cwd) {
+    return new Promise((resolve) => {
+        const proc = spawn(command, args, { cwd, shell: false, stdio: ["ignore", "pipe", "pipe"] });
+        const stdoutChunks = [];
+        const stderrChunks = [];
+        proc.stdout.on("data", (chunk) => stdoutChunks.push(chunk));
+        proc.stderr.on("data", (chunk) => stderrChunks.push(chunk));
+        proc.on("close", (code) => {
+            resolve({
+                stdout: Buffer.concat(stdoutChunks).toString("utf8").slice(0, 4096),
+                stderr: Buffer.concat(stderrChunks).toString("utf8").slice(0, 4096),
+                code: code ?? 1,
+            });
+        });
+        proc.on("error", (err) => {
+            resolve({ stdout: "", stderr: err.message, code: 1 });
+        });
+    });
+}
+async function executeQuery(body, schemaDir) {
+    const rt = getAskDbRuntimeConfig();
+    const { provider, databaseUrl, file } = rt.studio.execute;
+    if (!isRecord(body) || typeof body.sql !== "string" || body.sql.trim() === "") {
+        return { ok: false, error: "`sql` is required." };
+    }
+    const sql = body.sql;
+    const params = Array.isArray(body.params) ? body.params : [];
+    const projectRoot = findProjectRoot(schemaDir) ?? schemaDir;
+    const def = EXECUTE_DRIVER_REGISTRY[provider];
+    return def.execute({ connectionString: databaseUrl, file, sql, params, projectRoot });
+}
+function isLoopbackRequest(req) {
+    const socket = req.socket;
+    const remoteAddress = socket?.remoteAddress ?? "";
+    return (remoteAddress === "127.0.0.1" ||
+        remoteAddress === "::1" ||
+        remoteAddress === "::ffff:127.0.0.1" ||
+        remoteAddress === "localhost");
+}
 function readOptionalJson(path) {
     if (!existsSync(path))
         return undefined;
@@ -556,6 +1258,17 @@ function readOptionalJson(path) {
         return undefined;
     }
 }
+function parseConceptsBody(body) {
+    if (!isRecord(body) || !Array.isArray(body.concepts)) {
+        throw new StudioHttpError(400, "`concepts` array is required.");
+    }
+    for (const concept of body.concepts) {
+        if (!isRecord(concept) || typeof concept.id !== "string" || typeof concept.label !== "string") {
+            throw new StudioHttpError(400, "Each concept requires a string `id` and `label`.");
+        }
+    }
+    return body.concepts;
+}
 function parseTableDraftBody(body) {
     if (!isRecord(body) || !isRecord(body.draft)) {
         throw new StudioHttpError(400, "`draft` is required.");
@@ -604,9 +1317,26 @@ function parseAskBody(body) {
     if (mode !== "full" && mode !== "rag") {
         throw new StudioHttpError(400, "`mode` must be `full` or `rag`.");
     }
+    let tenantScope;
+    if (body.tenantScope !== undefined && body.tenantScope !== null) {
+        const parsed = tenantScopeSchema.safeParse(body.tenantScope);
+        if (!parsed.success) {
+            throw new StudioHttpError(400, `Invalid tenantScope: ${parsed.error.message}`);
+        }
+        tenantScope = parsed.data;
+    }
+    let tenantSqlMode;
+    if (typeof body.tenantSqlMode === "string") {
+        if (body.tenantSqlMode !== "sql-only" && body.tenantSqlMode !== "sql-params") {
+            throw new StudioHttpError(400, "`tenantSqlMode` must be `sql-only` or `sql-params`.");
+        }
+        tenantSqlMode = body.tenantSqlMode;
+    }
     return {
         question: body.question.trim(),
         useRag: mode === "rag",
+        tenantScope,
+        tenantSqlMode,
     };
 }
 function parseRagQuery(body) {
@@ -676,8 +1406,10 @@ function createRequestUsageCollector() {
 function normalizeGenerationUsage(value) {
     if (!isRecord(value))
         return undefined;
-    const promptTokens = readFiniteNumber(value.promptTokens);
-    const completionTokens = readFiniteNumber(value.completionTokens);
+    // ai@6 LanguageModelUsage uses inputTokens/outputTokens rather than the
+    // legacy promptTokens/completionTokens names.
+    const promptTokens = readFiniteNumber(value.promptTokens) ?? readFiniteNumber(value.inputTokens);
+    const completionTokens = readFiniteNumber(value.completionTokens) ?? readFiniteNumber(value.outputTokens);
     const totalTokens = readFiniteNumber(value.totalTokens);
     if (promptTokens === undefined && completionTokens === undefined && totalTokens === undefined) {
         return undefined;