Plugin Authoring Guide

bun create @useatlas/plugin my-datasource --type datasource

bun run atlas -- plugin create my-datasource --type datasource

my-datasource/
├── src/
│   ├── index.ts         # Plugin entry point
│   └── index.test.ts    # Test scaffold
├── package.json
├── tsconfig.json
└── README.md

// src/config.ts
import { z } from "zod";

export const ConfigSchema = z.object({
  url: z
    .string()
    .min(1, "URL must not be empty")
    .refine(
      (u) => u.startsWith("postgresql://") || u.startsWith("postgres://"),
      "URL must start with postgresql:// or postgres://",
    ),
  poolSize: z.number().int().positive().max(500).optional(),
});

export type PluginConfig = z.infer<typeof ConfigSchema>;

// src/connection.ts
import type { PluginDBConnection, PluginQueryResult } from "@useatlas/plugin-sdk";
import type { PluginConfig } from "./config";

export function createConnection(config: PluginConfig): PluginDBConnection {
  let Pool: typeof import("pg").Pool;
  try {
    ({ Pool } = require("pg"));
  } catch (err) {
    const isNotFound =
      err != null &&
      typeof err === "object" &&
      "code" in err &&
      (err as NodeJS.ErrnoException).code === "MODULE_NOT_FOUND";
    if (isNotFound) {
      throw new Error("This plugin requires the pg package. Install it with: bun add pg");
    }
    throw err;
  }

  const pool = new Pool({
    connectionString: config.url,
    max: config.poolSize ?? 10,
  });

  return {
    async query(sql: string, timeoutMs?: number): Promise<PluginQueryResult> {
      const client = await pool.connect();
      try {
        if (timeoutMs) {
          await client.query(`SET statement_timeout = ${timeoutMs}`);
        }
        const result = await client.query(sql);
        return {
          columns: result.fields.map((f) => f.name),
          rows: result.rows,
        };
      } finally {
        client.release();
      }
    },

    async close(): Promise<void> {
      await pool.end();
    },
  };
}

// src/index.ts
import { createPlugin } from "@useatlas/plugin-sdk";
import type { AtlasDatasourcePlugin, PluginHealthResult } from "@useatlas/plugin-sdk";
import { ConfigSchema, type PluginConfig } from "./config";
import { createConnection } from "./connection";

export function buildPlugin(config: PluginConfig): AtlasDatasourcePlugin<PluginConfig> {
  let cachedConnection: ReturnType<typeof createConnection> | undefined;

  return {
    id: "my-datasource",
    types: ["datasource"] as const,
    version: "1.0.0",
    name: "My DataSource",
    config,

    connection: {
      create: () => {
        if (!cachedConnection) {
          cachedConnection = createConnection(config);
        }
        return cachedConnection;
      },
      dbType: "postgres",
    },

    entities: [],

    dialect: "This datasource uses PostgreSQL. Use DATE_TRUNC() for date truncation.",

    // Called once during server startup. Throw to block startup (for fatal configuration errors).
    async initialize(ctx) {
      ctx.logger.info("My datasource plugin initialized");
    },

    // Called by `atlas doctor` and the admin API. Always return a result — never throw.
    // Return `{ healthy: false, message: '...' }` for recoverable issues.
    async healthCheck(): Promise<PluginHealthResult> {
      const start = performance.now();
      try {
        const conn = createConnection(config);
        await conn.query("SELECT 1", 5000);
        await conn.close();
        return { healthy: true, latencyMs: Math.round(performance.now() - start) };
      } catch (err) {
        return {
          healthy: false,
          message: err instanceof Error ? err.message : String(err),
          latencyMs: Math.round(performance.now() - start),
        };
      }
    },
  };
}

export const myPlugin = createPlugin({
  configSchema: ConfigSchema,
  create: buildPlugin,
});

import { defineConfig } from "@atlas/api/lib/config";
import { myPlugin } from "./plugins/my-datasource/src/index";

export default defineConfig({
  plugins: [
    myPlugin({ url: process.env.MY_DB_URL! }),
  ],
});

bun test plugins/my-datasource/src/index.test.ts

{
  "name": "atlas-plugin-my-datasource",
  "peerDependencies": {
    "@useatlas/plugin-sdk": ">=0.0.1",
    "pg": ">=8.0.0"
  },
  "peerDependenciesMeta": {
    "pg": { "optional": true }
  },
  "devDependencies": {
    "@useatlas/plugin-sdk": "^0.0.2"
  }
}

