mcp-data-agent 0.1.0

package/packages/core/src/query-safety.ts

@@ -0,0 +1,125 @@
+/**
+ * MCP Data Agent — @mcp-data-agent/core
+ * Created by Arslan Habib
+ *
+ * ---------------------------------------------------------------------------
+ * FILE: query-safety.ts
+ * PURPOSE:
+ *   Defense-in-depth SQL validation before any query reaches a database
+ *   adapter. Blocks writes, multi-statements, locking clauses, and caps
+ *   result size / timeout.
+ *
+ * IMPORTANT:
+ *   Regex validation is NOT enough alone. Host apps should also:
+ *   - Use read-only DB credentials for AI analytics
+ *   - Use bound parameters (never string-concat AI text into SQL)
+ *   - Enforce tenant isolation at the DB (e.g. Postgres RLS)
+ * ---------------------------------------------------------------------------
+ */
+
+import { QueryValidationError } from "./errors.js";
+
+/**
+ * Keywords that indicate a mutating or dangerous statement.
+ * Matched as whole words (case-insensitive).
+ */
+const FORBIDDEN_KEYWORDS =
+  /\b(insert|update|delete|drop|alter|truncate|create|grant|revoke|copy|call|execute|merge|replace|attach|detach|pragma|vacuum|reindex|cluster|comment|security|owner|load|import|export|into\s+outfile|load_file)\b/i;
+
+/** Detects a second statement after a semicolon (SQL injection pattern). */
+const MULTI_STATEMENT = /;\s*\S/;
+
+/** Result of a successful validation pass. */
+export interface ValidatedQuery {
+  sql: string;
+  limit: number;
+  timeoutMs: number;
+}
+
+/** Inputs accepted by validateReadOnlyQuery(). */
+export interface ValidateQueryOptions {
+  sql: string;
+  limit?: number;
+  timeoutMs?: number;
+  defaultLimit?: number;
+  defaultTimeoutMs?: number;
+  maxLimit?: number;
+  maxTimeoutMs?: number;
+}
+
+/**
+ * Validate that a query is a single read-only SELECT (or WITH … SELECT).
+ * Throws QueryValidationError on anything unsafe.
+ */
+export function validateReadOnlyQuery(
+  options: ValidateQueryOptions,
+): ValidatedQuery {
+  const raw = options.sql?.trim() ?? "";
+  if (!raw) {
+    throw new QueryValidationError("Query SQL is required");
+  }
+
+  // Block "SELECT 1; DROP TABLE ..." style attacks
+  if (MULTI_STATEMENT.test(raw)) {
+    throw new QueryValidationError("Multiple SQL statements are not allowed");
+  }
+
+  // Allow a single trailing semicolon; strip it for further checks
+  const sql = raw.replace(/;\s*$/, "");
+
+  if (FORBIDDEN_KEYWORDS.test(sql)) {
+    throw new QueryValidationError(
+      "Only read-only SELECT queries are allowed by default",
+    );
+  }
+
+  // Must start with SELECT or WITH (common table expressions)
+  const normalized = sql.replace(/^\s*\(/, "").trimStart();
+  const isSelect =
+    /^select\b/i.test(normalized) ||
+    /^with\b/i.test(normalized);
+
+  if (!isSelect) {
+    throw new QueryValidationError(
+      "Query must be a SELECT or WITH (CTE) statement",
+    );
+  }
+
+  // SELECT INTO can create tables — blocked in Phase 1
+  if (/\binto\b/i.test(sql) && !/\binto\s+temp(orary)?\b/i.test(sql)) {
+    throw new QueryValidationError("SELECT INTO is not allowed");
+  }
+
+  // Row locks are not needed for analytics and can block writers
+  if (/\bfor\s+update\b/i.test(sql) || /\block\s+in\s+share\b/i.test(sql)) {
+    throw new QueryValidationError("Locking clauses are not allowed");
+  }
+
+  // Clamp limit and timeout to safe ranges
+  const maxLimit = options.maxLimit ?? 1000;
+  const maxTimeout = options.maxTimeoutMs ?? 30_000;
+  const defaultLimit = options.defaultLimit ?? 100;
+  const defaultTimeout = options.defaultTimeoutMs ?? 10_000;
+
+  const limit = Math.min(
+    Math.max(1, options.limit ?? defaultLimit),
+    maxLimit,
+  );
+  const timeoutMs = Math.min(
+    Math.max(100, options.timeoutMs ?? defaultTimeout),
+    maxTimeout,
+  );
+
+  return { sql, limit, timeoutMs };
+}
+
+/**
+ * Ensure a LIMIT clause exists without rewriting complex SQL aggressively.
+ * If the query already has LIMIT, leave it; otherwise append one.
+ */
+export function ensureLimit(sql: string, limit: number): string {
+  if (/\blimit\s+\d+/i.test(sql)) {
+    return sql;
+  }
+  return `${sql} LIMIT ${limit}`;
+}

package/packages/core/src/schema-safety.ts

@@ -0,0 +1,118 @@
+/**
+ * MCP Data Agent — @mcp-data-agent/core
+ * Created by Arslan Habib
+ *
+ * ---------------------------------------------------------------------------
+ * FILE: schema-safety.ts
+ * PURPOSE:
+ *   Prevents sensitive columns (passwords, tokens, salaries, etc.) from
+ *   reaching the LLM. Works in two stages:
+ *     1) mark columns as sensitive (patterns + developer protectField list)
+ *     2) redact them from schema and from query result rows
+ *
+ * WHY:
+ *   Even a valid SELECT might return password_hash. We strip those keys
+ *   before returning data to AI tools.
+ * ---------------------------------------------------------------------------
+ */
+
+import type {
+  ColumnSchema,
+  DatabaseSchema,
+  ProtectFieldOptions,
+  TableSchema,
+} from "./types.js";
+
+/**
+ * Heuristic name patterns treated as sensitive by default.
+ * Developers can add more via agent.protectField().
+ */
+const DEFAULT_SENSITIVE_PATTERNS = [
+  /password/i,
+  /passwd/i,
+  /secret/i,
+  /api[_-]?key/i,
+  /access[_-]?token/i,
+  /refresh[_-]?token/i,
+  /private[_-]?key/i,
+  /credit[_-]?card/i,
+  /card[_-]?number/i,
+  /ssn/i,
+  /salary/i,
+  /password_hash/i,
+];
+
+/** Builds a stable key: "schema.table.column" (lowercased). */
+function fieldKey(schema: string | undefined, table: string, column: string): string {
+  return `${schema ?? "public"}.${table}.${column}`.toLowerCase();
+}
+
+/** Returns true if the column name looks sensitive based on default patterns. */
+export function isLikelySensitiveColumn(columnName: string): boolean {
+  return DEFAULT_SENSITIVE_PATTERNS.some((pattern) => pattern.test(columnName));
+}
+
+/** Converts protectField() options into a Set for O(1) lookups. */
+export function buildProtectedFieldSet(
+  fields: ProtectFieldOptions[] = [],
+): Set<string> {
+  return new Set(
+    fields.map((f) => fieldKey(f.schema, f.table, f.column)),
+  );
+}
+
+/**
+ * Walks the schema and sets `sensitive: true` on columns that match
+ * patterns or the developer's protected-field list.
+ */
+export function markSensitiveColumns(
+  schema: DatabaseSchema,
+  protectedFields: Set<string>,
+): DatabaseSchema {
+  const tables = schema.tables.map((table) => {
+    const columns = table.columns.map((col) => {
+      const key = fieldKey(table.schema, table.name, col.name);
+      const sensitive =
+        col.sensitive === true ||
+        protectedFields.has(key) ||
+        isLikelySensitiveColumn(col.name);
+      return sensitive ? { ...col, sensitive: true } : col;
+    });
+    return { ...table, columns };
+  });
+  return { ...schema, tables };
+}
+
+/**
+ * Returns a copy of the schema with sensitive columns removed entirely.
+ * This is what database.schema returns to the AI.
+ */
+export function redactSchemaForAi(schema: DatabaseSchema): DatabaseSchema {
+  const tables: TableSchema[] = schema.tables.map((table) => {
+    const columns: ColumnSchema[] = table.columns.filter((c) => !c.sensitive);
+    return { ...table, columns };
+  });
+  return { ...schema, tables };
+}
+
+/**
+ * Removes sensitive keys from each result row (by column name, case-insensitive).
+ * Applied after database.query so leaked SELECT columns never reach the LLM.
+ */
+export function redactResultRows(
+  rows: Record<string, unknown>[],
+  sensitiveColumnNames: Set<string>,
+): Record<string, unknown>[] {
+  if (sensitiveColumnNames.size === 0) {
+    return rows;
+  }
+  return rows.map((row) => {
+    const next: Record<string, unknown> = {};
+    for (const [key, value] of Object.entries(row)) {
+      if (!sensitiveColumnNames.has(key.toLowerCase())) {
+        next[key] = value;
+      }
+    }
+    return next;
+  });
+}

package/packages/core/src/tenant.ts

@@ -0,0 +1,58 @@
+/**
+ * MCP Data Agent — @mcp-data-agent/core
+ * Created by Arslan Habib
+ *
+ * ---------------------------------------------------------------------------
+ * FILE: tenant.ts
+ * PURPOSE:
+ *   Multi-tenant safety helpers. SaaS apps isolate data by tenant_id; this
+ *   module ensures tenant context is present and that AI/tool arguments
+ *   cannot spoof or override which tenant is in scope.
+ *
+ * RULE:
+ *   tenantId comes ONLY from the authenticated application session.
+ * ---------------------------------------------------------------------------
+ */
+
+import { TenantIsolationError } from "./errors.js";
+import type { TenantContext } from "./types.js";
+
+/** Default column name used by many SaaS schemas for row-level tenancy. */
+export const DEFAULT_TENANT_COLUMN = "tenant_id";
+
+/**
+ * Validates that tenant.tenantId is a non-empty string.
+ * Called on every tool invocation via MCPDataAgent.getContext().
+ */
+export function assertTenantContext(tenant: TenantContext): void {
+  if (!tenant.tenantId || typeof tenant.tenantId !== "string") {
+    throw new TenantIsolationError(
+      "tenantId must be derived from the authenticated session",
+    );
+  }
+  if (tenant.tenantId.trim().length === 0) {
+    throw new TenantIsolationError("tenantId cannot be empty");
+  }
+}
+
+/** Returns the configured tenant column name, or the default "tenant_id". */
+export function getTenantColumn(tenant: TenantContext): string {
+  return tenant.tenantColumn ?? DEFAULT_TENANT_COLUMN;
+}
+
+/**
+ * Reject AI/tool attempts to override tenant identity via parameters.
+ * Example attack: { sql: "...", tenantId: "other_gym" } — this throws.
+ */
+export function rejectTenantOverride(
+  params: Record<string, unknown>,
+  reservedKeys: string[] = ["tenantId", "tenant_id", "tenant"],
+): void {
+  for (const key of reservedKeys) {
+    if (Object.prototype.hasOwnProperty.call(params, key)) {
+      throw new TenantIsolationError(
+        `Parameter "${key}" is not allowed; tenant context comes from the session`,
+      );
+    }
+  }
+}

package/packages/core/src/tool-registry.ts

@@ -0,0 +1,61 @@
+/**
+ * MCP Data Agent — @mcp-data-agent/core
+ * Created by Arslan Habib
+ *
+ * ---------------------------------------------------------------------------
+ * FILE: tool-registry.ts
+ * PURPOSE:
+ *   In-memory registry of tools the agent can invoke. Built-in tools
+ *   (database.schema, database.query) and custom tools (members.search, …)
+ *   are stored here by unique name.
+ *
+ * RESPONSIBILITIES:
+ *   - Prevent duplicate tool names
+ *   - Reject mutating tools that misuse the read-only database.query permission
+ * ---------------------------------------------------------------------------
+ */
+
+import { ConfigurationError } from "./errors.js";
+import type { ToolDefinition } from "./types.js";
+
+/** Map of tool name → ToolDefinition. Owned by MCPDataAgent. */
+export class ToolRegistry {
+  private readonly tools = new Map<string, ToolDefinition>();
+
+  /**
+   * Registers a new tool. Throws if the name is missing or already taken.
+   * Mutating tools must use an explicit write permission (not database.query).
+   */
+  register<TInput = unknown, TOutput = unknown>(
+    tool: ToolDefinition<TInput, TOutput>,
+  ): void {
+    if (!tool.name?.trim()) {
+      throw new ConfigurationError("Tool name is required");
+    }
+    if (this.tools.has(tool.name)) {
+      throw new ConfigurationError(`Tool already registered: ${tool.name}`);
+    }
+    // database.query is reserved for read-only analytics
+    if (tool.mutates && tool.permission === "database.query") {
+      throw new ConfigurationError(
+        "Mutating tools must use an explicit write permission, not database.query",
+      );
+    }
+    this.tools.set(tool.name, tool as ToolDefinition);
+  }
+
+  /** Lookup a single tool by name (undefined if not registered). */
+  get(name: string): ToolDefinition | undefined {
+    return this.tools.get(name);
+  }
+
+  /** All registered tools (used by MCP bridge to expose custom tools). */
+  list(): ToolDefinition[] {
+    return [...this.tools.values()];
+  }
+
+  /** Returns true if a tool with this name exists. */
+  has(name: string): boolean {
+    return this.tools.has(name);
+  }
+}

package/packages/core/src/types.ts

@@ -0,0 +1,250 @@
+/**
+ * MCP Data Agent — @mcp-data-agent/core
+ * Created by Arslan Habib
+ *
+ * ---------------------------------------------------------------------------
+ * FILE: types.ts
+ * PURPOSE:
+ *   Shared TypeScript contracts for the whole framework. Every package
+ *   (postgres, mcp, host apps) depends on these shapes so adapters, tools,
+ *   and the agent speak the same language.
+ *
+ * WHAT LIVES HERE:
+ *   - Auth / tenant / request context (who is calling, which tenant)
+ *   - Database schema + query request/result types
+ *   - DatabaseAdapter interface (what every DB driver must implement)
+ *   - Tool definitions, audit events, entities, agent options
+ *
+ * SECURITY NOTES:
+ *   - AuthContext and TenantContext must come from the host session.
+ *   - Never populate tenantId from AI/tool arguments.
+ * ---------------------------------------------------------------------------
+ */
+
+/** Permission string, e.g. "members.read" or "database.query". */
+export type Permission = string;
+
+/**
+ * Who is authenticated for this request.
+ * Provided by the host application after it verifies the user session.
+ */
+export interface AuthContext {
+  /** Authenticated application user id — never trust AI-supplied values. */
+  userId: string;
+  /** Optional display name for audit logs. */
+  displayName?: string;
+  /** Roles assigned by the host application. */
+  roles?: string[];
+  /** Effective permission set for this session. */
+  permissions: Permission[];
+  /** Opaque session identifier from the host app. */
+  sessionId?: string;
+  /** Additional host-app claims (non-sensitive). */
+  claims?: Record<string, string | number | boolean>;
+}
+
+/**
+ * Which SaaS tenant this request belongs to.
+ * Must be derived from the authenticated session — never from the LLM.
+ */
+export interface TenantContext {
+  /** Tenant id derived from the authenticated session — never from AI params. */
+  tenantId: string;
+  /** Column used for tenant isolation (default: tenant_id). */
+  tenantColumn?: string;
+}
+
+/**
+ * Full per-request context passed into tools and adapters.
+ * Combines identity (auth) + tenancy + optional question/correlation ids.
+ */
+export interface RequestContext {
+  auth: AuthContext;
+  tenant: TenantContext;
+  /** Correlation id for audit trails. */
+  requestId?: string;
+  /** Original natural-language question when available. */
+  question?: string;
+}
+
+/** Normalized column types exposed to AI-facing schema (DB-agnostic). */
+export type ColumnDataType =
+  | "string"
+  | "number"
+  | "boolean"
+  | "date"
+  | "datetime"
+  | "json"
+  | "binary"
+  | "unknown";
+
+/** One column in a table schema returned by an adapter. */
+export interface ColumnSchema {
+  name: string;
+  dataType: ColumnDataType;
+  /** Original DB type name (e.g. "timestamptz", "varchar"). */
+  nativeType: string;
+  nullable: boolean;
+  isPrimaryKey: boolean;
+  isForeignKey: boolean;
+  /** Present when this column references another table. */
+  references?: {
+    table: string;
+    column: string;
+  };
+  description?: string;
+  /** When true, field is withheld from AI-facing schema/tools. */
+  sensitive?: boolean;
+}
+
+/** One table (or collection) in the inspected database schema. */
+export interface TableSchema {
+  name: string;
+  /** DB schema/namespace, e.g. "public" in PostgreSQL. */
+  schema?: string;
+  columns: ColumnSchema[];
+  description?: string;
+  primaryKey?: string[];
+}
+
+/** Full schema snapshot returned by DatabaseAdapter.getSchema(). */
+export interface DatabaseSchema {
+  tables: TableSchema[];
+  /** e.g. "postgresql", "mysql". */
+  dialect: string;
+  /** ISO timestamp of when the schema was inspected. */
+  inspectedAt: string;
+}
+
+/** Phase 1 only allows read (SELECT) operations by default. */
+export type QueryOperation = "select";
+
+/**
+ * A query the agent asks an adapter to run.
+ * SQL must already be validated as read-only before reaching the adapter.
+ */
+export interface QueryRequest {
+  /** Validated SQL (or dialect query) — must be read-only by default. */
+  sql: string;
+  /** Bound parameters — never interpolate user/AI text into SQL. */
+  params?: unknown[];
+  /** Max rows returned to the caller. */
+  limit?: number;
+  /** Query timeout in milliseconds. */
+  timeoutMs?: number;
+}
+
+export interface QueryResultColumn {
+  name: string;
+  dataType: ColumnDataType;
+}
+
+/** Structured result returned to tools / MCP / the LLM. */
+export interface QueryResult {
+  columns: QueryResultColumn[];
+  rows: Record<string, unknown>[];
+  rowCount: number;
+  /** True if more rows existed than we returned (limit applied). */
+  truncated: boolean;
+  executionMs?: number;
+}
+
+/**
+ * DatabaseAdapter — pluggable DB driver contract.
+ * Postgres/MySQL/MongoDB adapters all implement this so the core stays
+ * database-agnostic.
+ */
+export interface DatabaseAdapter {
+  readonly name: string;
+  getSchema(options?: { includeSensitive?: boolean }): Promise<DatabaseSchema>;
+  query(request: QueryRequest, context: RequestContext): Promise<QueryResult>;
+  /** Optional cleanup (e.g. close connection pool). */
+  close?(): Promise<void>;
+}
+
+/**
+ * A named capability the AI can call (built-in or custom).
+ * Each tool declares the permission required to run it.
+ */
+export interface ToolDefinition<TInput = unknown, TOutput = unknown> {
+  name: string;
+  description: string;
+  /** Permission required to invoke this tool (e.g. database.query). */
+  permission: Permission;
+  /** Whether the tool may mutate data. Default false. */
+  mutates?: boolean;
+  /** Optional JSON-Schema-like description of inputs. */
+  inputSchema?: Record<string, unknown>;
+  handler: (input: TInput, context: RequestContext) => Promise<TOutput>;
+}
+
+/**
+ * One audit-log entry for an AI tool invocation.
+ * Never put passwords, tokens, or connection strings in these fields.
+ */
+export interface AuditEvent {
+  timestamp: string;
+  requestId?: string;
+  userId: string;
+  tenantId: string;
+  question?: string;
+  tool: string;
+  parameters?: unknown;
+  status: "success" | "denied" | "error";
+  errorMessage?: string;
+  resultMeta?: {
+    rowCount?: number;
+    truncated?: boolean;
+    executionMs?: number;
+  };
+}
+
+/** Host-provided callback that receives audit events. */
+export type AuditLogger = (event: AuditEvent) => void | Promise<void>;
+
+/**
+ * Business meaning for a table — helps the AI understand domain language
+ * (e.g. "Member" maps to table "members").
+ */
+export interface EntityDefinition {
+  name: string;
+  table: string;
+  description?: string;
+  schema?: string;
+  relationships?: Array<{
+    name: string;
+    entity: string;
+    type: "one-to-one" | "one-to-many" | "many-to-one" | "many-to-many";
+  }>;
+}
+
+/** Marks a column that must never be shown to AI tools. */
+export interface ProtectFieldOptions {
+  table: string;
+  column: string;
+  schema?: string;
+  reason?: string;
+}
+
+/**
+ * Options passed to createMCPDataAgent().
+ * This is the main configuration surface for host applications.
+ */
+export interface AgentOptions {
+  database: DatabaseAdapter;
+  /**
+   * Resolve auth + tenant for each request.
+   * Tenant MUST come from the authenticated session, never from AI parameters.
+   */
+  resolveContext: () => Promise<RequestContext> | RequestContext;
+  /** Default max rows for analytical queries. */
+  defaultLimit?: number;
+  /** Default query timeout in ms. */
+  defaultTimeoutMs?: number;
+  /** Fields that must never be exposed to AI tools. */
+  protectedFields?: ProtectFieldOptions[];
+  /** Optional audit sink. Secrets must never be logged. */
+  audit?: AuditLogger;
+  /** Allow write tools only when explicitly enabled. Default false. */
+  allowMutations?: boolean;
+}

package/packages/core/tsconfig.json

@@ -0,0 +1,10 @@
+{
+  "extends": "../../tsconfig.base.json",
+  "compilerOptions": {
+    "rootDir": "src",
+    "outDir": "dist",
+    "composite": true
+  },
+  "include": ["src/**/*.ts"],
+  "exclude": ["src/**/*.test.ts", "dist"]
+}

package/packages/mcp/README.md

@@ -0,0 +1,7 @@
+# @mcp-data-agent/mcp
+
+MCP server bridge for **MCP Data Agent**.
+
+**Created by Arslan Habib**
+
+Uses `@modelcontextprotocol/server` (MCP TypeScript SDK v2).

package/packages/mcp/package.json

@@ -0,0 +1,46 @@
+{
+  "name": "@mcp-data-agent/mcp",
+  "version": "0.1.0",
+  "description": "MCP server bridge for MCP Data Agent. Created by Arslan Habib.",
+  "author": {
+    "name": "Arslan Habib"
+  },
+  "creator": "Arslan Habib",
+  "license": "UNLICENSED",
+  "type": "module",
+  "main": "./dist/index.js",
+  "types": "./dist/index.d.ts",
+  "exports": {
+    ".": {
+      "types": "./dist/index.d.ts",
+      "import": "./dist/index.js"
+    }
+  },
+  "files": [
+    "dist",
+    "README.md"
+  ],
+  "scripts": {
+    "build": "tsc -p tsconfig.json",
+    "clean": "rm -rf dist *.tsbuildinfo",
+    "typecheck": "tsc -p tsconfig.json --noEmit"
+  },
+  "engines": {
+    "node": ">=20"
+  },
+  "peerDependencies": {
+    "@mcp-data-agent/core": "0.1.0",
+    "@modelcontextprotocol/server": "^2.0.0",
+    "zod": "^3.24.0 || ^4.0.0"
+  },
+  "devDependencies": {
+    "@mcp-data-agent/core": "*",
+    "@modelcontextprotocol/server": "^2.0.0",
+    "zod": "^4.0.0"
+  },
+  "keywords": [
+    "mcp",
+    "model-context-protocol",
+    "mcp-data-agent"
+  ]
+}

package/packages/mcp/src/index.ts

@@ -0,0 +1,15 @@
+/**
+ * MCP Data Agent — @mcp-data-agent/mcp
+ * Created by Arslan Habib
+ *
+ * ---------------------------------------------------------------------------
+ * FILE: index.ts
+ * PURPOSE:
+ *   Public entry point for the MCP bridge package.
+ *
+ *     import { createMcpDataAgentServer } from "@mcp-data-agent/mcp";
+ * ---------------------------------------------------------------------------
+ */
+
+export { createMcpDataAgentServer } from "./server.js";
+export type { CreateMcpServerOptions } from "./server.js";