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

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
package/dist/server.js CHANGED
@@ -1,26 +1,41 @@
1
- import { existsSync, readFileSync, rmSync } from "node:fs";
1
+ import { randomUUID } from "node:crypto";
2
+ import { existsSync, readFileSync, rmSync, writeFileSync } from "node:fs";
2
3
  import { createServer } from "node:http";
3
- import { extname, relative, join, resolve } from "node:path";
4
+ import { extname, relative, join, resolve, dirname, sep } from "node:path";
5
+ import { spawn } from "node:child_process";
4
6
  import { fileURLToPath } from "node:url";
5
7
  import { generateText as defaultGenerateText } from "ai";
6
- import { getAskDbRuntimeConfig } from "@askdb/config";
7
- import { ask, askDbAiKeyMissingMessage, createAskDbEmbeddingModel, createAskDbLanguageModelFromEnv, isBuiltInDialectId, loadSchema, resolveAskDbAiConfig, resolveAskDbEmbeddingConfig, suggestEnrichment, } from "@askdb/core";
8
- import { buildSchemaIndex, chunkContentHash, chunkSchema, createRetriever, loadChunkerSourcesFromDir, } from "@askdb/rag";
9
- import { createAiSdkEmbedder } from "@askdb/rag/embedders/ai-sdk";
10
- import { createFileStore } from "@askdb/rag/stores/file";
11
- import { buildDefaultTableBody, buildFrontmatter, buildSuggestionContext, buildSuggestionTarget, buildTableDraft, loadWorkspace, replaceH2Section, replaceTableDescription, saveTable, } from "@askdb/enrich";
8
+ import { bootstrapAskDbEnv, getAskDbRuntimeConfig } from "@askdb/config";
9
+ import { createAiRegistry, } from "@askdb/ai";
10
+ import { anthropicProvider } from "@askdb/ai-anthropic";
11
+ import { azureProvider } from "@askdb/ai-azure";
12
+ import { googleProvider } from "@askdb/ai-google";
13
+ import { openaiProvider } from "@askdb/ai-openai";
14
+ import { ask, isBuiltInDialectId, loadSchema, suggestEnrichment, tenantScopeSchema, writeTenantPolicyMarkdown, tenantPolicyFrontmatterSchema, } from "@askdb/core";
15
+ import { buildSchemaIndex, chunkContentHash, chunkSchema, createAiSdkEmbedder, createFileStore, createMemoryStore, createPgvectorStore, createRetriever, loadChunkerSourcesFromDir, } from "@askdb/rag";
16
+ import { buildDefaultTableBody, buildFrontmatter, buildSuggestionContext, buildSuggestionTarget, buildTableDraft, loadWorkspace, replaceH2Section, replaceTableDescription, saveConcepts, saveTable, } from "@askdb/enrich";
17
+ import { EXECUTE_DRIVER_REGISTRY, isDriverInstalled } from "./execute-registry.js";
18
+ import { resolveStudioIntrospectionPlan, runStudioIntrospection } from "./introspection.js";
19
+ import { probeSetupState, writeSetupConfig, SetupError } from "./setup.js";
20
+ const ai = createAiRegistry([openaiProvider, azureProvider, googleProvider, anthropicProvider]);
12
21
  const CLIENT_DIR = fileURLToPath(new URL("./client/", import.meta.url));
13
22
  const STUDIO_RAG_MOCK_DIMENSIONS = 64;
14
23
  const STUDIO_RAG_MOCK_EMBEDDER_ID = `studio:mock-lexical-${STUDIO_RAG_MOCK_DIMENSIONS}`;
15
24
  const DEFAULT_EMBEDDING_MODEL = "text-embedding-3-small";