import { describe, test, expect } from "bun:test";
import { myPlugin } from "./index";

describe("my-datasource plugin", () => {
  test("validates config schema", () => {
    // Test that invalid config is rejected
    expect(() => myPlugin({ url: "" })).toThrow();
  });

  test("creates plugin with valid config", () => {
    const plugin = myPlugin({ url: "postgresql://localhost/test" });
    expect(plugin.id).toBe("my-datasource");
    expect(plugin.type).toBe("datasource");
  });

  test("health check reports status", async () => {
    const plugin = myPlugin({ url: "postgresql://localhost/test" });
    const health = await plugin.healthCheck?.();
    expect(health).toHaveProperty("healthy");
  });
});

import { definePlugin } from "@useatlas/plugin-sdk";

export default definePlugin({
  id: "company-glossary",
  types: ["context"],
  version: "1.0.0",
  name: "Company Glossary",

  contextProvider: {
    // Cache the loaded context to avoid re-reading on every request
    _cache: null as string | null,

    async load() {
      if (this._cache) return this._cache;

      // Load from any source: filesystem, database, API, etc.
      const terms = [
        { term: "ARR", definition: "Annual Recurring Revenue — sum of all active subscription values annualized" },
        { term: "MRR", definition: "Monthly Recurring Revenue — ARR / 12" },
        { term: "churn", definition: "Percentage of customers who cancel within a billing period" },
      ];

      const lines = terms.map((t) => `- **${t.term}**: ${t.definition}`);
      this._cache = `## Company Glossary\n\n${lines.join("\n")}`;
      return this._cache;
    },

    async refresh() {
      // Clear cache so next load() re-reads from source
      this._cache = null;
    },
  },

  async initialize(ctx) {
    ctx.logger.info("Company glossary context plugin initialized");
  },
});

export default definePlugin({
  id: "my-webhook",
  types: ["interaction"],
  version: "1.0.0",
  routes(app) {
    app.post("/webhooks/my-service", async (c) => {
      return c.json({ ok: true });
    });
  },
});

import { z } from "zod";
import { createPlugin } from "@useatlas/plugin-sdk";
import { tool } from "@useatlas/plugin-sdk/ai";
import type { AtlasActionPlugin, PluginAction } from "@useatlas/plugin-sdk";

const ticketConfigSchema = z.object({
  apiUrl: z.string().url(),
  apiKey: z.string().min(1, "apiKey must not be empty"),
  defaultPriority: z.enum(["low", "medium", "high"]).default("medium"),
});

type TicketConfig = z.infer<typeof ticketConfigSchema>;

export const ticketPlugin = createPlugin<TicketConfig, AtlasActionPlugin<TicketConfig>>({
  configSchema: ticketConfigSchema,

  create(config) {
    const action: PluginAction = {
      name: "createSupportTicket",
      description: "Create a support ticket from analysis findings",
      tool: tool({
        description: "Create a support ticket. Requires user approval before execution.",
        inputSchema: z.object({
          title: z.string().max(200).describe("Short summary of the issue"),
          body: z.string().describe("Detailed description with relevant data"),
          priority: z
            .enum(["low", "medium", "high"])
            .optional()
            .describe(`Priority level. Defaults to "${config.defaultPriority}"`),
        }),
        execute: async ({ title, body, priority }) => {
          // This only runs AFTER the user approves in the chat UI
          const response = await fetch(`${config.apiUrl}/tickets`, {
            method: "POST",
            headers: {
              Authorization: `Bearer ${config.apiKey}`,
              "Content-Type": "application/json",
            },
            body: JSON.stringify({
              title,
              body,
              priority: priority ?? config.defaultPriority,
            }),
          });

          if (!response.ok) {
            throw new Error(`Ticket API returned ${response.status}`);
          }

          const ticket = (await response.json()) as { id: string; url: string };
          return { ticketId: ticket.id, url: ticket.url };
        },
      }),
      actionType: "ticket:create",
      reversible: false,
      defaultApproval: "manual",
      requiredCredentials: ["apiKey"],
      // ^ Values must match environment variable names (e.g. process.env.apiKey).
      // At startup, Atlas checks these env vars exist and logs a warning for any
      // that are missing (see validateActionCredentials in the ToolRegistry).
      // Missing credentials do not block startup — they produce warnings only.
    };

    return {
      id: "ticket-action",
      types: ["action"] as const,
      version: "1.0.0",
      name: "Support Ticket Action",
      config,
      actions: [action],
    };
  },
});

