@lightdash-tools/mcp 0.2.6 → 0.3.1

package/src/tools/index.ts

@@ -16,6 +16,7 @@ import { registerMetricsTools } from './metrics.js';
 import { registerSchedulersTools } from './schedulers.js';
 import { registerTagsTools } from './tags.js';
 import { registerContentTools } from './content.js';
+import { registerAiAgentTools } from './ai-agents.js';
 
 export function registerTools(server: McpServer, client: LightdashClient): void {
   registerProjectTools(server, client);
@@ -30,4 +31,5 @@ export function registerTools(server: McpServer, client: LightdashClient): void
   registerSchedulersTools(server, client);
   registerTagsTools(server, client);
   registerContentTools(server, client);
+  registerAiAgentTools(server, client);
 }

package/src/tools/shared.test.ts

@@ -1,7 +1,19 @@
-import { describe, it, expect, vi } from 'vitest';
-import { registerToolSafe, READ_ONLY_DEFAULT, WRITE_DESTRUCTIVE } from './shared';
+import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest';
+import { registerToolSafe, READ_ONLY_DEFAULT, WRITE_DESTRUCTIVE, WRITE_IDEMPOTENT } from './shared';
 import { SafetyMode } from '@lightdash-tools/common';
-import { setStaticSafetyMode } from '../config.js';
+import { setStaticSafetyMode, setStaticAllowedProjectUuids, setDryRunMode } from '../config.js';
+
+// Silence audit log output during tests
+vi.mock('@lightdash-tools/common', async (importOriginal) => {
+  // eslint-disable-next-line @typescript-eslint/no-explicit-any
+  const actual = await importOriginal<any>();
+  return {
+    ...actual,
+    getSessionId: () => 'test-session',
+    logAuditEntry: vi.fn(),
+    initAuditLog: vi.fn(),
+  };
+});
 
 describe('registerToolSafe', () => {
   const mockServer = {
@@ -10,8 +22,25 @@ describe('registerToolSafe', () => {
 
   const mockHandler = vi.fn().mockResolvedValue({ content: [{ type: 'text', text: 'success' }] });
 
+  beforeEach(() => {
+    mockServer.registerTool.mockClear();
+    mockHandler.mockClear();
+    // Reset globals to safe defaults
+    setStaticSafetyMode(SafetyMode.WRITE_DESTRUCTIVE);
+    setStaticAllowedProjectUuids([]);
+    setDryRunMode(false);
+    process.env.LIGHTDASH_TOOL_SAFETY_MODE = SafetyMode.WRITE_DESTRUCTIVE;
+    delete process.env.LIGHTDASH_TOOLS_ALLOWED_PROJECTS;
+    delete process.env.LIGHTDASH_DRY_RUN;
+  });
+
+  afterEach(() => {
+    delete process.env.LIGHTDASH_TOOL_SAFETY_MODE;
+  });
+
   it('should allow read-only tool in read-only mode', async () => {
     process.env.LIGHTDASH_TOOL_SAFETY_MODE = SafetyMode.READ_ONLY;
+    setStaticSafetyMode(SafetyMode.WRITE_DESTRUCTIVE); // static = allow all
 
     registerToolSafe(
       mockServer,
@@ -36,6 +65,7 @@ describe('registerToolSafe', () => {
 
   it('should block destructive tool in read-only mode', async () => {
     process.env.LIGHTDASH_TOOL_SAFETY_MODE = SafetyMode.READ_ONLY;
+    setStaticSafetyMode(SafetyMode.WRITE_DESTRUCTIVE);
 
     registerToolSafe(
       mockServer,
@@ -48,7 +78,7 @@ describe('registerToolSafe', () => {
       mockHandler,
     );
 
-    const [, options, handler] = mockServer.registerTool.mock.calls[1];
+    const [, options, handler] = mockServer.registerTool.mock.calls[0];
 
     expect(options.description).toContain('[DISABLED in read-only mode]');
 
@@ -59,6 +89,7 @@ describe('registerToolSafe', () => {
 
   it('should allow destructive tool in write-destructive mode', async () => {
     process.env.LIGHTDASH_TOOL_SAFETY_MODE = SafetyMode.WRITE_DESTRUCTIVE;
+    setStaticSafetyMode(SafetyMode.WRITE_DESTRUCTIVE);
 
     registerToolSafe(
       mockServer,
@@ -71,7 +102,7 @@ describe('registerToolSafe', () => {
       mockHandler,
     );
 
-    const [, options, handler] = mockServer.registerTool.mock.calls[2];
+    const [, options, handler] = mockServer.registerTool.mock.calls[0];
 
     expect(options.description).toBe('Delete something 2');
 
@@ -81,9 +112,7 @@ describe('registerToolSafe', () => {
 
   describe('static filtering (safety-mode)', () => {
     it('should skip registration if tool is more permissive than binded mode', () => {
-      // Set binded mode to READ_ONLY
       setStaticSafetyMode(SafetyMode.READ_ONLY);
-
       mockServer.registerTool.mockClear();
 
       registerToolSafe(
@@ -102,7 +131,6 @@ describe('registerToolSafe', () => {
 
     it('should allow registration if tool matches binded mode', () => {
       setStaticSafetyMode(SafetyMode.READ_ONLY);
-
       mockServer.registerTool.mockClear();
 
       registerToolSafe(
@@ -119,12 +147,8 @@ describe('registerToolSafe', () => {
       expect(mockServer.registerTool).toHaveBeenCalled();
     });
 
-    it('should allow everything if binded mode is undefined', () => {
-      // This is a bit tricky since it's a global. We might need a way to reset it.
-      // For now, let's assume we can just pass a permissive mode or it was undefined initially.
-      // Since we don't have a reset, let's just test that it works when set to DESTRUCTIVE.
+    it('should allow everything if binded mode is write-destructive', () => {
       setStaticSafetyMode(SafetyMode.WRITE_DESTRUCTIVE);
-
       mockServer.registerTool.mockClear();
 
       registerToolSafe(
@@ -141,4 +165,181 @@ describe('registerToolSafe', () => {
       expect(mockServer.registerTool).toHaveBeenCalled();
     });
   });
+
+  describe('project UUID allowlist', () => {
+    it('should allow calls when allowlist is empty (all projects permitted)', async () => {
+      setStaticAllowedProjectUuids([]);
+
+      registerToolSafe(
+        mockServer,
+        'list_charts',
+        { description: 'List charts', inputSchema: {}, annotations: READ_ONLY_DEFAULT },
+        mockHandler,
+      );
+
+      const [, , handler] = mockServer.registerTool.mock.calls[0];
+      const result = await handler({ projectUuid: 'any-uuid' });
+      expect(result.isError).toBeUndefined();
+      expect(result.content[0].text).toBe('success');
+    });
+
+    it('should allow calls for a singular projectUuid in the allowlist', async () => {
+      setStaticAllowedProjectUuids(['uuid-allowed', 'uuid-other']);
+
+      registerToolSafe(
+        mockServer,
+        'list_charts_allowed',
+        { description: 'List charts', inputSchema: {}, annotations: READ_ONLY_DEFAULT },
+        mockHandler,
+      );
+
+      const [, , handler] = mockServer.registerTool.mock.calls[0];
+      const result = await handler({ projectUuid: 'uuid-allowed' });
+      expect(result.isError).toBeUndefined();
+      expect(result.content[0].text).toBe('success');
+    });
+
+    it('should block calls for a singular projectUuid not in the allowlist', async () => {
+      setStaticAllowedProjectUuids(['uuid-allowed']);
+
+      registerToolSafe(
+        mockServer,
+        'list_charts_blocked',
+        { description: 'List charts', inputSchema: {}, annotations: READ_ONLY_DEFAULT },
+        mockHandler,
+      );
+
+      const [, , handler] = mockServer.registerTool.mock.calls[0];
+      const result = await handler({ projectUuid: 'uuid-denied' });
+      expect(result.isError).toBe(true);
+      expect(result.content[0].text).toContain('not in the list of allowed projects');
+      expect(result.content[0].text).toContain('uuid-denied');
+    });
+
+    it('should allow calls with no projectUuid arg even when allowlist is set', async () => {
+      setStaticAllowedProjectUuids(['uuid-allowed']);
+
+      registerToolSafe(
+        mockServer,
+        'list_projects_no_uuid',
+        { description: 'List projects', inputSchema: {}, annotations: READ_ONLY_DEFAULT },
+        mockHandler,
+      );
+
+      const [, , handler] = mockServer.registerTool.mock.calls[0];
+      // No projectUuid in args → allowlist does not apply
+      const result = await handler({});
+      expect(result.isError).toBeUndefined();
+      expect(result.content[0].text).toBe('success');
+    });
+
+    it('should allow when all projectUuids[] are in the allowlist', async () => {
+      setStaticAllowedProjectUuids(['uuid-a', 'uuid-b']);
+
+      registerToolSafe(
+        mockServer,
+        'search_content_allowed',
+        { description: 'Search content', inputSchema: {}, annotations: READ_ONLY_DEFAULT },
+        mockHandler,
+      );
+
+      const [, , handler] = mockServer.registerTool.mock.calls[0];
+      const result = await handler({ projectUuids: ['uuid-a', 'uuid-b'] });
+      expect(result.isError).toBeUndefined();
+      expect(result.content[0].text).toBe('success');
+    });
+
+    it('should block when any UUID in projectUuids[] is not in the allowlist', async () => {
+      setStaticAllowedProjectUuids(['uuid-a']);
+
+      registerToolSafe(
+        mockServer,
+        'search_content_blocked',
+        { description: 'Search content', inputSchema: {}, annotations: READ_ONLY_DEFAULT },
+        mockHandler,
+      );
+
+      const [, , handler] = mockServer.registerTool.mock.calls[0];
+      const result = await handler({ projectUuids: ['uuid-a', 'uuid-denied'] });
+      expect(result.isError).toBe(true);
+      expect(result.content[0].text).toContain('uuid-denied');
+      expect(result.content[0].text).toContain('not in the list of allowed projects');
+    });
+
+    it('should block when all projectUuids[] are outside the allowlist', async () => {
+      setStaticAllowedProjectUuids(['uuid-allowed']);
+
+      registerToolSafe(
+        mockServer,
+        'search_content_all_blocked',
+        { description: 'Search content', inputSchema: {}, annotations: READ_ONLY_DEFAULT },
+        mockHandler,
+      );
+
+      const [, , handler] = mockServer.registerTool.mock.calls[0];
+      const result = await handler({ projectUuids: ['uuid-x', 'uuid-y'] });
+      expect(result.isError).toBe(true);
+      expect(result.content[0].text).toContain('not in the list of allowed projects');
+    });
+  });
+
+  describe('dry-run mode', () => {
+    it('should not affect read-only tools in dry-run mode', async () => {
+      setDryRunMode(true);
+
+      registerToolSafe(
+        mockServer,
+        'list_things_dry',
+        { description: 'List things', inputSchema: {}, annotations: READ_ONLY_DEFAULT },
+        mockHandler,
+      );
+
+      const [, options, handler] = mockServer.registerTool.mock.calls[0];
+      expect(options.description).not.toContain('[DRY-RUN]');
+
+      const result = await handler({});
+      expect(result.content[0].text).toBe('success');
+    });
+
+    it('should simulate write-idempotent tools in dry-run mode', async () => {
+      setDryRunMode(true);
+      process.env.LIGHTDASH_TOOL_SAFETY_MODE = SafetyMode.WRITE_DESTRUCTIVE;
+
+      registerToolSafe(
+        mockServer,
+        'upsert_thing_dry',
+        { description: 'Upsert thing', inputSchema: {}, annotations: WRITE_IDEMPOTENT },
+        mockHandler,
+      );
+
+      const [, options, handler] = mockServer.registerTool.mock.calls[0];
+      expect(options.description).toContain('[DRY-RUN]');
+
+      const result = await handler({ projectUuid: 'uuid-x', slug: 'my-chart' });
+      expect(result.isError).toBeUndefined();
+      expect(result.content[0].text).toContain('[DRY-RUN]');
+      expect(result.content[0].text).toContain('No changes were made');
+      // Verify the underlying handler was NOT called
+      expect(mockHandler).not.toHaveBeenCalled();
+    });
+
+    it('should simulate destructive tools in dry-run mode', async () => {
+      setDryRunMode(true);
+      process.env.LIGHTDASH_TOOL_SAFETY_MODE = SafetyMode.WRITE_DESTRUCTIVE;
+
+      registerToolSafe(
+        mockServer,
+        'delete_thing_dry',
+        { description: 'Delete thing', inputSchema: {}, annotations: WRITE_DESTRUCTIVE },
+        mockHandler,
+      );
+
+      const [, options, handler] = mockServer.registerTool.mock.calls[0];
+      expect(options.description).toContain('[DRY-RUN]');
+
+      const result = await handler({ projectUuid: 'uuid-x' });
+      expect(result.content[0].text).toContain('No changes were made');
+      expect(mockHandler).not.toHaveBeenCalled();
+    });
+  });
 });

package/src/tools/shared.ts

@@ -1,13 +1,32 @@
 /**
  * Shared types and helpers for MCP tool registration.
+ *
+ * Guardrail layers applied by registerToolSafe (outer → inner):
+ *   1. Audit log wrapper   — captures timing and outcome for every call.
+ *   2. Project allowlist   — rejects calls targeting disallowed project UUIDs at runtime.
+ *   3. Dry-run wrapper     — simulates writes without executing them (registration-time).
+ *   4. Safety-mode wrapper — disables tools that exceed the configured safety level.
+ *   5. Raw handler         — the actual tool implementation.
  */
 
 import type { LightdashClient } from '@lightdash-tools/client';
-import { isAllowed, READ_ONLY_DEFAULT } from '@lightdash-tools/common';
+import {
+  isAllowed,
+  areAllProjectsAllowed,
+  extractProjectUuids,
+  READ_ONLY_DEFAULT,
+  logAuditEntry,
+  getSessionId,
+} from '@lightdash-tools/common';
 import type { ToolAnnotations } from '@lightdash-tools/common';
 import type { z } from 'zod';
 import { toMcpErrorMessage } from '../errors.js';
-import { getStaticSafetyMode, getSafetyMode } from '../config.js';
+import {
+  getStaticSafetyMode,
+  getSafetyMode,
+  getAllowedProjectUuids,
+  isDryRunMode,
+} from '../config.js';
 
 /** Prefix for all MCP tool names (disambiguation when multiple servers are connected). */
 export const TOOL_PREFIX = 'lightdash_tools__';
@@ -41,7 +60,27 @@ function mergeAnnotations(overrides?: ToolAnnotations): ToolAnnotations {
   return { ...DEFAULT_ANNOTATIONS, ...overrides };
 }
 
-/** Registers a tool with prefix and annotations. shortName is TOOL_PREFIX + shortName. Pass annotations explicitly (e.g. READ_ONLY_DEFAULT, WRITE_IDEMPOTENT, or WRITE_DESTRUCTIVE). */
+/**
+ * Internal marker attached to responses produced by a guardrail (safety-mode block,
+ * dry-run simulation, or project-allowlist denial). The audit wrapper reads this flag
+ * to set status = 'blocked', then strips it before returning to the MCP client.
+ */
+type BlockedContent = TextContent & { readonly _lightdashBlocked: true };
+
+function isGuardrailBlocked(result: TextContent): result is BlockedContent {
+  return (
+    '_lightdashBlocked' in result &&
+    (result as Record<string, unknown>)['_lightdashBlocked'] === true
+  );
+}
+
+/**
+ * Registers a tool with prefix and annotations, applying all guardrail layers.
+ * shortName is prefixed to become TOOL_PREFIX + shortName.
+ * Pass annotations explicitly (e.g. READ_ONLY_DEFAULT, WRITE_IDEMPOTENT, or WRITE_DESTRUCTIVE).
+ *
+ * CLI flag --allowed-projects always takes priority over LIGHTDASH_TOOLS_ALLOWED_PROJECTS.
+ */
 export function registerToolSafe(
   server: unknown,
   shortName: string,
@@ -51,23 +90,26 @@ export function registerToolSafe(
   const name = TOOL_PREFIX + shortName;
   const annotations = mergeAnnotations(options.annotations);
 
-  // Static Filtering: Skip registration if not allowed in static safety mode
+  // ── Static Filtering ──────────────────────────────────────────────────────
+  // Skip registration entirely if the tool exceeds the static safety mode.
   const staticMode = getStaticSafetyMode();
   if (staticMode && !isAllowed(staticMode, annotations)) {
     return;
   }
 
-  // Dynamic Enforcement: Wrap handler if not allowed in current safety mode (env)
+  // ── Safety-mode wrapper ───────────────────────────────────────────────────
+  // Tool is registered but calls are rejected at runtime when the dynamic mode
+  // does not permit the operation.
   const mode = getSafetyMode();
   const isToolAllowed = isAllowed(mode, annotations);
+  const isReadOnly = !!annotations.readOnlyHint;
 
-  // If not allowed, wrap handler to return an error and update description
-  let finalHandler = handler;
+  let finalHandler: ToolHandler = handler;
   let finalDescription = options.description;
 
   if (!isToolAllowed) {
     finalDescription = `[DISABLED in ${mode} mode] ${options.description}`;
-    finalHandler = async () => ({
+    finalHandler = async (): Promise<BlockedContent> => ({
       content: [
         {
           type: 'text',
@@ -75,9 +117,94 @@ export function registerToolSafe(
         },
       ],
       isError: true,
+      _lightdashBlocked: true,
+    });
+  } else if (isDryRunMode() && !isReadOnly) {
+    // ── Dry-run wrapper ─────────────────────────────────────────────────────
+    // Write operations are simulated; no API calls are made.
+    finalDescription = `[DRY-RUN] ${options.description}`;
+    finalHandler = async (args): Promise<BlockedContent> => ({
+      content: [
+        {
+          type: 'text',
+          text: `[DRY-RUN] Tool '${name}' would be called with: ${JSON.stringify(args, null, 2)}. No changes were made.`,
+        },
+      ],
+      _lightdashBlocked: true,
     });
   }
 
+  // ── Project allowlist wrapper ─────────────────────────────────────────────
+  // Reject calls targeting project UUIDs not in the configured allowlist.
+  // Covers both singular (projectUuid) and plural (projectUuids[]) arg shapes.
+  // CLI --allowed-projects takes priority over LIGHTDASH_TOOLS_ALLOWED_PROJECTS.
+  const allowedProjects = getAllowedProjectUuids();
+  if (allowedProjects.length > 0) {
+    const innerHandler = finalHandler;
+    finalHandler = async (args, extra): Promise<TextContent> => {
+      const projectUuids = extractProjectUuids(args);
+      const deniedUuids = projectUuids.filter(
+        (uuid) => !areAllProjectsAllowed(allowedProjects, [uuid]),
+      );
+      if (deniedUuids.length > 0) {
+        return {
+          content: [
+            {
+              type: 'text',
+              text: `Error: Project(s) [${deniedUuids.join(', ')}] are not in the list of allowed projects. Allowed: [${allowedProjects.join(', ')}].`,
+            },
+          ],
+          isError: true,
+          _lightdashBlocked: true,
+        } as BlockedContent;
+      }
+      return innerHandler(args, extra);
+    };
+  }
+
+  // ── Audit log wrapper ─────────────────────────────────────────────────────
+  // Outermost layer: records timing and outcome for every call.
+  const auditedInner = finalHandler;
+  finalHandler = async (args, extra): Promise<TextContent> => {
+    const start = Date.now();
+    const projectUuids = extractProjectUuids(args);
+    let status: 'success' | 'error' | 'blocked' = 'success';
+    let result: TextContent;
+
+    try {
+      result = await auditedInner(args, extra);
+      if (isGuardrailBlocked(result)) {
+        status = 'blocked';
+      } else if (result.isError) {
+        status = 'error';
+      }
+    } catch (err) {
+      status = 'error';
+      logAuditEntry({
+        timestamp: new Date().toISOString(),
+        sessionId: getSessionId(),
+        tool: name,
+        projectUuids: projectUuids.length > 0 ? projectUuids : undefined,
+        status,
+        durationMs: Date.now() - start,
+      });
+      throw err;
+    }
+
+    logAuditEntry({
+      timestamp: new Date().toISOString(),
+      sessionId: getSessionId(),
+      tool: name,
+      projectUuids: projectUuids.length > 0 ? projectUuids : undefined,
+      status,
+      durationMs: Date.now() - start,
+    });
+
+    // Strip the internal marker before returning to the MCP client.
+    const { content, isError } = result;
+    return { content, isError };
+  };
+
   const mergedOptions: ToolOptions = {
     ...options,
     description: finalDescription,