25
+ let studioPgvectorStoreFactoryForTests;
26
+ export function setStudioPgvectorStoreFactoryForTests(factory) {
27
+ studioPgvectorStoreFactoryForTests = factory;
28
+ }
16
29
  export function createStudioServer(options) {
17
30
  const schemaDir = resolve(options.schema);
18
- if (!existsSync(schemaDir)) {
31
+ const setupReason = options.setupReason ?? null;
32
+ if (!setupReason && !existsSync(schemaDir)) {
19
33
  throw new Error(`schema directory not found: ${schemaDir}`);
20
34
  }
21
35
  const state = {
22
36
  schemaDir,
23
- workspace: loadWorkspace(schemaDir),
37
+ workspace: setupReason ? null : loadWorkspace(schemaDir),
38
+ setupReason,
24
39
  };
25
40
  return createServer(async (req, res) => {
26
41
  try {
@@ -31,39 +46,122 @@ export function createStudioServer(options) {
31
46
  if (req.method === "GET" && url.pathname.startsWith("/assets/")) {
32
47
  return serveClientFile(res, decodeURIComponent(url.pathname.slice(1)));
33
48
  }
49
+ if (req.method === "GET" && url.pathname === "/api/setup/status") {
50
+ return writeJson(res, 200, buildSetupStatus(state));
51
+ }
52
+ if (req.method === "POST" && url.pathname === "/api/setup/config") {
53
+ if (!isLoopbackRequest(req)) {
54
+ return writeJson(res, 403, { error: { message: "Setup is only available from loopback clients." } });
55
+ }
56
+ const body = await readJson(req);
57
+ const result = handleSetupConfig(state, body);
58
+ return writeJson(res, 200, result);
59
+ }
60
+ if (req.method === "POST" && url.pathname === "/api/setup/introspect") {
61
+ if (!isLoopbackRequest(req)) {
62
+ return writeJson(res, 403, { error: { message: "Setup is only available from loopback clients." } });
63
+ }
64
+ const result = await handleSetupIntrospect(state);
65
+ return writeJson(res, 200, result);
66
+ }
67
+ // Everything below needs a loaded workspace — route setup-mode clients to the wizard.
68
+ if (state.setupReason && url.pathname.startsWith("/api/")) {
69
+ return writeJson(res, 409, {
70
+ error: { message: "Studio needs setup. Open Studio in the browser to finish the guided setup." },
71
+ setup: buildSetupStatus(state),
72
+ });
73
+ }
34
74
  if (req.method === "GET" && url.pathname === "/api/workspace") {
35
- return writeJson(res, 200, serializeWorkspace(state.workspace));
75
+ return writeJson(res, 200, serializeWorkspace(requireWorkspace(state)));
76
+ }
77
+ if (req.method === "GET" && url.pathname === "/api/introspect/status") {
78
+ requireWorkspace(state);
79
+ return writeJson(res, 200, buildIntrospectionPlanDto());
80
+ }
81
+ if (req.method === "POST" && url.pathname === "/api/introspect") {
82
+ if (!isLoopbackRequest(req)) {
83
+ return writeJson(res, 403, { error: { message: "Resync is only available from loopback clients." } });
84
+ }
85
+ const result = await handleResync(state);
86
+ return writeJson(res, 200, result);
36
87
  }
37
88
  if (req.method === "POST" && url.pathname.startsWith("/api/tables/")) {
38
89
  const tableId = decodeURIComponent(url.pathname.slice("/api/tables/".length));
39
90
  const body = await readJson(req);
40
91
  const draft = parseTableDraftBody(body);
41
92
  saveDraft(state, tableId, draft);
42
- return writeJson(res, 200, serializeWorkspace(state.workspace));
93
+ return writeJson(res, 200, serializeWorkspace(requireWorkspace(state)));
94
+ }
95
+ if (req.method === "POST" && url.pathname === "/api/concepts") {
96
+ const body = await readJson(req);
97
+ const concepts = parseConceptsBody(body);
98
+ saveConceptsDraft(state, concepts);
99
+ return writeJson(res, 200, serializeWorkspace(requireWorkspace(state)));
100
+ }
101
+ if (req.method === "POST" && url.pathname === "/api/tenant-policy") {
102
+ const body = await readJson(req);
103
+ const { frontmatter, bodyText } = parseTenantPolicyBody(body);
104
+ saveTenantPolicy(state, frontmatter, bodyText);
105
+ return writeJson(res, 200, serializeWorkspace(requireWorkspace(state)));
106
+ }
107
+ if (req.method === "POST" && url.pathname === "/api/suggest-tenant-policy") {
108
+ const draft = await suggestTenantPolicyDraft(state);
109
+ return writeJson(res, 200, draft);
43
110
  }
44
111
  if (req.method === "POST" && url.pathname === "/api/suggest") {
45
112
  const body = await readJson(req);
46
113
  const source = parseSuggestSource(body);
47
- const candidates = await suggestForSource(state.workspace, source);
114
+ const candidates = await suggestForSource(requireWorkspace(state), source);
48
115
  return writeJson(res, 200, { candidates });
49
116
  }
50
117
  if (req.method === "GET" && url.pathname === "/api/rag/status") {
51
- return writeJson(res, 200, getRagStatus(state.schemaDir));
118
+ return writeJson(res, 200, await getRagStatus(state));
52
119
  }
53
120
  if (req.method === "POST" && url.pathname === "/api/rag/index") {
54
- const result = await indexRag(state.schemaDir);
121
+ const result = await indexRag(state);
55
122
  return writeJson(res, 200, result);
56
123
  }
57
124
  if (req.method === "POST" && url.pathname === "/api/rag/query") {
58
125
  const body = await readJson(req);
59
126
  const query = parseRagQuery(body);
60
- const result = await queryRag(state.schemaDir, query);
127
+ const result = await queryRag(state, query);
61
128
  return writeJson(res, 200, result);
62
129
  }
63
130
  if (req.method === "POST" && url.pathname === "/api/ask") {
64
131
  const body = await readJson(req);
65
132
  const options = parseAskBody(body);
66
- const result = await askSampleQuestion(state.schemaDir, options);
133
+ const result = await askSampleQuestion(state, options);
134
+ return writeJson(res, 200, result);
135
+ }
136
+ if (req.method === "GET" && url.pathname === "/api/history") {
137
+ return writeJson(res, 200, readPlaygroundHistory(state.schemaDir));
138
+ }
139
+ if (req.method === "POST" && url.pathname === "/api/history") {
140
+ const body = await readJson(req);
141
+ const entry = parsePlaygroundHistoryEntry(body);
142
+ appendPlaygroundHistory(state.schemaDir, entry);
143
+ return writeJson(res, 200, { ok: true });
144
+ }
145
+ if (req.method === "DELETE" && url.pathname.startsWith("/api/history/")) {
146
+ const id = decodeURIComponent(url.pathname.slice("/api/history/".length));
147
+ const result = deletePlaygroundHistoryEntry(state.schemaDir, id);
148
+ return writeJson(res, 200, result);
149
+ }
150
+ if (req.method === "GET" && url.pathname === "/api/execute/status") {
151
+ const result = await getExecuteStatus(state.schemaDir);
152
+ return writeJson(res, 200, result);
153
+ }
154
+ if (req.method === "POST" && url.pathname === "/api/execute/install-driver") {
155
+ if (!isLoopbackRequest(req)) {
156
+ return writeJson(res, 403, { error: { message: "Driver install is only available from loopback clients." } });
157
+ }
158
+ const body = await readJson(req);
159
+ const result = await installExecuteDriver(body, state.schemaDir);
160
+ return writeJson(res, 200, result);
161
+ }
162
+ if (req.method === "POST" && url.pathname === "/api/execute") {
163
+ const body = await readJson(req);
164
+ const result = await executeQuery(body, state.schemaDir);
67
165
  return writeJson(res, 200, result);
68
166
  }
69
167
  if (req.method === "GET" && !url.pathname.startsWith("/api/")) {
@@ -72,7 +170,7 @@ export function createStudioServer(options) {
72
170
  writeJson(res, 404, { error: { message: "Not found" } });
73
171
  }
74
172
  catch (error) {
75
- const status = error instanceof StudioHttpError ? error.status : 500;
173
+ const status = error instanceof StudioHttpError || error instanceof SetupError ? error.status : 500;
76
174
  writeJson(res, status, {
77
175
  error: { message: error instanceof Error ? error.message : String(error) },
78
176
  });
@@ -83,9 +181,7 @@ export function serializeWorkspace(workspace) {
83
181
  const rt = getAskDbRuntimeConfig();
84
182
  const aiConfig = (() => {
85
183
  try {
86
- return resolveAskDbAiConfig(rt.ai.aiEnv, {
87
- modelEnvVar: rt.ai.studioModel ? "ASKDB_STUDIO_MODEL" : undefined,
88
- });
184
+ return ai.resolveAiConfig(rt.ai.aiEnv);
89
185
  }
90
186
  catch {
91
187
  // A misconfigured AI env (e.g. azure without resourceName) shouldn't crash the workspace
@@ -93,9 +189,15 @@ export function serializeWorkspace(workspace) {
93
189
  return undefined;
94
190
  }
95
191
  })();
192
+ const schema = loadSchema(workspace.schemaDir);
193
+ const schemaProvider = "provider" in schema && typeof schema.provider === "string" ? schema.provider : undefined;
194
+ const dialect = rt.nlToSql.dialect ??
195
+ (schemaProvider && isBuiltInDialectId(schemaProvider) ? schemaProvider : "postgres");
96
196
  return {
97
197
  schemaDir: workspace.schemaDir,
198
+ schemaPathRelative: toRelativeSchemaPath(workspace.schemaDir),
98
199
  schemaId: workspace.physical.schemaId,
200
+ dialect,
99
201
  warnings: workspace.warnings,
100
202
  aiConfigured: Boolean(aiConfig),
101
203
  model: aiConfig?.model ?? "gpt-4o-mini",
@@ -113,13 +215,211 @@ export function serializeWorkspace(workspace) {
113
215
  };
114
216
  }),
115
217
  concepts: workspace.concepts?.frontmatter.concepts ?? [],
218
+ tenantPolicy: schema.tenantPolicy ?? null,
116
219
  };
117
220
  }
221
+ function toRelativeSchemaPath(schemaDir) {
222
+ const rel = relative(process.cwd(), schemaDir);
223
+ if (rel === "")
224
+ return ".";
225
+ if (rel.startsWith("..") || rel.includes(`..${sep}`))
226
+ return schemaDir;
227
+ return rel.startsWith(".") ? rel : `./${rel}`;
228
+ }
229
+ // ---------------------------------------------------------------------------
230
+ // Setup wizard + server-side introspection (resync)
231
+ // ---------------------------------------------------------------------------
232
+ function requireWorkspace(state) {
233
+ if (!state.workspace) {
234
+ throw new StudioHttpError(409, "Studio is in setup mode — finish the guided setup first.");
235
+ }
236
+ return state.workspace;
237
+ }
238
+ function buildSetupStatus(state) {
239
+ const projectDir = process.cwd();
240
+ if (!state.setupReason) {
241
+ return {
242
+ needed: false,
243
+ reason: null,
244
+ projectDir,
245
+ configPath: null,
246
+ outputDir: state.schemaDir,
247
+ configLoadError: null,
248
+ };
249
+ }
250
+ const probe = probeSetupState(projectDir);
251
+ return {
252
+ needed: true,
253
+ reason: probe.hasConfig ? "no-artifact" : "no-config",
254
+ projectDir,
255
+ configPath: probe.configPath,
256
+ outputDir: probe.outputDir,
257
+ configLoadError: probe.loadError,
258
+ };
259
+ }
260
+ function handleSetupConfig(state, body) {
261
+ if (!state.setupReason) {
262
+ throw new StudioHttpError(409, "Studio is already set up — edit askdb.config.ts directly.");
263
+ }
264
+ const input = parseSetupConfigBody(body);
265
+ const result = writeSetupConfig(process.cwd(), input);
266
+ // Config now exists; the artifact usually doesn't yet.
267
+ const probe = probeSetupState(process.cwd());
268
+ if (probe.hasArtifact && probe.outputDir) {
269
+ state.schemaDir = probe.outputDir;
270
+ state.workspace = loadWorkspace(probe.outputDir);
271
+ state.setupReason = null;
272
+ }
273
+ else {
274
+ state.setupReason = "no-artifact";
275
+ }
276
+ return {
277
+ configPath: result.configPath,
278
+ envExamplePath: result.envExamplePath,
279
+ envVars: result.envVars,
280
+ installed: result.installed,
281
+ manualInstallCommand: result.manualInstallCommand,
282
+ packageJsonCreated: result.packageJsonCreated,
283
+ status: buildSetupStatus(state),
284
+ };
285
+ }
286
+ async function handleSetupIntrospect(state) {
287
+ if (!state.setupReason) {
288
+ throw new StudioHttpError(409, "Studio is already set up — use Resync schema instead.");
289
+ }
290
+ // probeSetupState re-bootstraps, picking up any .env created since server start.
291
+ const probe = probeSetupState(process.cwd());
292
+ if (!probe.hasConfig) {
293
+ throw new StudioHttpError(409, "Write askdb.config.ts first (setup step 1).");
294
+ }
295
+ if (probe.loadError || !probe.outputDir) {
296
+ throw new StudioHttpError(409, `askdb.config.ts exists but could not be loaded: ${probe.loadError ?? "unknown error"}. ` +
297
+ "Install the project dependencies shown in setup step 1, then retry.");
298
+ }
299
+ const outDir = probe.outputDir;
300
+ const run = await runStudioIntrospection({
301
+ outDir,
302
+ hasExistingArtifact: existsSync(join(outDir, "schema.json")),
303
+ });
304
+ state.schemaDir = outDir;
305
+ state.workspace = loadWorkspace(outDir);
306
+ state.setupReason = null;
307
+ return {
308
+ ok: true,
309
+ engine: run.engine,
310
+ schemaId: run.schemaId,
311
+ tables: run.tables,
312
+ warnings: run.warnings,
313
+ workspace: serializeWorkspace(state.workspace),
314
+ };
315
+ }
316
+ async function handleResync(state) {
317
+ const workspace = requireWorkspace(state);
318
+ // Refresh the runtime snapshot so config edits made while Studio runs are honored.
319
+ try {
320
+ bootstrapAskDbEnv({ cwd: process.cwd() });
321
+ }
322
+ catch {
323
+ // Keep the snapshot the server started with when the config disappeared mid-session.
324
+ }
325
+ const run = await runStudioIntrospection({
326
+ outDir: state.schemaDir,
327
+ schemaId: workspace.physical.schemaId,
328
+ hasExistingArtifact: true,
329
+ });
330
+ state.workspace = loadWorkspace(state.schemaDir);
331
+ return {
332
+ ok: true,
333
+ engine: run.engine,
334
+ schemaId: run.schemaId,
335
+ tables: run.tables,
336
+ warnings: run.warnings,
337
+ workspace: serializeWorkspace(state.workspace),
338
+ };
339
+ }
340
+ function buildIntrospectionPlanDto() {
341
+ const plan = resolveStudioIntrospectionPlan();
342
+ return plan.ok
343
+ ? { ok: true, engine: plan.engine, sourceLabel: plan.sourceLabel }
344
+ : { ok: false, engine: plan.engine, error: plan.error };
345
+ }
346
+ function parseSetupConfigBody(body) {
347
+ if (!isRecord(body)) {
348
+ throw new StudioHttpError(400, "Request body must be a JSON object.");
349
+ }
350
+ const databases = ["postgres", "mysql", "sqlite", "sqlserver", "prisma"];
351
+ const aiProviders = ["openai", "anthropic", "google", "azure", "foundry"];
352
+ const ragStores = ["file", "memory", "pgvector"];
353
+ const executeProviders = ["postgres", "mysql", "sqlite", "sqlserver"];
354
+ if (typeof body.database !== "string" || !databases.includes(body.database)) {
355
+ throw new StudioHttpError(400, `\`database\` must be one of: ${databases.join(", ")}.`);
356
+ }
357
+ if (typeof body.aiProvider !== "string" ||
358
+ !aiProviders.includes(body.aiProvider)) {
359
+ throw new StudioHttpError(400, `\`aiProvider\` must be one of: ${aiProviders.join(", ")}.`);
360
+ }
361
+ if (body.ragStore !== undefined && !ragStores.includes(body.ragStore)) {
362
+ throw new StudioHttpError(400, `\`ragStore\` must be one of: ${ragStores.join(", ")}.`);
363
+ }
364
+ if (body.studioExecuteProvider !== undefined &&
365
+ !executeProviders.includes(body.studioExecuteProvider)) {
366
+ throw new StudioHttpError(400, `\`studioExecuteProvider\` must be one of: ${executeProviders.join(", ")}.`);
367
+ }
368
+ if (body.studioExecute !== undefined && typeof body.studioExecute !== "boolean") {
369
+ throw new StudioHttpError(400, "`studioExecute` must be a boolean.");
370
+ }
371
+ const optionalString = (key) => {
372
+ const value = body[key];
373
+ if (value === undefined || value === null || value === "")
374
+ return undefined;
375
+ if (typeof value !== "string")
376
+ throw new StudioHttpError(400, `\`${key}\` must be a string.`);
377
+ return value;
378
+ };
379
+ return {
380
+ database: body.database,
381
+ aiProvider: body.aiProvider,
382
+ connectionEnv: optionalString("connectionEnv"),
383
+ sqliteFile: optionalString("sqliteFile"),
384
+ prismaSchema: optionalString("prismaSchema"),
385
+ aiKeyEnv: optionalString("aiKeyEnv"),
386
+ aiModelEnv: optionalString("aiModelEnv"),
387
+ schemaOut: optionalString("schemaOut"),
388
+ ragStore: body.ragStore,
389
+ pgvectorEnv: optionalString("pgvectorEnv"),
390
+ studioExecute: body.studioExecute,
391
+ studioExecuteProvider: body.studioExecuteProvider,
392
+ studioExecuteConnectionEnv: optionalString("studioExecuteConnectionEnv"),
393
+ studioExecuteSqliteFile: optionalString("studioExecuteSqliteFile"),
394
+ };
395
+ }
396
+ function saveConceptsDraft(state, concepts) {
397
+ saveConcepts(requireWorkspace(state), { concepts });
398
+ state.workspace = loadWorkspace(state.schemaDir);
399
+ }
400
+ function saveTenantPolicy(state, frontmatter, body) {
401
+ const filePath = join(state.schemaDir, "tenant-policy.md");
402
+ const md = writeTenantPolicyMarkdown(frontmatter, body);
403
+ writeFileSync(filePath, md, "utf8");
404
+ state.workspace = loadWorkspace(state.schemaDir);
405
+ }
406
+ function parseTenantPolicyBody(body) {
407
+ if (!isRecord(body) || !isRecord(body.frontmatter)) {
408
+ throw new StudioHttpError(400, "`frontmatter` object is required.");
409
+ }
410
+ const parsed = tenantPolicyFrontmatterSchema.safeParse(body.frontmatter);
411
+ if (!parsed.success) {
412
+ throw new StudioHttpError(400, `Invalid tenant policy: ${parsed.error.message}`);
413
+ }
414
+ const bodyText = typeof body.body === "string" ? body.body : "";
415
+ return { frontmatter: parsed.data, bodyText };
416
+ }
118
417
  function saveDraft(state, tableId, draft) {
119
- const table = state.workspace.tables.find((candidate) => candidate.physical.id === tableId);
418
+ const workspace = requireWorkspace(state);
419
+ const table = workspace.tables.find((candidate) => candidate.physical.id === tableId);
120
420
  if (!table)
121
421
  throw new StudioHttpError(404, `No such table: ${tableId}`);
122
- const frontmatter = buildFrontmatter(table.physical, state.workspace.physical.schemaId, draft);
422
+ const frontmatter = buildFrontmatter(table.physical, workspace.physical.schemaId, draft);
123
423
  let body = table.parsed
124
424
  ? replaceTableDescription(table.parsed.body, draft.description)
125
425
  : buildDefaultTableBody(table.physical.name, draft.description);
@@ -129,43 +429,123 @@ function saveDraft(state, tableId, draft) {
129
429
  if (draft.exampleQuestions !== undefined) {
130
430
  body = replaceH2Section(body, "Example questions", draft.exampleQuestions);
131
431
  }
132
- saveTable(state.workspace, tableId, frontmatter, body);
432
+ saveTable(workspace, tableId, frontmatter, body);
133
433
  state.workspace = loadWorkspace(state.schemaDir);
134
434
  }
135
435
  async function suggestForSource(workspace, source) {
136
436
  const rt = getAskDbRuntimeConfig();
137
- const modelEnvVar = rt.ai.studioModel
138
- ? "ASKDB_STUDIO_MODEL"
139
- : rt.ai.tuiModel
140
- ? "ASKDB_TUI_MODEL"
141
- : undefined;
142
- const model = await createAskDbLanguageModelFromEnv(rt.ai.aiEnv, {
143
- modelEnvVar,
144
- });
437
+ const model = await ai.createLanguageModelFromEnv(rt.ai.aiEnv);
145
438
  if (!model) {
146
- throw new StudioHttpError(400, askDbAiKeyMissingMessage("AI enrichment suggestions"));
439
+ throw new StudioHttpError(400, ai.keyMissingMessage("AI enrichment suggestions"));
147
440
  }
148
441
  const candidates = await suggestEnrichment(buildSuggestionTarget(workspace, source), buildSuggestionContext(workspace, source.tableId), model);
149
442
  return candidates.map((candidate) => ({ text: candidate.text }));
150
443
  }
151
- async function askSampleQuestion(schemaDir, options) {
444
+ async function suggestTenantPolicyDraft(state) {
445
+ const rt = getAskDbRuntimeConfig();
446
+ const model = await ai.createLanguageModelFromEnv(rt.ai.aiEnv);
447
+ if (!model) {
448
+ throw new StudioHttpError(400, ai.keyMissingMessage("AI tenant policy suggestion"));
449
+ }
450
+ // Build a compact DDL-like representation of the schema for the prompt
451
+ const workspace = requireWorkspace(state);
452
+ const schemaId = workspace.physical.schemaId;
453
+ const schemaDdl = workspace.tables
454
+ .map((t) => {
455
+ const cols = t.physical.columns
456
+ .map((c) => {
457
+ const parts = [c.name, c.type];
458
+ if (c.primaryKey)
459
+ parts.push("PRIMARY KEY");
460
+ if (!c.nullable)
461
+ parts.push("NOT NULL");
462
+ return ` ${parts.join(" ")}`;
463
+ })
464
+ .join(",\n");
465
+ const rels = (t.physical.relationships ?? [])
466
+ .map((r) => ` -- FK: ${r.from} -> ${r.to}`)
467
+ .join("\n");
468
+ return `CREATE TABLE ${t.physical.schema}.${t.physical.name} (\n${cols}\n);\n${rels}`.trim();
469
+ })
470
+ .join("\n\n");
471
+ const tableIds = workspace.tables.map((t) => t.physical.id);
472
+ const systemPrompt = `You are an expert database architect specializing in multi-tenant SaaS systems.
473
+ You analyze database schemas and produce a tenant-policy configuration in a specific JSON format.
474
+
475
+ The tenant policy identifies:
476
+ 1. **roots** — Tables whose rows represent tenants (organizations, companies, accounts, etc).
477
+ Each root has: id (table ID like "table:public.orgs"), tenantIdColumn (column ref like "table:public.orgs#id"), label (human-readable name).
478
+ If a root has a parent root (e.g. sub_org belongs to org), include parent: { root: "<parent-root-id>", foreignKey: "<fk-column-ref>" }.
479
+ 2. **hierarchy** — Edges showing parent→child relationships between roots or tenant-scoped tables.
480
+ Each edge: { parent: "<table-id>", child: "<table-id>", foreignKey: "<child-fk-column-ref>" }.
481
+ 3. **scopedTables** — Tables that belong to a tenant via a direct column or a JOIN chain.
482
+ Each entry: { id: "<table-id>", scopeThrough: [{ root: "<root-id>", column: "<column-ref>" }] } for direct FK,
483
+ or { id: "<table-id>", scopeThrough: [{ root: "<root-id>", join: [{ from: "<col-ref>", to: "<col-ref>" }] }] } for indirect.
484
+ 4. **polymorphicTables** — Tables with a type/id polymorphic association pattern.
485
+ Each entry: { id: "<table-id>", typeColumn: "<col-ref>", idColumn: "<col-ref>", mapping: { "<type_value>": "<target-table-id>" } }.
486
+ 5. **globalTables** — Tables shared across all tenants (lookup tables, reference data). Array of table IDs.
487
+ 6. **enforcement** — "strict" (reject queries on unknown tables) or "warn" (allow with warnings).
488
+
489
+ Column references use the format "table:<schema>.<table>#<column>".
490
+ Table IDs use the format "table:<schema>.<table>".
491
+
492
+ Available table IDs: ${JSON.stringify(tableIds)}
493
+
494
+ Return a JSON object with two keys:
495
+ - "frontmatter": the TenantPolicyFrontmatter object (schemaId, enforcement, roots, hierarchy?, scopedTables?, polymorphicTables?, globalTables?)
496
+ - "body": a markdown body string with sections "## Hierarchy", "## Scope rules", "## Sensitive interactions" explaining the policy in plain English.
497
+
498
+ Return ONLY valid JSON, no markdown fences or explanations.`;
499
+ const userPrompt = `Analyze this database schema and draft a complete multi-tenant policy.
500
+
501
+ Schema ID: ${schemaId}
502
+
503
+ ${schemaDdl}
504
+
505
+ Look for:
506
+ - Tables that represent organizations, companies, accounts, tenants, or workspaces — these are roots.
507
+ - Foreign key columns like org_id, company_id, tenant_id, account_id that scope data to a tenant.
508
+ - Tables with no tenant FK that are clearly shared/global (lookup tables, enums, config).
509
+ - Polymorphic association patterns (type_column + id_column pointing to multiple parent tables).
510
+ - Multi-level hierarchies (org → team → user, company → department → employee).
511
+
512
+ If the schema has no clear multi-tenant pattern, still make your best guess at what would be the tenant root.`;
513
+ const result = await defaultGenerateText({
514
+ model,
515
+ system: systemPrompt,
516
+ prompt: userPrompt,
517
+ });
518
+ let parsed;
519
+ try {
520
+ parsed = JSON.parse(result.text);
521
+ }
522
+ catch {
523
+ throw new StudioHttpError(500, `AI returned invalid JSON for tenant policy suggestion. Raw text: ${result.text.slice(0, 500)}`);
524
+ }
525
+ // Validate frontmatter shape
526
+ const validation = tenantPolicyFrontmatterSchema.safeParse(parsed.frontmatter);
527
+ if (!validation.success) {
528
+ throw new StudioHttpError(500, `AI-generated tenant policy frontmatter is invalid: ${validation.error.issues.map((i) => i.message).join("; ")}`);
529
+ }
530
+ parsed.frontmatter = validation.data;
531
+ return parsed;
532
+ }
533
+ async function askSampleQuestion(state, options) {
152
534
  const rt = getAskDbRuntimeConfig();
153
535
  const mockSql = rt.dev.mockSql;
154
- const aiConfig = mockSql ? undefined : resolveAskDbAiConfig(rt.ai.aiEnv);
536
+ const aiConfig = mockSql ? undefined : ai.resolveAiConfig(rt.ai.aiEnv);
155
537
  if (!mockSql && !aiConfig) {
156
- throw new StudioHttpError(400, `${askDbAiKeyMissingMessage("Sample NL-to-SQL generation")} Set ASKDB_MOCK_SQL to bypass the live model.`);
538
+ throw new StudioHttpError(400, `${ai.keyMissingMessage("Sample NL-to-SQL generation")} Set ASKDB_MOCK_SQL to bypass the live model.`);
157
539
  }
158
- const schema = loadSchema(schemaDir);
540
+ const schema = loadSchema(state.schemaDir);
159
541
  const retrievedChunks = [];
160
542
  const usage = createRequestUsageCollector();
161
543
  const ragIndex = options.useRag
162
- ? await createCurrentStudioRagIndex(schemaDir, "using RAG for sample generation", usage)
544
+ ? await createCurrentStudioRagIndex(state, "using RAG for sample generation", usage)
163
545
  : undefined;
164
546
  const model = mockSql
165
547
  ? undefined
166
- : (await createAskDbLanguageModelFromEnv(rt.ai.aiEnv, {
167
- modelEnvVar: rt.ai.studioModel ? "ASKDB_STUDIO_MODEL" : undefined,
168
- }));
548
+ : (await ai.createLanguageModelFromEnv(rt.ai.aiEnv));
169
549
  const schemaProvider = "provider" in schema && typeof schema.provider === "string"
170
550
  ? schema.provider
171
551
  : undefined;
@@ -177,6 +557,8 @@ async function askSampleQuestion(schemaDir, options) {
177
557
  model,
178
558
  dialect,
179
559
  explain: true,
560
+ ...(options.tenantScope ? { tenantScope: options.tenantScope } : {}),
561
+ ...(options.tenantSqlMode ? { tenantSqlMode: options.tenantSqlMode } : {}),
180
562
  ...(ragIndex
181
563
  ? {
182
564
  retriever: async (params) => {
@@ -185,7 +567,7 @@ async function askSampleQuestion(schemaDir, options) {
185
567
  results = await ragIndex.retriever(params);
186
568
  }
187
569
  catch (error) {
188
- throw formatStudioRagEmbeddingError(error, ragIndex.config);
570
+ throw formatStudioRagOperationError(error, ragIndex.config);
189
571
  }
190
572
  retrievedChunks.splice(0, retrievedChunks.length, ...results);
191
573
  return results;
@@ -202,6 +584,8 @@ async function askSampleQuestion(schemaDir, options) {
202
584
  : {
203
585
  generateText: createTrackedGenerateText(usage),
204
586
  },
587
+ }).finally(async () => {
588
+ await ragIndex?.dispose();
205
589
  });
206
590
  return {
207
591
  sql: result.sql,
@@ -211,90 +595,115 @@ async function askSampleQuestion(schemaDir, options) {
211
595
  enabled: options.useRag,
212
596
  chunks: options.useRag ? retrievedChunks.map(serializeRagResult) : [],
213
597
  },
598
+ tenant: options.tenantScope
599
+ ? {
600
+ enabled: true,
601
+ sqlMode: options.tenantSqlMode ?? "sql-only",
602
+ bindings: result.tenantBindings ?? [],
603
+ params: result.tenantParams ?? [],
604
+ }
605
+ : null,
214
606
  usage: usage.toDto(),
215
607
  };
216
608
  }
217
- function getRagStatus(schemaDir) {
609
+ async function getRagStatus(state) {
218
610
  const config = resolveStudioRagEmbedderConfig();
219
- const sources = loadChunkerSourcesFromDir(schemaDir);
611
+ const sources = loadChunkerSourcesFromDir(state.schemaDir);
220
612
  const chunkResult = chunkSchema(sources);
221
- const lockPath = join(schemaDir, "schema.lock.json");
222
- const embeddingsJsonPath = join(schemaDir, "schema.embeddings.json");
223
- const embeddingsBinPath = join(schemaDir, "schema.embeddings.bin");
613
+ const lockPath = join(state.schemaDir, "schema.lock.json");
224
614
  const lock = readOptionalJson(lockPath);
225
- const embeddings = readOptionalJson(embeddingsJsonPath);
226
615
  const currentHashes = Object.fromEntries(chunkResult.chunks.map((chunk) => [chunk.id, chunkContentHash(chunk.text)]));
227
616
  const lockHashes = lock?.hashes ?? {};
228
617
  const hashIds = Object.keys(currentHashes);
229
- const stale = !lock ||
230
- !existsSync(embeddingsJsonPath) ||
231
- !existsSync(embeddingsBinPath) ||
232
- lock.embedderId !== config.embedderId ||
233
- embeddings?.dimensions !== config.dimensions ||
234
- Object.keys(lockHashes).length !== hashIds.length ||
235
- hashIds.some((id) => lockHashes[id] !== currentHashes[id]);
236
- return {
237
- schemaId: sources.schema.schemaId,
238
- embedder: {
239
- kind: config.kind,
240
- label: config.label,
241
- configured: config.configured,
242
- expectedId: config.embedderId,
243
- indexedId: lock?.embedderId ?? null,
244
- provider: config.kind === "ai-sdk" ? config.provider : null,
245
- model: config.kind === "ai-sdk" ? config.model : null,
246
- baseUrl: config.kind === "ai-sdk" ? config.baseUrl ?? null : null,
247
- },
248
- embedderId: lock?.embedderId ?? config.embedderId,
249
- expectedEmbedderId: config.embedderId,
250
- hasIndex: Boolean(lock && existsSync(embeddingsJsonPath) && existsSync(embeddingsBinPath)),
251
- stale,
252
- updatedAt: lock?.updatedAt ?? null,
253
- chunksTotal: chunkResult.chunks.length,
254
- chunksIndexed: Array.isArray(embeddings?.records) ? embeddings.records.length : 0,
255
- dimensions: typeof embeddings?.dimensions === "number" ? embeddings.dimensions : config.dimensions,
256
- expectedDimensions: config.dimensions,
257
- sensitiveExcluded: chunkResult.stats.sensitiveExcluded,
258
- sensitiveIncluded: chunkResult.stats.sensitiveIncluded,
259
- files: {
260
- lock: existsSync(lockPath),
261
- embeddingsJson: existsSync(embeddingsJsonPath),
262
- embeddingsBin: existsSync(embeddingsBinPath),
263
- },
264
- };
618
+ const store = await openStudioRagStore(state, config.dimensions);
619
+ try {
620
+ const chunksIndexed = await countStudioRagStoreChunks(store, sources.schema.schemaId);
621
+ const stale = !lock ||
622
+ lock.embedderId !== config.embedderId ||
623
+ Object.keys(lockHashes).length !== hashIds.length ||
624
+ hashIds.some((id) => lockHashes[id] !== currentHashes[id]) ||
625
+ chunksIndexed !== hashIds.length;
626
+ const fileArtifacts = store.kind === "file" && store.basePath
627
+ ? {
628
+ lock: existsSync(lockPath),
629
+ embeddingsJson: existsSync(`${store.basePath}.embeddings.json`),
630
+ embeddingsBin: existsSync(`${store.basePath}.embeddings.bin`),
631
+ }
632
+ : {
633
+ lock: existsSync(lockPath),
634
+ embeddingsJson: false,
635
+ embeddingsBin: false,
636
+ };
637
+ return {
638
+ schemaId: sources.schema.schemaId,
639
+ store: {
640
+ kind: store.kind,
641
+ basePath: store.basePath ?? null,
642
+ table: store.table ?? null,
643
+ indexStrategy: store.indexStrategy ?? null,
644
+ },
645
+ embedder: {
646
+ kind: config.kind,
647
+ label: config.label,
648
+ configured: config.configured,
649
+ expectedId: config.embedderId,
650
+ indexedId: lock?.embedderId ?? null,
651
+ provider: config.kind === "ai-sdk" ? config.provider : null,
652
+ model: config.kind === "ai-sdk" ? config.model : null,
653
+ baseUrl: config.kind === "ai-sdk" ? config.baseUrl ?? null : null,
654
+ },
655
+ embedderId: lock?.embedderId ?? config.embedderId,
656
+ expectedEmbedderId: config.embedderId,
657
+ hasIndex: Boolean(lock && chunksIndexed > 0),
658
+ stale,
659
+ updatedAt: lock?.updatedAt ?? null,
660
+ chunksTotal: chunkResult.chunks.length,
661
+ chunksIndexed,
662
+ dimensions: config.dimensions,
663
+ expectedDimensions: config.dimensions,
664
+ sensitiveExcluded: chunkResult.stats.sensitiveExcluded,
665
+ sensitiveIncluded: chunkResult.stats.sensitiveIncluded,
666
+ files: fileArtifacts,
667
+ };
668
+ }
669
+ finally {
670
+ await store.dispose();
671
+ }
265
672
  }
266
- async function indexRag(schemaDir) {
673
+ async function indexRag(state) {
267
674
  const config = resolveStudioRagEmbedderConfig();
268
675
  if (!config.configured) {
269
676
  throw new StudioHttpError(400, studioRagAiSdkKeyMissingMessage());
270
677
  }
271
678
  const usage = createRequestUsageCollector();
272
- clearIncompatibleRagStore(schemaDir, config);
273
- const sources = loadChunkerSourcesFromDir(schemaDir);
274
- const store = createFileStore({ basePath: join(schemaDir, "schema") });
679
+ clearIncompatibleRagStore(state, config);
680
+ const sources = loadChunkerSourcesFromDir(state.schemaDir);
681
+ const store = await openStudioRagStore(state, config.dimensions);
275
682
  let result;
276
683
  try {
277
684
  result = await buildSchemaIndex({
278
685
  schema: sources,
279
686
  embedder: await createStudioRagEmbedder(config, usage),
280
- store,
687
+ store: store.store,
281
688
  embedderId: config.embedderId,
282
- lockFilePath: join(schemaDir, "schema.lock.json"),
689
+ lockFilePath: join(state.schemaDir, "schema.lock.json"),
283
690
  });
284
691
  }
285
692
  catch (error) {
286
- throw formatStudioRagEmbeddingError(error, config);
693
+ throw formatStudioRagOperationError(error, config);
694
+ }
695
+ finally {
696
+ await store.dispose();
287
697
  }
288
- store.flush();
289
698
  return {
290
- status: getRagStatus(schemaDir),
699
+ status: await getRagStatus(state),
291
700
  stats: result.stats,
292
701
  usage: usage.toDto(),
293
702
  };
294
703
  }
295
- async function queryRag(schemaDir, query) {
704
+ async function queryRag(state, query) {
296
705
  const usage = createRequestUsageCollector();
297
- const ragIndex = await createCurrentStudioRagIndex(schemaDir, "querying chunks", usage);
706
+ const ragIndex = await createCurrentStudioRagIndex(state, "querying chunks", usage);
298
707
  let results;
299
708
  try {
300
709
  results = await ragIndex.retriever({
@@ -307,7 +716,10 @@ async function queryRag(schemaDir, query) {
307
716
  });
308
717
  }
309
718
  catch (error) {
310
- throw formatStudioRagEmbeddingError(error, ragIndex.config);
719
+ throw formatStudioRagOperationError(error, ragIndex.config);
720
+ }
721
+ finally {
722
+ await ragIndex.dispose();
311
723
  }
312
724
  return {
313
725
  question: query.question,
@@ -316,12 +728,12 @@ async function queryRag(schemaDir, query) {
316
728
  usage: usage.toDto(),
317
729
  };
318
730
  }
319
- async function createCurrentStudioRagIndex(schemaDir, action, usage) {
731
+ async function createCurrentStudioRagIndex(state, action, usage) {
320
732
  const config = resolveStudioRagEmbedderConfig();
321
733
  if (!config.configured) {
322
734
  throw new StudioHttpError(400, studioRagAiSdkKeyMissingMessage());
323
735
  }
324
- const status = getRagStatus(schemaDir);
736
+ const status = (await getRagStatus(state));
325
737
  if (!status.hasIndex) {
326
738
  throw new StudioHttpError(400, `Build the RAG index before ${action}.`);
327
739
  }
@@ -331,6 +743,7 @@ async function createCurrentStudioRagIndex(schemaDir, action, usage) {
331
743
  if (!status.schemaId || typeof status.chunksTotal !== "number") {
332
744
  throw new StudioHttpError(500, "Studio RAG status is missing schema metadata.");
333
745
  }
746
+ const store = await openStudioRagStore(state, config.dimensions);
334
747
  return {
335
748
  config,
336
749
  status: {
@@ -339,10 +752,85 @@ async function createCurrentStudioRagIndex(schemaDir, action, usage) {
339
752
  },
340
753
  retriever: createRetriever({
341
754
  embedder: await createStudioRagEmbedder(config, usage),
342
- store: createFileStore({ basePath: join(schemaDir, "schema") }),
755
+ store: store.store,
343
756
  }),
757
+ dispose: async () => {
758
+ await store.dispose();
759
+ },
760
+ };
761
+ }
762
+ function resolveStudioRagStoreConfig(state) {
763
+ const rt = getAskDbRuntimeConfig();
764
+ const kind = rt.structured.rag.store;
765
+ if (kind === "memory")
766
+ return { kind };
767
+ if (kind === "file") {
768
+ const basePath = rt.structured.rag.storeConfig.file?.basePath?.trim();
769
+ return { kind, basePath: basePath ? resolve(basePath) : join(state.schemaDir, "schema") };
770
+ }
771
+ const connectionString = pickFlat(rt.flat, "ASKDB_PGVECTOR_URL");
772
+ return {
773
+ kind,
774
+ connectionString,
775
+ table: rt.structured.rag.storeConfig.pgvector?.table?.trim() || undefined,
776
+ indexStrategy: pickFlat(rt.flat, "ASKDB_PGVECTOR_INDEX_STRATEGY"),
777
+ };
778
+ }
779
+ async function openStudioRagStore(state, dimensions) {
780
+ const config = resolveStudioRagStoreConfig(state);
781
+ if (config.kind === "memory") {
782
+ state.ragMemoryStore ??= createMemoryStore();
783
+ return {
784
+ kind: "memory",
785
+ store: state.ragMemoryStore,
786
+ dispose: async () => { },
787
+ };
788
+ }
789
+ if (config.kind === "file") {
790
+ const store = createFileStore({ basePath: config.basePath });
791
+ return {
792
+ kind: "file",
793
+ store,
794
+ basePath: config.basePath,
795
+ dispose: async () => {
796
+ store.flush();
797
+ },
798
+ };
799
+ }
800
+ if (!config.connectionString) {
801
+ throw new StudioHttpError(400, 'Studio pgvector RAG requires `ASKDB_PGVECTOR_URL` via `askdb.config.ts`.');
802
+ }
803
+ const pgvectorFactory = studioPgvectorStoreFactoryForTests ?? createPgvectorStore;
804
+ const store = pgvectorFactory({
805
+ connectionString: config.connectionString,
806
+ table: config.table,
807
+ dimensions,
808
+ ...(config.indexStrategy ? { indexStrategy: config.indexStrategy } : {}),
809
+ });
810
+ await store.ensureSchema();
811
+ return {
812
+ kind: "pgvector",
813
+ store,
814
+ table: config.table,
815
+ indexStrategy: config.indexStrategy,
816
+ dispose: async () => {
817
+ await store.close();
818
+ },
344
819
  };
345
820
  }
821
+ async function countStudioRagStoreChunks(store, schemaId) {
822
+ if (typeof store.store.count === "function") {
823
+ return store.store.count({ schemaId });
824
+ }
825
+ if (typeof store.store.size === "function") {
826
+ return store.store.size();
827
+ }
828
+ return 0;
829
+ }
830
+ function pickFlat(flat, key) {
831
+ const value = flat[key];
832
+ return typeof value === "string" && value.trim() !== "" ? value.trim() : undefined;
833
+ }
346
834
  function pickEnv(env, key) {
347
835
  const v = env[key];
348
836
  return typeof v === "string" && v.trim() !== "" ? v.trim() : undefined;
@@ -350,7 +838,7 @@ function pickEnv(env, key) {
350
838
  function resolveStudioRagEmbedderConfig() {
351
839
  const rt = getAskDbRuntimeConfig();
352
840
  const base = rt.ai.aiEnv;
353
- const explicitKind = pickEnv(base, "ASKDB_RAG_EMBEDDER") ?? pickEnv(base, "ASKDB_STUDIO_RAG_EMBEDDER");
841
+ const explicitKind = pickEnv(base, "ASKDB_RAG_EMBEDDER");
354
842
  const kind = explicitKind?.toLowerCase();
355
843
  if (kind === "mock") {
356
844
  return {
@@ -365,8 +853,8 @@ function resolveStudioRagEmbedderConfig() {
365
853
  throw new StudioHttpError(400, `Unsupported Studio RAG embedder: ${kind}`);
366
854
  }
367
855
  const env = buildStudioRagEmbeddingEnv(kind, base);
368
- const aiConfig = resolveAskDbEmbeddingConfig(env, {
369
- modelEnvVar: "ASKDB_STUDIO_RAG_EMBEDDER_MODEL",
856
+ const aiConfig = ai.resolveEmbeddingConfig(env, {
857
+ modelEnvVar: "ASKDB_RAG_EMBEDDER_MODEL",
370
858
  modelDefault: DEFAULT_EMBEDDING_MODEL,
371
859
  });
372
860
  if (!aiConfig && kind === undefined) {
@@ -380,8 +868,7 @@ function resolveStudioRagEmbedderConfig() {
380
868
  }
381
869
  const provider = aiConfig?.provider ?? fallbackStudioRagProvider(kind, base);
382
870
  const model = aiConfig?.model ?? DEFAULT_EMBEDDING_MODEL;
383
- const dimensionOverride = readPositiveIntegerEnv(pickEnv(base, "ASKDB_RAG_EMBEDDER_DIMENSIONS") ??
384
- pickEnv(base, "ASKDB_STUDIO_RAG_EMBEDDER_DIMENSIONS"));
871
+ const dimensionOverride = readPositiveIntegerEnv(pickEnv(base, "ASKDB_RAG_EMBEDDER_DIMENSIONS"));
385
872
  const dimensions = dimensionOverride ?? defaultEmbeddingDimensions(model);
386
873
  return {
387
874
  kind: "ai-sdk",
@@ -397,14 +884,12 @@ function resolveStudioRagEmbedderConfig() {
397
884
  };
398
885
  }
399
886
  function buildStudioRagEmbeddingEnv(kind, base) {
400
- const apiKeyOverride = pickEnv(base, "ASKDB_RAG_EMBEDDER_API_KEY") ?? pickEnv(base, "ASKDB_STUDIO_RAG_EMBEDDER_API_KEY");
401
- const modelOverride = pickEnv(base, "ASKDB_RAG_EMBEDDER_MODEL") ?? pickEnv(base, "ASKDB_STUDIO_RAG_EMBEDDER_MODEL");
402
- const baseUrlOverride = pickEnv(base, "ASKDB_RAG_EMBEDDER_BASE_URL") ?? pickEnv(base, "ASKDB_STUDIO_RAG_EMBEDDER_BASE_URL");
887
+ const apiKeyOverride = pickEnv(base, "ASKDB_RAG_EMBEDDER_API_KEY");
888
+ const baseUrlOverride = pickEnv(base, "ASKDB_RAG_EMBEDDER_BASE_URL");
403
889
  return {
404
890
  ...base,
405
891
  ...(kind === "openai" ? { ASKDB_AI_PROVIDER: "openai" } : {}),
406
892
  ...(apiKeyOverride ? { ASKDB_AI_API_KEY: apiKeyOverride } : {}),
407
- ...(modelOverride ? { ASKDB_STUDIO_RAG_EMBEDDER_MODEL: modelOverride } : {}),
408
893
  ...(baseUrlOverride ? { ASKDB_AI_BASE_URL: baseUrlOverride } : {}),
409
894
  };
410
895
  }
@@ -426,7 +911,7 @@ async function createStudioRagEmbedder(config, usage) {
426
911
  if (!config.aiConfig) {
427
912
  throw new StudioHttpError(400, studioRagAiSdkKeyMissingMessage());
428
913
  }
429
- const model = await createAskDbEmbeddingModel(config.aiConfig, {
914
+ const model = await ai.createEmbeddingModel(config.aiConfig, {
430
915
  dimensions: config.requestDimensions,
431
916
  });
432
917
  return createAiSdkEmbedder({
@@ -441,11 +926,16 @@ async function createStudioRagEmbedder(config, usage) {
441
926
  },
442
927
  });
443
928
  }
444
- function formatStudioRagEmbeddingError(error, config) {
929
+ function formatStudioRagOperationError(error, config) {
930
+ if (error instanceof StudioHttpError)
931
+ return error;
445
932
  if (config.kind === "mock") {
446
933
  return new StudioHttpError(500, error instanceof Error ? error.message : String(error));
447
934
  }
448
935
  const apiError = findApiCallError(error);
936
+ if (!apiError) {
937
+ return new StudioHttpError(500, error instanceof Error ? error.message : String(error));
938
+ }
449
939
  const parts = [
450
940
  `Studio RAG embedding request failed for provider ${config.provider}, model ${config.model}.`,
451
941
  ];
@@ -488,15 +978,18 @@ function truncateForMessage(value) {
488
978
  return undefined;
489
979
  return trimmed.length > 500 ? `${trimmed.slice(0, 500)}...` : trimmed;
490
980
  }
491
- function clearIncompatibleRagStore(schemaDir, config) {
492
- const embeddingsJsonPath = join(schemaDir, "schema.embeddings.json");
981
+ function clearIncompatibleRagStore(state, config) {
982
+ const store = resolveStudioRagStoreConfig(state);
983
+ if (store.kind !== "file")
984
+ return;
985
+ const embeddingsJsonPath = `${store.basePath}.embeddings.json`;
493
986
  const embeddings = readOptionalJson(embeddingsJsonPath);
494
987
  if (embeddings?.dimensions === undefined || embeddings.dimensions === config.dimensions)
495
988
  return;
496
989
  for (const path of [
497
- join(schemaDir, "schema.embeddings.json"),
498
- join(schemaDir, "schema.embeddings.bin"),
499
- join(schemaDir, "schema.lock.json"),
990
+ `${store.basePath}.embeddings.json`,
991
+ `${store.basePath}.embeddings.bin`,
992
+ join(state.schemaDir, "schema.lock.json"),
500
993
  ]) {
501
994
  rmSync(path, { force: true });
502
995
  }
@@ -546,6 +1039,234 @@ function stableTokenHash(token) {
546
1039
  }
547
1040
  return hash >>> 0;
548
1041
  }
1042
+ // ---------------------------------------------------------------------------
1043
+ // Playground history helpers
1044
+ // ---------------------------------------------------------------------------
1045
+ const HISTORY_MAX_STORED = 200;
1046
+ const HISTORY_MAX_RETURNED = 50;
1047
+ function playgroundHistoryPath(schemaDir) {
1048
+ return join(schemaDir, "playground-history.json");
1049
+ }
1050
+ function readPlaygroundHistory(schemaDir) {
1051
+ const raw = readOptionalJson(playgroundHistoryPath(schemaDir));
1052
+ const entries = Array.isArray(raw) ? raw : [];
1053
+ return { entries: entries.slice(0, HISTORY_MAX_RETURNED) };
1054
+ }
1055
+ function appendPlaygroundHistory(schemaDir, entry) {
1056
+ const histPath = playgroundHistoryPath(schemaDir);
1057
+ const raw = readOptionalJson(histPath);
1058
+ const existing = Array.isArray(raw)
1059
+ ? raw
1060
+ : [];
1061
+ const newEntry = {
1062
+ ...entry,
1063
+ id: randomUUID(),
1064
+ timestamp: new Date().toISOString(),
1065
+ };
1066
+ const updated = [newEntry, ...existing].slice(0, HISTORY_MAX_STORED);
1067
+ writeFileSync(histPath, JSON.stringify(updated, null, 2), "utf8");
1068
+ }
1069
+ function deletePlaygroundHistoryEntry(schemaDir, id) {
1070
+ const histPath = playgroundHistoryPath(schemaDir);
1071
+ const raw = readOptionalJson(histPath);
1072
+ const existing = Array.isArray(raw)
1073
+ ? raw
1074
+ : [];
1075
+ const filtered = existing.filter((e) => e.id !== id);
1076
+ if (filtered.length === existing.length) {
1077
+ return { ok: false, error: "not found" };
1078
+ }
1079
+ writeFileSync(histPath, JSON.stringify(filtered, null, 2), "utf8");
1080
+ return { ok: true };
1081
+ }
1082
+ function parsePlaygroundHistoryEntry(body) {
1083
+ if (!isRecord(body)) {
1084
+ throw new StudioHttpError(400, "Request body must be a JSON object.");
1085
+ }
1086
+ if (typeof body.question !== "string" || body.question.trim() === "") {
1087
+ throw new StudioHttpError(400, "`question` is required.");
1088
+ }
1089
+ if (body.mode !== "full" && body.mode !== "rag") {
1090
+ throw new StudioHttpError(400, "`mode` must be `full` or `rag`.");
1091
+ }
1092
+ if (typeof body.sqlMode !== "string") {
1093
+ throw new StudioHttpError(400, "`sqlMode` is required.");
1094
+ }
1095
+ if (typeof body.sql !== "string") {
1096
+ throw new StudioHttpError(400, "`sql` is required.");
1097
+ }
1098
+ return body;
1099
+ }
1100
+ // ---------------------------------------------------------------------------
1101
+ // Execute endpoint helpers
1102
+ // ---------------------------------------------------------------------------
1103
+ async function getExecuteStatus(schemaDir) {
1104
+ const rt = getAskDbRuntimeConfig();
1105
+ const { provider, databaseUrl, file } = rt.studio.execute;
1106
+ const def = EXECUTE_DRIVER_REGISTRY[provider];
1107
+ const projectRoot = findProjectRoot(schemaDir) ?? schemaDir;
1108
+ const installed = isDriverInstalled(def.packageName, projectRoot);
1109
+ const connectionKind = provider === "sqlite" ? "file" : "url";
1110
+ const configured = connectionKind === "file" ? Boolean(file) : Boolean(databaseUrl);
1111
+ return {
1112
+ provider,
1113
+ label: def.label,
1114
+ configured,
1115
+ connectionKind,
1116
+ packageName: def.packageName,
1117
+ installed,
1118
+ installCommand: def.installCommand,
1119
+ canInstallFromStudio: true,
1120
+ manualInstallReason: null,
1121
+ };
1122
+ }
1123
+ async function installExecuteDriver(body, schemaDir) {
1124
+ const rt = getAskDbRuntimeConfig();
1125
+ const configuredProvider = rt.studio.execute.provider;
1126
+ // Accept the configured provider or an explicit allowlisted provider.
1127
+ let provider = configuredProvider;
1128
+ if (isRecord(body) && typeof body.provider === "string") {
1129
+ const requested = body.provider;
1130
+ if (!(requested in EXECUTE_DRIVER_REGISTRY)) {
1131
+ return {
1132
+ ok: false,
1133
+ provider: configuredProvider,
1134
+ packageName: EXECUTE_DRIVER_REGISTRY[configuredProvider].packageName,
1135
+ command: [],
1136
+ error: `Unknown provider: ${body.provider}`,
1137
+ installed: false,
1138
+ };
1139
+ }
1140
+ provider = requested;
1141
+ }
1142
+ const def = EXECUTE_DRIVER_REGISTRY[provider];
1143
+ const { packageManager, command, args, manualReason } = detectPackageManager(schemaDir, def.packageName);
1144
+ if (manualReason) {
1145
+ return {
1146
+ ok: false,
1147
+ provider,
1148
+ packageName: def.packageName,
1149
+ command: [],
1150
+ error: manualReason,
1151
+ installed: false,
1152
+ };
1153
+ }
1154
+ // Run the install.
1155
+ const { stdout, stderr, code } = await spawnCommand(command, args, schemaDir);
1156
+ const projectRoot = findProjectRoot(schemaDir) ?? schemaDir;
1157
+ const installed = isDriverInstalled(def.packageName, projectRoot);
1158
+ return {
1159
+ ok: code === 0,
1160
+ provider,
1161
+ packageName: def.packageName,
1162
+ command: [command, ...args],
1163
+ stdout,
1164
+ stderr,
1165
+ error: code !== 0 ? `${packageManager} exited with code ${code}` : undefined,
1166
+ installed,
1167
+ };
1168
+ }
1169
+ function detectPackageManager(cwd, packageName) {
1170
+ // Walk up from schemaDir to find the project root's lockfile.
1171
+ const projectRoot = findProjectRoot(cwd);
1172
+ if (!projectRoot) {
1173
+ return {
1174
+ packageManager: "",
1175
+ command: "",
1176
+ args: [],
1177
+ manualReason: `Could not locate a package.json above ${cwd}. Run the install manually.`,
1178
+ };
1179
+ }
1180
+ if (existsSync(join(projectRoot, "pnpm-lock.yaml"))) {
1181
+ return {
1182
+ packageManager: "pnpm",
1183
+ command: "pnpm",
1184
+ args: ["add", packageName],
1185
+ manualReason: null,
1186
+ };
1187
+ }
1188
+ if (existsSync(join(projectRoot, "package-lock.json"))) {
1189
+ return {
1190
+ packageManager: "npm",
1191
+ command: "npm",
1192
+ args: ["install", packageName],
1193
+ manualReason: null,
1194
+ };
1195
+ }
1196
+ if (existsSync(join(projectRoot, "yarn.lock"))) {
1197
+ return {
1198
+ packageManager: "yarn",
1199
+ command: "yarn",
1200
+ args: ["add", packageName],
1201
+ manualReason: null,
1202
+ };
1203
+ }
1204
+ if (existsSync(join(projectRoot, "bun.lockb")) || existsSync(join(projectRoot, "bun.lock"))) {
1205
+ return {
1206
+ packageManager: "bun",
1207
+ command: "bun",
1208
+ args: ["add", packageName],
1209
+ manualReason: null,
1210
+ };
1211
+ }
1212
+ return {
1213
+ packageManager: "",
1214
+ command: "",
1215
+ args: [],
1216
+ manualReason: `No recognized lockfile found in ${projectRoot}. Run: pnpm add ${packageName}`,
1217
+ };
1218
+ }
1219
+ function findProjectRoot(startDir) {
1220
+ let current = resolve(startDir);
1221
+ for (let i = 0; i < 20; i++) {
1222
+ if (existsSync(join(current, "package.json")))
1223
+ return current;
1224
+ const parent = dirname(current);
1225
+ if (parent === current)
1226
+ return null;
1227
+ current = parent;
1228
+ }
1229
+ return null;
1230
+ }
1231
+ function spawnCommand(command, args, cwd) {
1232
+ return new Promise((resolve) => {
1233
+ const proc = spawn(command, args, { cwd, shell: false, stdio: ["ignore", "pipe", "pipe"] });
1234
+ const stdoutChunks = [];
1235
+ const stderrChunks = [];
1236
+ proc.stdout.on("data", (chunk) => stdoutChunks.push(chunk));
1237
+ proc.stderr.on("data", (chunk) => stderrChunks.push(chunk));
1238
+ proc.on("close", (code) => {
1239
+ resolve({
1240
+ stdout: Buffer.concat(stdoutChunks).toString("utf8").slice(0, 4096),
1241
+ stderr: Buffer.concat(stderrChunks).toString("utf8").slice(0, 4096),
1242
+ code: code ?? 1,
1243
+ });
1244
+ });
1245
+ proc.on("error", (err) => {
1246
+ resolve({ stdout: "", stderr: err.message, code: 1 });
1247
+ });
1248
+ });
1249
+ }
1250
+ async function executeQuery(body, schemaDir) {
1251
+ const rt = getAskDbRuntimeConfig();
1252
+ const { provider, databaseUrl, file } = rt.studio.execute;
1253
+ if (!isRecord(body) || typeof body.sql !== "string" || body.sql.trim() === "") {
1254
+ return { ok: false, error: "`sql` is required." };
1255
+ }
1256
+ const sql = body.sql;
1257
+ const params = Array.isArray(body.params) ? body.params : [];
1258
+ const projectRoot = findProjectRoot(schemaDir) ?? schemaDir;
1259
+ const def = EXECUTE_DRIVER_REGISTRY[provider];
1260
+ return def.execute({ connectionString: databaseUrl, file, sql, params, projectRoot });
1261
+ }
1262
+ function isLoopbackRequest(req) {
1263
+ const socket = req.socket;
1264
+ const remoteAddress = socket?.remoteAddress ?? "";
1265
+ return (remoteAddress === "127.0.0.1" ||
1266
+ remoteAddress === "::1" ||
1267
+ remoteAddress === "::ffff:127.0.0.1" ||
1268
+ remoteAddress === "localhost");
1269
+ }
549
1270
  function readOptionalJson(path) {
550
1271
  if (!existsSync(path))
551
1272
  return undefined;
@@ -556,6 +1277,17 @@ function readOptionalJson(path) {
556
1277
  return undefined;
557
1278
  }
558
1279
  }
1280
+ function parseConceptsBody(body) {
1281
+ if (!isRecord(body) || !Array.isArray(body.concepts)) {
1282
+ throw new StudioHttpError(400, "`concepts` array is required.");
1283
+ }
1284
+ for (const concept of body.concepts) {
1285
+ if (!isRecord(concept) || typeof concept.id !== "string" || typeof concept.label !== "string") {
1286
+ throw new StudioHttpError(400, "Each concept requires a string `id` and `label`.");
1287
+ }
1288
+ }
1289
+ return body.concepts;
1290
+ }
559
1291
  function parseTableDraftBody(body) {
560
1292
  if (!isRecord(body) || !isRecord(body.draft)) {
561
1293
  throw new StudioHttpError(400, "`draft` is required.");
@@ -604,9 +1336,26 @@ function parseAskBody(body) {
604
1336
  if (mode !== "full" && mode !== "rag") {
605
1337
  throw new StudioHttpError(400, "`mode` must be `full` or `rag`.");
606
1338
  }
1339
+ let tenantScope;
1340
+ if (body.tenantScope !== undefined && body.tenantScope !== null) {
1341
+ const parsed = tenantScopeSchema.safeParse(body.tenantScope);
1342
+ if (!parsed.success) {
1343
+ throw new StudioHttpError(400, `Invalid tenantScope: ${parsed.error.message}`);
1344
+ }
1345
+ tenantScope = parsed.data;
1346
+ }
1347
+ let tenantSqlMode;
1348
+ if (typeof body.tenantSqlMode === "string") {
1349
+ if (body.tenantSqlMode !== "sql-only" && body.tenantSqlMode !== "sql-params") {
1350
+ throw new StudioHttpError(400, "`tenantSqlMode` must be `sql-only` or `sql-params`.");
1351
+ }
1352
+ tenantSqlMode = body.tenantSqlMode;
1353
+ }
607
1354
  return {
608
1355
  question: body.question.trim(),
609
1356
  useRag: mode === "rag",
1357
+ tenantScope,
1358
+ tenantSqlMode,
610
1359
  };
611
1360
  }
612
1361
  function parseRagQuery(body) {
@@ -676,8 +1425,10 @@ function createRequestUsageCollector() {
676
1425
  function normalizeGenerationUsage(value) {
677
1426
  if (!isRecord(value))
678
1427
  return undefined;
679
- const promptTokens = readFiniteNumber(value.promptTokens);
680
- const completionTokens = readFiniteNumber(value.completionTokens);
1428
+ // ai@6 LanguageModelUsage uses inputTokens/outputTokens rather than the
1429
+ // legacy promptTokens/completionTokens names.
1430
+ const promptTokens = readFiniteNumber(value.promptTokens) ?? readFiniteNumber(value.inputTokens);
1431
+ const completionTokens = readFiniteNumber(value.completionTokens) ?? readFiniteNumber(value.outputTokens);
681
1432
  const totalTokens = readFiniteNumber(value.totalTokens);
682
1433
  if (promptTokens === undefined && completionTokens === undefined && totalTokens === undefined) {
683
1434
  return undefined;