plugins: [
  ticketPlugin({
    apiUrl: process.env.TICKET_API_URL!,
    apiKey: process.env.TICKET_API_KEY!,
  }),
],

sandbox: {
  create(semanticRoot: string): PluginExploreBackend {
    return {
      async exec(command: string) {
        // Execute command in isolation, return { stdout, stderr, exitCode }
      },
      async close() { /* cleanup */ },
    };
  },
  priority: 60,
},
security: {
  networkIsolation: true,
  filesystemIsolation: true,
  unprivilegedExecution: true,
  description: "My isolation mechanism...",
},

import { createPlugin } from "@useatlas/plugin-sdk";
import { z } from "zod";

export const myPlugin = createPlugin({
  configSchema: z.object({ url: z.string().url() }),
  create: (config) => ({
    id: "my-plugin",
    types: ["datasource"] as const,
    version: "1.0.0",
    config,
    connection: { create: () => makeConnection(config.url), dbType: "postgres" },
  }),
});

// Usage in atlas.config.ts:
plugins: [myPlugin({ url: process.env.MY_URL! })]

import { definePlugin } from "@useatlas/plugin-sdk";

export default definePlugin({
  id: "my-context",
  types: ["context"],
  version: "1.0.0",
  contextProvider: {
    async load() { return "Additional context for the agent"; },
  },
});

// Usage in atlas.config.ts:
import myContext from "./plugins/my-context";
plugins: [myContext]

import type { $InferServerPlugin } from "@useatlas/plugin-sdk";
import type { clickhousePlugin } from "@useatlas/clickhouse";

type CH = $InferServerPlugin<typeof clickhousePlugin>;
// CH["Config"] → { url: string; database?: string }
// CH["Type"]   → "datasource"
// CH["Id"]     → string
// CH["DbType"] → "clickhouse"

export default definePlugin({
  id: "audit-logger",
  types: ["context"],
  version: "1.0.0",
  contextProvider: { async load() { return ""; } },

  hooks: {
    beforeQuery: [{
      matcher: (ctx) => ctx.sql.includes("sensitive_table"),
      handler: (ctx) => {
        console.log(`Query on sensitive table: ${ctx.sql}`);
        // Return { sql } to rewrite, throw to reject, or return void to pass through
      },
    }],
    afterQuery: [{
      handler: (ctx) => {
        console.log(`Query completed in ${ctx.durationMs}ms, ${ctx.result.rows.length} rows`);
      },
    }],
  },
});

export default definePlugin({
  id: "my-plugin",
  types: ["context"],
  version: "1.0.0",
  schema: {
    my_plugin_cache: {
      fields: {
        key: { type: "string", required: true, unique: true },
        value: { type: "string", required: true },
        updated_at: { type: "date" },
      },
    },
  },
  // ...
});

connection: { create: () => myConn, dbType: "postgres" },
entities: [
  { name: "users", yaml: "table: users\ndimensions:\n  id:\n    type: number" },
],
// Or as an async factory:
entities: async () => {
  const tables = await discoverTables();
  return tables.map(t => ({ name: t.name, yaml: generateYaml(t) }));
},

dialect: "This datasource uses ClickHouse. Use toStartOfMonth() for date truncation, not DATE_TRUNC().",

validate?(query: string): QueryValidationResult | Promise<QueryValidationResult>

interface QueryValidationResult {
  valid: boolean;
  /** User-facing rejection reason — appears in error responses and audit logs. */
  reason?: string;
}

connection: {
  create: () => mySalesforceConn,
  dbType: "salesforce",
  validate(query) {
    if (query.length > 20_000) {
      return { valid: false, reason: "SOQL query exceeds 20,000 character limit" };
    }
    if (/\b(DELETE|INSERT|UPDATE|UPSERT)\b/i.test(query)) {
      return { valid: false, reason: "Only SELECT queries are allowed" };
    }
    return { valid: true };
  },
},

connection: {
  create: () => myConn,
  dbType: "custom-api",
  async validate(query) {
    const res = await fetch("https://schema.internal/validate", {
      method: "POST",
      body: JSON.stringify({ query }),
      headers: { "Content-Type": "application/json" },
      signal: AbortSignal.timeout(5000),
    });
    if (!res.ok) return { valid: false, reason: "Schema service unavailable" };
    const body = await res.json() as { allowed: boolean; message?: string };
    return body.allowed
      ? { valid: true }
      : { valid: false, reason: body.message ?? "Query rejected" };
  },
},

