@useatlas/create 0.0.5 → 0.0.7

package/templates/docker/src/lib/agent.ts

@@ -2,8 +2,15 @@
  * The Atlas agent.
  *
  * Runs a single-agent loop driven by a ToolRegistry (default: explore,
- * executeSQL). The loop runs until the step limit (25) is reached or
- * the model stops issuing tool calls.
+ * executeSQL). The loop runs until the step limit is reached (configurable
+ * via `ATLAS_AGENT_MAX_STEPS`, default 25) or the model stops issuing
+ * tool calls.
+ *
+ * Effect migration (P10c):
+ * The agent function optionally reads its dependencies (model, tools,
+ * user context) from Effect Context when available, falling back to
+ * global singletons otherwise. This makes the agent testable via
+ * Layer.provide with mock services.
  */
 
 import {
@@ -12,20 +19,55 @@ import {
   streamText,
   type ModelMessage,
   type SystemModelMessage,
+  type ToolSet,
   type UIMessage,
 } from "ai";
-import { getModel, getProviderType, type ProviderType } from "./providers";
+import type { LanguageModel } from "ai";
+import { Effect } from "effect";
+import { getModel, getProviderType, getModelFromWorkspaceConfig, getWorkspaceProviderType, type ProviderType } from "./providers";
 import { defaultRegistry, type ToolRegistry } from "./tools/registry";
 import { getContextFragments, getDialectHints } from "./plugins/tools";
 import { connections, detectDBType, type ConnectionMetadata, type DBType } from "./db/connection";
-import { getCrossSourceJoins, type CrossSourceJoin } from "./semantic";
-import { getSemanticIndex } from "./semantic-index";
-import { createLogger } from "./logger";
-import { trace, SpanStatusCode } from "@opentelemetry/api";
+import { getCrossSourceJoins, type CrossSourceJoin, loadOrgWhitelist, getOrgSemanticIndex } from "./semantic";
+import { getSemanticIndex } from "./semantic/search";
+import { getConfig } from "./config";
+import { createLogger, getRequestContext } from "./logger";
+import { getSetting } from "./settings";
+import { hasInternalDB, internalExecute } from "./db/internal";
+import { logUsageEvent } from "./metering";
+import { buildLearnedPatternsSection } from "./learn/pattern-cache";
+import { dispatchMutableHook } from "./plugins/hooks";
+import { plugins } from "./plugins/registry";
+import {
+  trace,
+  SpanStatusCode,
+  context as otelContext,
+} from "@opentelemetry/api";
+import { AtlasAiModel, type AtlasAiModelShape } from "./effect/ai";
 
 const log = createLogger("agent");
 const tracer = trace.getTracer("atlas");
 
+const DEFAULT_MAX_STEPS = 25;
+const MIN_MAX_STEPS = 1;
+const MAX_MAX_STEPS = 100;
+
+let lastWarnedMaxSteps: string | undefined;
+
+/** Read agent max steps from settings cache (DB override > env var > default). */
+function getAgentMaxSteps(): number {
+  const raw = getSetting("ATLAS_AGENT_MAX_STEPS") ?? String(DEFAULT_MAX_STEPS);
+  const n = parseInt(raw, 10);
+  if (!Number.isFinite(n) || n < MIN_MAX_STEPS || n > MAX_MAX_STEPS) {
+    if (raw !== lastWarnedMaxSteps) {
+      log.warn({ value: raw }, `Invalid ATLAS_AGENT_MAX_STEPS value; using default ${DEFAULT_MAX_STEPS}`);
+      lastWarnedMaxSteps = raw;
+    }
+    return DEFAULT_MAX_STEPS;
+  }
+  return n;
+}
+
 const SYSTEM_PROMPT_PREFIX = `You are Atlas, an expert data analyst AI. You answer questions about data by exploring a semantic layer, writing SQL, and interpreting results.
 
 ## Your Workflow
@@ -67,7 +109,15 @@ When a SQL query fails, read the error carefully before retrying:
 - **Type mismatch** — You may need to CAST a column (e.g., CAST(value AS numeric)). Check the column type in the entity schema.
 - **Timeout** — Simplify the query: remove unnecessary JOINs, add WHERE filters to reduce the dataset, or break into smaller queries.
 - Never retry the exact same SQL. Always fix the identified issue first.
-- Max 2 retries per question — if the query still fails, explain the issue to the user.`;
+- Max 2 retries per question — if the query still fails, explain the issue to the user.
+
+## Suggested Follow-ups
+After each substantive answer, end your response with a <suggestions> block containing 2-3 contextual follow-up questions the user might ask next. Base them on the tables, metrics, and data from the current answer. Format:
+<suggestions>
+What is the trend over the last 12 months?
+How does this break down by region?
+Which accounts contribute the most?
+</suggestions>`;
 
 const MYSQL_DIALECT_GUIDE = `
 
@@ -161,7 +211,7 @@ const PYTHON_GUIDANCE = `
 
 **Chart guidance:** prefer \`_atlas_chart\` (interactive Recharts) for bar/line/pie charts. Use \`chart_path()\` only for advanced matplotlib visualizations that Recharts cannot render.`;
 
-function buildSystemPrompt(registry: ToolRegistry): string {
+function buildSystemPrompt(registry: ToolRegistry, orgSemanticIndex?: string, learnedPatternsSection?: string): string {
   let base = SYSTEM_PROMPT_PREFIX + "\n\n" + registry.describe() + "\n\n" + SYSTEM_PROMPT_SUFFIX;
 
   // Add Python guidance only when the tool is available
@@ -169,10 +219,19 @@ function buildSystemPrompt(registry: ToolRegistry): string {
     base += "\n" + PYTHON_GUIDANCE;
   }
 
-  // Append the pre-indexed semantic layer summary
-  const semanticIndex = getSemanticIndex();
-  if (semanticIndex) {
-    base += "\n\n" + semanticIndex;
+  // Append the pre-indexed semantic layer summary (respects config)
+  const indexEnabled = getConfig()?.semanticIndex?.enabled !== false;
+  if (indexEnabled) {
+    // Prefer org-scoped index if available, fall back to file-based
+    const semanticIndex = orgSemanticIndex || getSemanticIndex();
+    if (semanticIndex) {
+      base += "\n\n" + semanticIndex;
+    }
+  }
+
+  // Append learned patterns (if any)
+  if (learnedPatternsSection) {
+    base += "\n\n" + learnedPatternsSection;
   }
 
   // Append plugin context fragments (if any)
@@ -222,14 +281,22 @@ function buildSystemPrompt(registry: ToolRegistry): string {
  *   `providerOptions.anthropic.cacheControl` (~80% savings on steps 2+).
  * - Bedrock (non-Anthropic): returns a SystemModelMessage with
  *   `providerOptions.bedrock.cachePoint`.
- * - OpenAI / Ollama / Gateway: returns a plain string (OpenAI caches
- *   automatically for prompts >= 1024 tokens; others have no caching).
+ * - OpenAI / Ollama / OpenAI-compatible / Gateway: returns a plain string
+ *   (OpenAI caches automatically for prompts >= 1024 tokens; others have
+ *   no caching).
  */
 export function buildSystemParam(
   providerType: ProviderType,
   registry: ToolRegistry = defaultRegistry,
+  warnings?: string[],
+  orgSemanticIndex?: string,
+  learnedPatternsSection?: string,
 ): string | SystemModelMessage {
-  const content = buildSystemPrompt(registry);
+  let content = buildSystemPrompt(registry, orgSemanticIndex, learnedPatternsSection);
+
+  if (warnings && warnings.length > 0) {
+    content += "\n\n## Warnings\n\n" + warnings.map((w) => `- ${w}`).join("\n");
+  }
 
   switch (providerType) {
     case "anthropic":
@@ -251,6 +318,7 @@ export function buildSystemParam(
       };
     case "openai":
     case "ollama":
+    case "openai-compatible":
     case "gateway":
       return content;
     default: {
@@ -269,7 +337,7 @@ export function buildSystemParam(
  *
  * - Anthropic / Bedrock-Anthropic: `providerOptions.anthropic.cacheControl`
  * - Bedrock (non-Anthropic): `providerOptions.bedrock.cachePoint`
- * - OpenAI / Ollama / Gateway: no-op (OpenAI caches automatically)
+ * - OpenAI / Ollama / OpenAI-compatible / Gateway: no-op (OpenAI caches automatically)
  */
 export function applyCacheControl(
   messages: ModelMessage[],
@@ -308,6 +376,7 @@ export function applyCacheControl(
     }
     case "openai":
     case "ollama":
+    case "openai-compatible":
     case "gateway":
       return messages;
     default: {
@@ -317,27 +386,196 @@ export function applyCacheControl(
   }
 }
 
+/**
+ * Wrap each tool's execute function with beforeToolCall / afterToolCall
+ * plugin hook dispatch. No-op when no plugins are registered.
+ *
+ * Execution order for tools that have domain-specific hooks (e.g. executeSQL):
+ *   beforeToolCall → beforeQuery → execute → afterQuery → afterToolCall
+ *
+ * - beforeToolCall: can return `{ args }` to modify args, or throw to reject
+ * - afterToolCall: can return `{ result }` to modify the result, or throw to reject
+ */
+function wrapToolsWithHooks(
+  toolSet: ToolSet,
+  hookCtx: { userId?: string; conversationId?: string },
+): ToolSet {
+  // Optimization: skip wrapping when no plugins are registered at request time.
+  // Plugins are initialized at boot before any requests — this is safe.
+  if (plugins.size === 0) return toolSet;
+
+  const callCounter = { value: 0 };
+  const wrapped: ToolSet = {};
+
+  for (const [name, t] of Object.entries(toolSet)) {
+    if (!t.execute) {
+      wrapped[name] = t;
+      continue;
+    }
+
+    const origExecute = t.execute;
+    wrapped[name] = {
+      ...t,
+      execute: async (
+        args: Record<string, unknown>,
+        options: Parameters<NonNullable<typeof t.execute>>[1],
+      ) => {
+        callCounter.value++;
+        const ctx = {
+          toolName: name,
+          args,
+          context: { ...hookCtx, toolCallCount: callCounter.value },
+        };
+
+        // beforeToolCall — can modify args or throw to reject
+        let finalArgs: Record<string, unknown>;
+        try {
+          finalArgs = await dispatchMutableHook(
+            "beforeToolCall",
+            ctx,
+            "args",
+          ) as Record<string, unknown>;
+        } catch (err) {
+          log.warn(
+            { toolName: name, err: err instanceof Error ? err : new Error(String(err)) },
+            "Tool call rejected by plugin",
+          );
+          return `Tool call rejected by plugin: ${err instanceof Error ? err.message : String(err)}`;
+        }
+
+        const start = Date.now();
+        const result = await origExecute(finalArgs, options);
+        const durationMs = Date.now() - start;
+
+        // afterToolCall — can modify result or throw to reject
+        try {
+          return await dispatchMutableHook(
+            "afterToolCall",
+            { ...ctx, args: finalArgs, result, durationMs, context: { ...hookCtx, toolCallCount: callCounter.value } },
+            "result",
+          );
+        } catch (err) {
+          log.error(
+            { toolName: name, err: err instanceof Error ? err : new Error(String(err)) },
+            "Tool result rejected by plugin",
+          );
+          return `Tool result rejected by plugin: ${err instanceof Error ? err.message : String(err)}`;
+        }
+      },
+    };
+  }
+
+  return wrapped;
+}
+
 /**
  * Run the Atlas agent loop.
  *
  * @param messages - The conversation history from the chat UI.
  * @param tools - Optional custom {@link ToolRegistry}. Defaults to
  *   {@link defaultRegistry} (explore + executeSQL). The loop terminates
- *   when the step limit (25) is reached or the model stops issuing tool calls.
+ *   when the step limit is reached (configurable via `ATLAS_AGENT_MAX_STEPS`,
+ *   default 25) or the model stops issuing tool calls.
+ * @param conversationId - Optional conversation ID for token usage tracking.
+ *   When provided, the recorded token usage row links to this conversation.
  */
 export async function runAgent({
   messages,
   tools: toolRegistry = defaultRegistry,
+  conversationId,
+  warnings,
+  maxSteps: maxStepsOverride,
+  /** Optional pre-resolved AI model. When provided, skips provider resolution. */
+  aiModel: injectedAiModel,
 }: {
   messages: UIMessage[];
   tools?: ToolRegistry;
+  conversationId?: string;
+  warnings?: string[];
+  /** Override the default agent step limit (e.g. for demo mode). */
+  maxSteps?: number;
+  /** Pre-resolved AI model from Effect Context (P10c). */
+  aiModel?: AtlasAiModelShape;
 }) {
-  const model = getModel();
-  const providerType = getProviderType();
+  // Capture context eagerly — AsyncLocalStorage may have exited by the time onFinish fires
+  const reqCtx = getRequestContext();
+  const userId = reqCtx?.user?.id ?? null;
+  const orgId = reqCtx?.user?.activeOrganizationId;
+
+  // Resolve model: injected > workspace config (enterprise) > platform env vars
+  let model: LanguageModel;
+  let providerType: ProviderType;
+
+  if (injectedAiModel) {
+    // Model provided via Effect Context (P10c) — skip provider resolution
+    model = injectedAiModel.model;
+    providerType = injectedAiModel.providerType;
+  } else {
+    let workspaceConfig: { provider: import("@useatlas/types").ModelConfigProvider; model: string; apiKey: string; baseUrl: string | null } | null = null;
+    if (orgId && hasInternalDB()) {
+      try {
+        const { getWorkspaceModelConfigRaw } = await import("@atlas/ee/platform/model-routing");
+        workspaceConfig = await getWorkspaceModelConfigRaw(orgId);
+      } catch (err) {
+        log.debug({ orgId, err: err instanceof Error ? err.message : String(err) }, "Workspace model config not available — using platform default");
+      }
+    }
+
+    if (workspaceConfig) {
+      model = getModelFromWorkspaceConfig(workspaceConfig);
+      providerType = getWorkspaceProviderType(workspaceConfig.provider);
+      log.info({ orgId, provider: workspaceConfig.provider, model: workspaceConfig.model }, "Using workspace model config");
+    } else {
+      model = getModel();
+      providerType = getProviderType();
+    }
+  }
+
+  const resolvedModelId = typeof model === "string" ? model : model.modelId;
+
+  // Pre-load org-scoped semantic data and learned patterns before the agent loop.
+  // These are independent async operations — run in parallel to avoid waterfalls.
+  let orgSemanticIndex: string | undefined;
+  let learnedPatternsSection: string | undefined;
+
+  const orgSemanticPromise = (orgId && hasInternalDB())
+    ? Promise.all([loadOrgWhitelist(orgId), getOrgSemanticIndex(orgId)])
+        .then(([, idx]) => { orgSemanticIndex = idx || undefined; })
+        .catch((err: unknown) => {
+          log.error({ orgId, err: err instanceof Error ? err.message : String(err) }, "Failed to load org semantic data — agent will use file-based fallback");
+          if (!warnings) warnings = [];
+          warnings.push("Your organization's semantic layer could not be loaded. Using default configuration. Contact your admin if this persists.");
+        })
+    : Promise.resolve();
+
+  const learnedPatternsPromise = hasInternalDB()
+    ? (async () => {
+        const lastUserMsg = [...messages].reverse().find((m) => m.role === "user");
+        const question = lastUserMsg?.parts
+          ?.filter((p): p is { type: "text"; text: string } => p.type === "text")
+          .map((p) => p.text)
+          .join(" ") ?? "";
+        if (question) {
+          const section = await buildLearnedPatternsSection(orgId ?? null, question);
+          if (section) learnedPatternsSection = section;
+        }
+      })().catch((err: unknown) => {
+        log.warn({ orgId, err: err instanceof Error ? err.message : String(err) }, "Failed to load learned patterns — continuing without");
+      })
+    : Promise.resolve();
+
+  await Promise.all([orgSemanticPromise, learnedPatternsPromise]);
 
   const span = tracer.startSpan("atlas.agent", {
-    attributes: { provider: providerType, messageCount: messages.length },
+    attributes: {
+      "atlas.provider": providerType,
+      "atlas.model": resolvedModelId,
+      "atlas.message_count": messages.length,
+    },
   });
+  // Make the agent span the active context so tool spans (withSpan in
+  // sql.ts, explore.ts) become children in the trace hierarchy.
+  const agentCtx = trace.setSpan(otelContext.active(), span);
 
   let spanEnded = false;
   function endSpan(code: SpanStatusCode, message?: string) {
@@ -347,16 +585,21 @@ export async function runAgent({
     span.end();
   }
 
+  // Resolve async work before entering otelContext.with() (sync callback).
+  const modelMessages = await convertToModelMessages(messages);
+  const rawTools = toolRegistry.getAll();
+  const tools = wrapToolsWithHooks(rawTools, { userId: userId ?? undefined, conversationId });
+
   let result;
   try {
-    result = streamText({
+    result = otelContext.with(agentCtx, () => streamText({
       model,
-      system: buildSystemParam(providerType, toolRegistry),
-      messages: await convertToModelMessages(messages),
-      tools: toolRegistry.getAll(),
+      system: buildSystemParam(providerType, toolRegistry, warnings, orgSemanticIndex, learnedPatternsSection),
+      messages: modelMessages,
+      tools,
       temperature: 0.2,
       maxOutputTokens: 4096,
-      stopWhen: stepCountIs(25),
+      stopWhen: stepCountIs(maxStepsOverride ?? getAgentMaxSteps()),
       // totalMs: 180s for self-hosted (full agent loop budget).
       // On Vercel, maxDuration caps the serverless function at 300s (Pro plan).
       timeout: { totalMs: 180_000, stepMs: 30_000, chunkMs: 5_000 },
@@ -405,14 +648,61 @@ export async function runAgent({
           "agent finished",
         );
         span.setAttributes({
-          finishReason: finishReason ?? "",
-          totalSteps: steps.length,
-          totalInputTokens: totalUsage?.inputTokens ?? 0,
-          totalOutputTokens: totalUsage?.outputTokens ?? 0,
+          "atlas.finish_reason": finishReason ?? "",
+          "atlas.total_steps": steps.length,
+          "atlas.total_input_tokens": totalUsage?.inputTokens ?? 0,
+          "atlas.total_output_tokens": totalUsage?.outputTokens ?? 0,
         });
         endSpan(SpanStatusCode.OK);
+
+        // Persist token usage to internal DB (fire-and-forget).
+        // Shares the internalExecute circuit breaker with audit writes.
+        if (hasInternalDB() && totalUsage) {
+          try {
+            internalExecute(
+              `INSERT INTO token_usage (user_id, conversation_id, prompt_tokens, completion_tokens, model, provider, org_id)
+               VALUES ($1, $2, $3, $4, $5, $6, $7)`,
+              [
+                userId,
+                conversationId ?? null,
+                totalUsage.inputTokens ?? 0,
+                totalUsage.outputTokens ?? 0,
+                resolvedModelId,
+                providerType,
+                orgId ?? null,
+              ],
+            );
+          } catch (err) {
+            log.warn({ err: err instanceof Error ? err.message : String(err) }, "Failed to persist token usage");
+          }
+
+          // Log usage metering events for billing/overage tracking.
+          // Wrapped in its own try/catch to ensure a metering failure
+          // never disrupts the onFinish callback or stream finalization.
+          try {
+            const totalTokens = (totalUsage.inputTokens ?? 0) + (totalUsage.outputTokens ?? 0);
+            logUsageEvent({
+              workspaceId: orgId ?? null,
+              userId: userId ?? null,
+              eventType: "query",
+              quantity: 1,
+              metadata: { conversationId, model: resolvedModelId, steps: steps.length },
+            });
+            if (totalTokens > 0) {
+              logUsageEvent({
+                workspaceId: orgId ?? null,
+                userId: userId ?? null,
+                eventType: "token",
+                quantity: totalTokens,
+                metadata: { input: totalUsage.inputTokens ?? 0, output: totalUsage.outputTokens ?? 0 },
+              });
+            }
+          } catch (err) {
+            log.warn({ err: err instanceof Error ? err.message : String(err) }, "Failed to log usage metering events");
+          }
+        }
       },
-    });
+    }));
   } catch (err) {
     endSpan(
       SpanStatusCode.ERROR,
@@ -423,3 +713,45 @@ export async function runAgent({
 
   return result;
 }
+
+// ── Effect-based agent runner (P10c) ────────────────────────────────
+
+/**
+ * Run the Atlas agent as an Effect program.
+ *
+ * Reads the AI model from `AtlasAiModel` in the Effect Context and
+ * delegates to `runAgent`. This is the preferred entry point for
+ * Effect-based callers — it makes the agent testable via Layer.provide
+ * with a mock LLM.
+ *
+ * @example
+ * ```ts
+ * import { runAgentEffect } from "@atlas/api/lib/agent";
+ * import { createAiModelTestLayer } from "@atlas/api/lib/effect/ai";
+ *
+ * // In tests — provide a mock model
+ * const result = await Effect.runPromise(
+ *   runAgentEffect({ messages }).pipe(
+ *     Effect.provide(createAiModelTestLayer({ ... })),
+ *   ),
+ * );
+ * ```
+ */
+export function runAgentEffect(params: {
+  messages: UIMessage[];
+  tools?: ToolRegistry;
+  conversationId?: string;
+  warnings?: string[];
+  maxSteps?: number;
+}): Effect.Effect<ReturnType<typeof streamText>, Error, AtlasAiModel> {
+  return Effect.gen(function* () {
+    const aiModel = yield* AtlasAiModel;
+    return yield* Effect.tryPromise({
+      try: () => runAgent({ ...params, aiModel }),
+      catch: (err) =>
+        new Error(
+          `Agent execution failed: ${err instanceof Error ? err.message : String(err)}`,
+        ),
+    });
+  });
+}

package/templates/docker/src/lib/auth/admin-permissions.ts

@@ -0,0 +1,38 @@
+/**
+ * Admin-scoped access control for Better Auth's admin plugin.
+ *
+ * Defines user-level roles (`admin`, `platform_admin`) that govern
+ * who can perform administrative operations (ban users, set roles,
+ * impersonate, manage sessions, etc.).
+ *
+ * Separate from org-permissions.ts which handles organization-scoped
+ * RBAC (owner/admin/member within an org). This file handles the
+ * top-level user role that determines system-wide admin privileges.
+ *
+ * Exported for use in the server config (server.ts). A client-side
+ * mirror exists at packages/web/src/lib/auth/admin-permissions.ts —
+ * keep both files in sync.
+ */
+
+import { createAccessControl } from "better-auth/plugins/access";
+import { defaultStatements, adminAc } from "better-auth/plugins/admin/access";
+
+const statement = {
+  ...defaultStatements,
+} as const;
+
+export const adminAccessControl = createAccessControl(statement);
+
+/** Standard admin — full admin permissions over users and sessions. */
+export const adminRole = adminAccessControl.newRole({
+  ...adminAc.statements,
+});
+
+/**
+ * Platform admin — same permissions as admin.
+ * Distinguished by name for platform-operator routes that check
+ * `user.role === "platform_admin"` to gate cross-tenant operations.
+ */
+export const platformAdminRole = adminAccessControl.newRole({
+  ...adminAc.statements,
+});

package/templates/docker/src/lib/auth/audit.ts

@@ -19,6 +19,7 @@
 import { createLogger, getRequestContext } from "@atlas/api/lib/logger";
 import { hasInternalDB, internalExecute } from "@atlas/api/lib/db/internal";
 import { SENSITIVE_PATTERNS } from "@atlas/api/lib/security";
+import { recordQueryEvent } from "@atlas/api/lib/security/abuse";
 import type { DBType } from "@atlas/api/lib/db/connection";
 
 const log = createLogger("audit");
@@ -27,8 +28,8 @@ const PINO_SQL_LIMIT = 500;
 const DB_SQL_LIMIT = 2000;
 
 export type AuditEntry =
-  | { sql: string; durationMs: number; rowCount: number; success: true; sourceId?: string; sourceType?: DBType; targetHost?: string }
-  | { sql: string; durationMs: number; rowCount: null; success: false; error?: string; sourceId?: string; sourceType?: DBType; targetHost?: string };
+  | { sql: string; durationMs: number; rowCount: number; success: true; sourceId?: string; sourceType?: DBType; targetHost?: string; tablesAccessed?: string[]; columnsAccessed?: string[] }
+  | { sql: string; durationMs: number; rowCount: null; success: false; error?: string; sourceId?: string; sourceType?: DBType; targetHost?: string; tablesAccessed?: string[]; columnsAccessed?: string[] };
 
 function scrubError(error: string | undefined): string | undefined {
   if (!error) return undefined;
@@ -58,16 +59,27 @@ export function logQueryAudit(entry: AuditEntry): void {
       ...(entry.sourceId && { sourceId: entry.sourceId }),
       ...(entry.sourceType && { sourceType: entry.sourceType }),
       ...(entry.targetHost && { targetHost: entry.targetHost }),
+      ...(entry.tablesAccessed?.length && { tablesAccessed: entry.tablesAccessed }),
+      ...(entry.columnsAccessed?.length && { columnsAccessed: entry.columnsAccessed }),
     },
     entry.success ? "query_success" : "query_failure",
   );
 
+  // Record query event for abuse detection (per-workspace anomaly tracking)
+  const orgId = ctx?.user?.activeOrganizationId;
+  if (orgId) {
+    recordQueryEvent(orgId, {
+      success: entry.success,
+      tablesAccessed: entry.tablesAccessed,
+    });
+  }
+
   // Insert into audit_log when internal DB is available (SQL truncated to 2000 chars)
   if (hasInternalDB()) {
     try {
       internalExecute(
-        `INSERT INTO audit_log (user_id, user_label, auth_mode, sql, duration_ms, row_count, success, error, source_id, source_type, target_host)
-         VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11)`,
+        `INSERT INTO audit_log (user_id, user_label, auth_mode, sql, duration_ms, row_count, success, error, source_id, source_type, target_host, tables_accessed, columns_accessed, org_id)
+         VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14)`,
         [
           userId,
           userLabel,
@@ -80,6 +92,9 @@ export function logQueryAudit(entry: AuditEntry): void {
           entry.sourceId ?? null,
           entry.sourceType ?? null,
           entry.targetHost ?? null,
+          entry.tablesAccessed?.length ? JSON.stringify(entry.tablesAccessed) : null,
+          entry.columnsAccessed?.length ? JSON.stringify(entry.columnsAccessed) : null,
+          ctx?.user?.activeOrganizationId ?? null,
         ],
       );
     } catch (err) {

package/templates/docker/src/lib/auth/byot.ts

@@ -50,7 +50,7 @@ function extractRoleFromPayload(payload: Record<string, unknown>) {
     if (typeof value === "string") {
       const role = parseRole(value);
       if (role) return role;
-      log.warn({ claimPath, value, validRoles: ["viewer", "analyst", "admin"] }, "JWT role claim value is not a valid Atlas role — ignoring");
+      log.warn({ claimPath, value, validRoles: ["member", "admin", "owner"] }, "JWT role claim value is not a valid Atlas role — ignoring");
     } else if (value !== undefined && value !== null) {
       log.warn({ claimPath, type: typeof value }, "JWT role claim is not a string — ignoring");
     }
@@ -63,7 +63,7 @@ function extractRoleFromPayload(payload: Record<string, unknown>) {
     if (typeof value === "string") {
       const role = parseRole(value);
       if (role) return role;
-      log.warn({ claim, value, validRoles: ["viewer", "analyst", "admin"] }, "JWT role claim value is not a valid Atlas role — ignoring");
+      log.warn({ claim, value, validRoles: ["member", "admin", "owner"] }, "JWT role claim value is not a valid Atlas role — ignoring");
     } else if (value !== undefined && value !== null) {
       log.warn({ claim, type: typeof value }, "JWT role claim is not a string — ignoring");
     }
@@ -117,7 +117,7 @@ export async function validateBYOT(req: Request): Promise<AuthResult> {
     return {
       authenticated: true,
       mode: "byot",
-      user: createAtlasUser(sub, "byot", email || sub, role, payload as Record<string, unknown>),
+      user: createAtlasUser(sub, "byot", email || sub, { role, claims: payload as Record<string, unknown> }),
     };
   } catch (err) {
     // Infrastructure errors — JWKS endpoint issues are not client auth failures.