@useatlas/create 0.0.1

package/templates/nextjs-standalone/src/lib/tools/__tests__/custom-validation.test.ts

@@ -0,0 +1,238 @@
+/**
+ * Tests for custom query validation (non-SQL datasource plugins).
+ *
+ * When a datasource plugin registers a `validate` function on its connection,
+ * executeSQL uses that validator instead of the standard SQL validation pipeline.
+ */
+
+import { describe, expect, it, beforeEach, mock, type Mock } from "bun:test";
+
+// eslint-disable-next-line @typescript-eslint/no-explicit-any
+type AnyResult = any;
+
+// --- Mock dependencies ---
+
+mock.module("@atlas/api/lib/semantic", () => ({
+  getWhitelistedTables: () => new Set(["companies", "people"]),
+  _resetWhitelists: () => {},
+}));
+
+mock.module("@atlas/api/lib/auth/audit", () => ({
+  logQueryAudit: () => {},
+}));
+
+mock.module("@atlas/api/lib/tracing", () => ({
+  withSpan: async (
+    _name: string,
+    _attrs: Record<string, unknown>,
+    fn: () => Promise<unknown>,
+  ) => fn(),
+}));
+
+mock.module("@atlas/api/lib/db/source-rate-limit", () => ({
+  acquireSourceSlot: () => ({ acquired: true }),
+  decrementSourceConcurrency: () => {},
+}));
+
+mock.module("@atlas/api/lib/plugins/hooks", () => ({
+  dispatchHook: async () => {},
+  dispatchMutableHook: async (
+    _hookName: string,
+    ctx: { sql: string },
+    _field: string,
+  ) => ctx.sql,
+}));
+
+// --- Query mock ---
+
+let queryFn: Mock<(sql: string, timeout: number) => Promise<{ columns: string[]; rows: Record<string, unknown>[] }>>;
+
+const mockConn = {
+  query: (...args: [string, number]) => queryFn(...args),
+  close: async () => {},
+};
+
+// Custom validator: accepts SOQL-like "SELECT ... FROM ..." but rejects anything with DELETE/INSERT
+function soqlValidator(query: string): { valid: boolean; reason?: string } {
+  if (/\b(DELETE|INSERT|UPDATE)\b/i.test(query)) {
+    return { valid: false, reason: "SOQL only supports SELECT queries" };
+  }
+  return { valid: true };
+}
+
+// Connections mock with getValidator support
+let validatorMap: Map<string, ((q: string) => { valid: boolean; reason?: string }) | undefined>;
+
+mock.module("@atlas/api/lib/db/connection", () => ({
+  getDB: () => mockConn,
+  connections: {
+    get: () => mockConn,
+    getDefault: () => mockConn,
+    getDBType: () => "postgres",
+    getTargetHost: () => "localhost",
+    list: () => ["default", "salesforce-plugin"],
+    getValidator: (id: string) => validatorMap.get(id),
+  },
+  detectDBType: () => "postgres",
+}));
+
+// Import after mocks
+const { executeSQL } = await import("@atlas/api/lib/tools/sql");
+
+const exec = (sql: string, connectionId?: string) =>
+  executeSQL.execute!(
+    { sql, explanation: "test", connectionId },
+    { toolCallId: "test", messages: [], abortSignal: undefined as never },
+  ) as Promise<AnyResult>;
+
+describe("custom query validation", () => {
+  beforeEach(() => {
+    validatorMap = new Map();
+    queryFn = mock(() =>
+      Promise.resolve({ columns: ["id", "name"], rows: [{ id: 1, name: "Acme" }] }),
+    );
+  });
+
+  it("plugin with custom validator accepts valid queries", async () => {
+    validatorMap.set("salesforce-plugin", soqlValidator);
+
+    const result = await exec("SELECT Id, Name FROM Account", "salesforce-plugin");
+
+    expect(result.success).toBe(true);
+    expect(queryFn).toHaveBeenCalledTimes(1);
+  });
+
+  it("plugin with custom validator rejects invalid queries", async () => {
+    validatorMap.set("salesforce-plugin", soqlValidator);
+
+    const result = await exec("DELETE FROM Account WHERE Id = '001'", "salesforce-plugin");
+
+    expect(result.success).toBe(false);
+    expect(result.error).toContain("SOQL only supports SELECT queries");
+    expect(queryFn).not.toHaveBeenCalled();
+  });
+
+  it("falls back to validateSQL when no custom validator", async () => {
+    // No custom validator for default — standard SQL validation applies
+    validatorMap.set("default", undefined);
+
+    const result = await exec("SELECT * FROM companies");
+
+    expect(result.success).toBe(true);
+    expect(queryFn).toHaveBeenCalledTimes(1);
+  });
+
+  it("standard SQL validation rejects DML when no custom validator", async () => {
+    validatorMap.set("default", undefined);
+
+    const result = await exec("DROP TABLE companies");
+
+    expect(result.success).toBe(false);
+    expect(result.error).toContain("Forbidden SQL operation");
+    expect(queryFn).not.toHaveBeenCalled();
+  });
+
+  it("auto-LIMIT is skipped for custom-validated connections", async () => {
+    validatorMap.set("salesforce-plugin", soqlValidator);
+
+    await exec("SELECT Id FROM Account", "salesforce-plugin");
+
+    // The query passed to db.query should NOT have LIMIT appended
+    const executedQuery = queryFn.mock.calls[0]![0];
+    expect(executedQuery).not.toContain("LIMIT");
+  });
+
+  it("auto-LIMIT is applied for standard SQL connections", async () => {
+    validatorMap.set("default", undefined);
+
+    await exec("SELECT * FROM companies");
+
+    const executedQuery = queryFn.mock.calls[0]![0];
+    expect(executedQuery).toContain("LIMIT");
+  });
+
+  it("handles validator that throws an exception", async () => {
+    validatorMap.set("broken-plugin", () => { throw new Error("validator crashed"); });
+
+    const result = await exec("SELECT 1", "broken-plugin");
+
+    expect(result.success).toBe(false);
+    expect(result.error).toContain("internal validator error");
+    expect(queryFn).not.toHaveBeenCalled();
+  });
+
+  it("handles validator that returns undefined", async () => {
+    // eslint-disable-next-line @typescript-eslint/no-explicit-any
+    validatorMap.set("bad-plugin", (() => undefined) as any);
+
+    const result = await exec("SELECT 1", "bad-plugin");
+
+    expect(result.success).toBe(false);
+    expect(result.error).toContain("misconfigured");
+    expect(queryFn).not.toHaveBeenCalled();
+  });
+
+  it("handles validator that returns non-boolean valid field", async () => {
+    // eslint-disable-next-line @typescript-eslint/no-explicit-any
+    validatorMap.set("bad-valid", (() => ({ valid: "yes" })) as any);
+
+    const result = await exec("SELECT 1", "bad-valid");
+
+    expect(result.success).toBe(false);
+    expect(result.error).toContain("misconfigured");
+    expect(queryFn).not.toHaveBeenCalled();
+  });
+
+  it("normalizes SQL before passing to custom validator (trims whitespace and trailing semicolons)", async () => {
+    let receivedQuery = "";
+    validatorMap.set("trim-check", (q: string) => {
+      receivedQuery = q;
+      return { valid: true };
+    });
+
+    await exec("  SELECT Id FROM Account ;  ", "trim-check");
+
+    expect(receivedQuery).toBe("SELECT Id FROM Account");
+    expect(queryFn).toHaveBeenCalledTimes(1);
+  });
+
+  it("custom validator used for re-validation after hook mutation", async () => {
+    // Install a custom validator that rejects queries containing "FORBIDDEN"
+    const strictValidator = (q: string): { valid: boolean; reason?: string } => {
+      if (q.includes("FORBIDDEN")) return { valid: false, reason: "contains forbidden keyword" };
+      return { valid: true };
+    };
+    validatorMap.set("custom-conn", strictValidator);
+
+    // The hook rewrites to include "FORBIDDEN" — custom validator should catch it
+    mock.module("@atlas/api/lib/plugins/hooks", () => ({
+      dispatchHook: async () => {},
+      dispatchMutableHook: async (
+        _hookName: string,
+        _ctx: { sql: string },
+        _field: string,
+      ) => "SELECT FORBIDDEN FROM data",
+    }));
+
+    // Re-import to pick up the new mock
+    const { executeSQL: executeSQLv2 } = await import("@atlas/api/lib/tools/sql");
+
+    const result = await (executeSQLv2.execute!(
+      { sql: "SELECT Id FROM data", explanation: "query data", connectionId: "custom-conn" },
+      { toolCallId: "t7", messages: [], abortSignal: undefined as never },
+    ) as Promise<AnyResult>);
+
+    expect(result.success).toBe(false);
+    expect(result.error).toContain("contains forbidden keyword");
+
+    // Restore original mock
+    mock.module("@atlas/api/lib/plugins/hooks", () => ({
+      dispatchHook: async () => {},
+      dispatchMutableHook: async (
+        _hookName: string,
+        ctx: { sql: string },
+        _field: string,
+      ) => ctx.sql,
+    }));
+  });
+});