connection: {
  create: () => myConn,
  dbType: "snowflake",
  // Override auto-detected parser dialect (case-sensitive, e.g. "Snowflake" not "snowflake")
  parserDialect: "Snowflake",
  // Additional regex patterns to block beyond the base DML/DDL guard
  forbiddenPatterns: [/\bCOPY\s+INTO\b/i, /\bPUT\b/i],
},

async teardown() {
  await this.pool.end();
},

async initialize(ctx) {
  ctx.tools.register({
    name: "lookupInventory",
    description: "Check inventory levels for a product SKU",
    tool: tool({
      description: "Look up current inventory by SKU",
      inputSchema: z.object({ sku: z.string() }),
      execute: async ({ sku }) => fetchInventory(sku),
    }),
  });
},

async healthCheck(): Promise<PluginHealthResult> {
  const start = performance.now();
  let conn: PluginDBConnection | undefined;
  try {
    conn = createConnection(config);
    await conn.query("SELECT 1", 5000); // timeout via query parameter
    return { healthy: true, latencyMs: Math.round(performance.now() - start) };
  } catch (err) {
    return {
      healthy: false,
      message: err instanceof Error ? err.message : String(err),
      latencyMs: Math.round(performance.now() - start),
    };
  } finally {
    if (conn) await conn.close();
  }
}

async healthCheck(): Promise<PluginHealthResult> {
  const start = performance.now();
  try {
    const response = await fetch("https://api.example.com/health", {
      headers: { Authorization: `Bearer ${config.apiKey}` },
      signal: AbortSignal.timeout(5000), // 5s timeout
    });
    const latencyMs = Math.round(performance.now() - start);
    if (response.ok) return { healthy: true, latencyMs };
    return { healthy: false, message: `API returned ${response.status}`, latencyMs };
  } catch (err) {
    return {
      healthy: false,
      message: err instanceof Error ? err.message : String(err),
      latencyMs: Math.round(performance.now() - start),
    };
  }
}

async healthCheck(): Promise<PluginHealthResult> {
  const start = performance.now();
  const TIMEOUT = 30_000; // sandbox creation can be slow
  let sandbox: SandboxInstance | null = null; // hoist for cleanup on timeout
  let timer: ReturnType<typeof setTimeout>;
  try {
    const result = await Promise.race([
      (async () => {
        sandbox = await createSandbox(config); // your SDK's create method
        await sandbox.close(); // cleanup method varies by SDK (kill, stop, delete, etc.)
        sandbox = null;
        return "ok" as const;
      })(),
      new Promise<"timeout">((resolve) => {
        timer = setTimeout(() => resolve("timeout"), TIMEOUT);
      }),
    ]).finally(() => clearTimeout(timer!)); // always clean up the timer
    const latencyMs = Math.round(performance.now() - start);
    if (result === "timeout") {
      // Best-effort cleanup — sandbox may still be creating
      if (sandbox) {
        try { await sandbox.close(); } catch { /* best-effort */ }
      }
      return { healthy: false, message: `Timed out after ${TIMEOUT}ms`, latencyMs };
    }
    return { healthy: true, latencyMs };
  } catch (err) {
    if (sandbox) {
      try { await sandbox.close(); } catch { /* best-effort */ }
    }
    return {
      healthy: false,
      message: err instanceof Error ? err.message : String(err),
      latencyMs: Math.round(performance.now() - start),
    };
  }
}

async healthCheck(): Promise<PluginHealthResult> {
  const start = performance.now();
  try {
    // Verify local resources exist
    const files = fs.readdirSync(dir).filter((f) => f.endsWith(".yml"));
    const latencyMs = Math.round(performance.now() - start);
    if (files.length === 0) {
      return { healthy: false, message: "No entity files found", latencyMs };
    }
    return { healthy: true, latencyMs };
  } catch (err) {
    return {
      healthy: false,
      message: err instanceof Error ? err.message : String(err),
      latencyMs: Math.round(performance.now() - start),
    };
  }
}

// Good
myPlugin({ apiKey: process.env.MY_API_KEY! })

// Bad -- hidden dependency on env var name
// inside plugin: process.env.MY_API_KEY