@useatlas/create 0.0.5 → 0.0.7

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

@@ -1,240 +0,0 @@
-/**
- * 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),
-    getParserDialect: () => undefined,
-    getForbiddenPatterns: () => [],
-  },
-  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/docker/src/lib/tools/__tests__/explore-backend.test.ts

@@ -1,267 +0,0 @@
-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
-    });
-  });
-});