package/templates/nextjs-standalone/src/lib/tools/__tests__/explore-backend.test.ts

@@ -0,0 +1,267 @@
+import { describe, expect, it, beforeEach, afterEach, spyOn } from "bun:test";
+
+// ---------------------------------------------------------------------------
+// Helpers: We need fresh module state for each test since explore.ts has
+// module-level variables (_nsjailAvailable, _nsjailFailed, backendPromise).
+// We use dynamic imports with cache-busting to get fresh module instances.
+// ---------------------------------------------------------------------------
+
+let testCounter = 0;
+
+/** Import a fresh copy of explore.ts with all module state reset. */
+async function freshExploreModule() {
+  // Bun caches modules by resolved path. We can't bust the cache directly,
+  // so we rely on mock.module and re-import. Instead, we'll test the exported
+  // functions by understanding their stateful behavior and resetting env vars.
+  //
+  // Since we can't easily reset module-level let bindings from outside,
+  // we'll structure tests to work with the module's caching behavior.
+  testCounter++;
+  // Use a unique query param to bust the module cache
+  const mod = await import(
+    `@atlas/api/lib/tools/explore?t=${testCounter}`
+  );
+  return mod;
+}
+
+// ---------------------------------------------------------------------------
+// Tests for useNsjail / getExploreBackendType / getExploreBackend
+// ---------------------------------------------------------------------------
+
+describe("explore backend selection", () => {
+  const originalEnv = { ...process.env };
+
+  beforeEach(() => {
+    // Clean env for each test
+    delete process.env.ATLAS_RUNTIME;
+    delete process.env.VERCEL;
+    delete process.env.ATLAS_SANDBOX;
+    delete process.env.ATLAS_SANDBOX_URL;
+    delete process.env.ATLAS_NSJAIL_PATH;
+  });
+
+  afterEach(() => {
+    process.env = { ...originalEnv };
+  });
+
+  describe("useNsjail via getExploreBackendType", () => {
+    it("returns 'nsjail' when ATLAS_SANDBOX=nsjail (regardless of binary)", async () => {
+      process.env.ATLAS_SANDBOX = "nsjail";
+      const mod = await freshExploreModule();
+      expect(mod.getExploreBackendType()).toBe("nsjail");
+    });
+
+    it("returns 'nsjail' when nsjail binary is available on PATH", async () => {
+      // Set a PATH that has nsjail available
+      process.env.ATLAS_NSJAIL_PATH = "/usr/local/bin/nsjail";
+      // Mock fs.accessSync to report nsjail exists
+      const fs = await import("fs");
+      const spy = spyOn(fs, "accessSync").mockImplementation(() => {});
+
+      const mod = await freshExploreModule();
+      expect(mod.getExploreBackendType()).toBe("nsjail");
+
+      spy.mockRestore();
+    });
+
+    it("returns 'just-bash' when nsjail is not available", async () => {
+      delete process.env.ATLAS_SANDBOX;
+      delete process.env.ATLAS_NSJAIL_PATH;
+      process.env.PATH = "/usr/bin:/bin";
+      // Mock fs.accessSync to always throw (nsjail not found)
+      const fs = await import("fs");
+      const spy = spyOn(fs, "accessSync").mockImplementation(() => {
+        throw Object.assign(new Error("ENOENT"), { code: "ENOENT" });
+      });
+
+      const mod = await freshExploreModule();
+      expect(mod.getExploreBackendType()).toBe("just-bash");
+
+      spy.mockRestore();
+    });
+
+    it("returns 'vercel-sandbox' when ATLAS_RUNTIME=vercel", async () => {
+      process.env.ATLAS_RUNTIME = "vercel";
+      const mod = await freshExploreModule();
+      expect(mod.getExploreBackendType()).toBe("vercel-sandbox");
+    });
+
+    it("returns 'vercel-sandbox' when VERCEL env is set", async () => {
+      process.env.VERCEL = "1";
+      const mod = await freshExploreModule();
+      expect(mod.getExploreBackendType()).toBe("vercel-sandbox");
+    });
+
+    it("vercel-sandbox takes priority over nsjail", async () => {
+      process.env.ATLAS_RUNTIME = "vercel";
+      process.env.ATLAS_SANDBOX = "nsjail";
+      const mod = await freshExploreModule();
+      expect(mod.getExploreBackendType()).toBe("vercel-sandbox");
+    });
+  });
+
+  describe("useSidecar via getExploreBackendType", () => {
+    it("returns 'sidecar' when ATLAS_SANDBOX_URL is set and nsjail unavailable", async () => {
+      process.env.ATLAS_SANDBOX_URL = "http://localhost:8080";
+      delete process.env.ATLAS_NSJAIL_PATH;
+      process.env.PATH = "/usr/bin:/bin";
+      const fs = await import("fs");
+      const spy = spyOn(fs, "accessSync").mockImplementation(() => {
+        throw Object.assign(new Error("ENOENT"), { code: "ENOENT" });
+      });
+
+      const mod = await freshExploreModule();
+      expect(mod.getExploreBackendType()).toBe("sidecar");
+
+      spy.mockRestore();
+    });
+
+    it("sidecar takes priority over nsjail auto-detect when ATLAS_SANDBOX_URL is set", async () => {
+      process.env.ATLAS_SANDBOX_URL = "http://localhost:8080";
+      process.env.ATLAS_NSJAIL_PATH = "/usr/local/bin/nsjail";
+      const fs = await import("fs");
+      const spy = spyOn(fs, "accessSync").mockImplementation(() => {});
+
+      const mod = await freshExploreModule();
+      expect(mod.getExploreBackendType()).toBe("sidecar");
+
+      spy.mockRestore();
+    });
+
+    it("explicit nsjail (ATLAS_SANDBOX=nsjail) still beats sidecar", async () => {
+      process.env.ATLAS_SANDBOX = "nsjail";
+      process.env.ATLAS_SANDBOX_URL = "http://localhost:8080";
+      const mod = await freshExploreModule();
+      expect(mod.getExploreBackendType()).toBe("nsjail");
+    });
+
+    it("nsjail auto-detect works when no ATLAS_SANDBOX_URL is set", async () => {
+      delete process.env.ATLAS_SANDBOX_URL;
+      process.env.ATLAS_NSJAIL_PATH = "/usr/local/bin/nsjail";
+      const fs = await import("fs");
+      const spy = spyOn(fs, "accessSync").mockImplementation(() => {});
+
+      const mod = await freshExploreModule();
+      expect(mod.getExploreBackendType()).toBe("nsjail");
+
+      spy.mockRestore();
+    });
+
+    it("vercel-sandbox takes priority over sidecar", async () => {
+      process.env.ATLAS_RUNTIME = "vercel";
+      process.env.ATLAS_SANDBOX_URL = "http://localhost:8080";
+      const mod = await freshExploreModule();
+      expect(mod.getExploreBackendType()).toBe("vercel-sandbox");
+    });
+  });
+
+  describe("nsjail failure handling", () => {
+    it("falls back to just-bash after nsjail init failure (_nsjailFailed)", async () => {
+      // We can't directly set _nsjailFailed, but we can trigger the fallback
+      // by having nsjail available but failing to initialize.
+      // The getExploreBackendType checks _nsjailFailed flag.
+      // Since each fresh import resets the flag, we need to trigger the
+      // failure path through getExploreBackend first.
+
+      process.env.ATLAS_NSJAIL_PATH = "/usr/local/bin/nsjail";
+      const fs = await import("fs");
+      const spy = spyOn(fs, "accessSync").mockImplementation(
+        (p: import("fs").PathLike) => {
+          const path = String(p);
+          if (path === "/usr/local/bin/nsjail") return;
+          // Semantic root not readable — causes createNsjailBackend to throw
+          throw Object.assign(new Error("ENOENT"), { code: "ENOENT" });
+        },
+      );
+
+      const mod = await freshExploreModule();
+      // Initially nsjail is detected as available
+      expect(mod.getExploreBackendType()).toBe("nsjail");
+
+      // Trigger the explore tool to attempt nsjail init (it will fail)
+      // We access the internal getExploreBackend via the tool's execute.
+      // The tool.execute wraps getExploreBackend and returns error string on failure.
+      const result = await mod.explore.execute(
+        { command: "ls" },
+        { toolCallId: "test", messages: [], abortSignal: new AbortController().signal },
+      );
+      // The error should mention the nsjail failure or backend issue
+      expect(typeof result).toBe("string");
+
+      // After failure with auto-detected nsjail (not ATLAS_SANDBOX=nsjail),
+      // _nsjailFailed is set to true, so it falls back to just-bash
+      // But the next getExploreBackendType should reflect just-bash
+      // Note: this only works if the module state is shared
+      // Since _nsjailFailed is set in getExploreBackend's catch block
+      expect(mod.getExploreBackendType()).toBe("just-bash");
+
+      spy.mockRestore();
+    });
+
+    it("throws (does NOT fall back) when ATLAS_SANDBOX=nsjail and binary is missing", async () => {
+      process.env.ATLAS_SANDBOX = "nsjail";
+      process.env.PATH = "";
+      delete process.env.ATLAS_NSJAIL_PATH;
+
+      const fs = await import("fs");
+      const spy = spyOn(fs, "accessSync").mockImplementation(() => {
+        throw Object.assign(new Error("ENOENT"), { code: "ENOENT" });
+      });
+
+      const mod = await freshExploreModule();
+      // getExploreBackendType says nsjail because ATLAS_SANDBOX=nsjail
+      // bypasses binary check
+      expect(mod.getExploreBackendType()).toBe("nsjail");
+
+      // But actually trying to create the backend will fail with a throw,
+      // NOT fall back to just-bash
+      const result = await mod.explore.execute(
+        { command: "ls" },
+        { toolCallId: "test", messages: [], abortSignal: new AbortController().signal },
+      );
+      expect(typeof result).toBe("string");
+      expect(result).toContain("nsjail was explicitly requested");
+
+      spy.mockRestore();
+    });
+  });
+
+  describe("useNsjail unexpected error logging", () => {
+    it("logs unexpected errors that are not MODULE_NOT_FOUND", async () => {
+      // We need to make the require("./explore-nsjail") throw a non-MODULE_NOT_FOUND error.
+      // This is hard to test with real modules, but we can verify the code path exists
+      // by checking that _nsjailAvailable is set to false on unexpected errors.
+      // The console.error call is the observable side effect.
+
+      const consoleSpy = spyOn(console, "error").mockImplementation(() => {});
+
+      // Remove ATLAS_SANDBOX so useNsjail goes through auto-detection
+      delete process.env.ATLAS_SANDBOX;
+      delete process.env.ATLAS_NSJAIL_PATH;
+      process.env.PATH = "";
+
+      const fs = await import("fs");
+      const fsSpy = spyOn(fs, "accessSync").mockImplementation(() => {
+        throw Object.assign(new Error("ENOENT"), { code: "ENOENT" });
+      });
+
+      const mod = await freshExploreModule();
+      // This should exercise the useNsjail auto-detection path.
+      // Since nsjail is not found, it returns just-bash.
+      const result = mod.getExploreBackendType();
+      expect(result).toBe("just-bash");
+
+      consoleSpy.mockRestore();
+      fsSpy.mockRestore();
+    });
+  });
+
+  describe("invalidateExploreBackend", () => {
+    it("clears cached backend so next call recreates it", async () => {
+      const mod = await freshExploreModule();
+      // Just verify the function exists and is callable
+      expect(typeof mod.invalidateExploreBackend).toBe("function");
+      mod.invalidateExploreBackend(); // Should not throw
+    });
+  });
+});