Skip to main content
Glama
siddhashutosh

AegisQuery

README.md
# AegisQuery β€” a governed analytics MCP server

> A production-grade **Model Context Protocol** server that gives LLM agents **safe, read-only** access to an analytics data warehouse β€” mediated by a layered safety pipeline (AST validation, column/row governance, PII masking, cost limits, result truncation, audit) over both **stdio** and an **OAuth-secured Streamable HTTP** transport.

Built in TypeScript against the official `@modelcontextprotocol/sdk`, targeting MCP spec revision **2025-11-25** and designed **stateless-first** for the direction the protocol is heading.

Its differentiator is the security story: it reproduces the class of bug that defeated Anthropic's *own* reference Postgres MCP server β€” a stacked-statement injection (`COMMIT; DROP SCHEMA public CASCADE;`) that bypassed a naive `READ ONLY` transaction β€” and then **defeats it** with structural AST validation plus defense-in-depth. See [`docs/`](docs/) for the full SRS / HLD / LLD / Design documents.

## πŸ”΄ Live demo

**[web-beta-rosy-0ovpsr1q8m.vercel.app](https://web-beta-rosy-0ovpsr1q8m.vercel.app)** β€” a coffee-themed web console (light + dark) that runs the **same governance engine** on Vercel serverless. Switch roles (analyst / regional / finance / admin) and watch the *same* query behave differently: PII masked for an analyst but raw for an admin, `cost` denied for an analyst but visible to finance, and adversarial SQL rejected on the spot. Source in [`web/`](web/).

---

## Quickstart

```bash
npm install
npm run seed          # build a local SQLite warehouse (deterministic synthetic data)
npm run smoke         # end-to-end verification incl. an adversarial SQL corpus
npm run eval          # golden-task scorecard vs an independent oracle

# run the server
npm run build && npm start          # stdio (local, env-key auth)
npm start -- --http                 # Streamable HTTP (remote, bearer/OAuth auth)
```

Point an MCP client (e.g. Claude Desktop, MCP Inspector) at `node dist/index.js`, with `AEGIS_API_KEY=devkey` in the environment.

## What it exposes

**Tools** (workflow-oriented, not endpoint mirrors):

| Tool | Purpose |
|------|---------|
| `search_catalog` | Rank tables/columns relevant to a natural-language question (with access levels). |
| `revenue_by_period` | Curated metric β€” gross/net revenue by day/week/month. No model-authored SQL. |
| `preview_query_cost` | Validate a query and return plan + cost + violations **without** executing. |
| `run_analytics_query` | Governed ad-hoc read-only `SELECT`. Full safety pipeline applies. |

**Resources** β€” the warehouse schema (`aegis://catalog`, `aegis://schema/{table}`), annotated per-policy with `allow` / `mask` / `deny`.
**Prompts** β€” `analyze_question`, steering agents to discover schema, prefer curated metrics, and treat results as data.

## The safety pipeline

Every query passes the same ordered stages before a single row is returned:

```
parse (AST) β†’ read-only guard β†’ reference extraction β†’ table/column allowlist
  β†’ PII masking β†’ row-policy injection β†’ cost guard (LIMIT/scan budget)
  β†’ execute (read-only role) β†’ truncate to token budget β†’ audit
```

- **AST read-only guard** β€” only a *single* `SELECT`/`WITH`/`UNION` passes; multiple statements parse to an array and are rejected. This structurally prevents the stacked-statement bypass.
- **Defense in depth** β€” the SQLite connection is opened read-only, so even a parser gap cannot mutate data.
- **Governance** β€” deny-by-default table/column allowlists; PII columns masked (`substr(x,1,2)||'***'`) and refused if used outside the projection; restricted columns (cost, url) require an elevated role.
- **Row-level policy** β€” mandatory `WHERE` predicates injected from principal attributes (e.g. a `regional_analyst` is scoped to their region).
- **Cost & context** β€” planner-based scan estimate rejects runaway queries; every result is `LIMIT`-capped and truncated to a ~25k-token budget.
- **Audit** β€” one structured JSON record per invocation (principal, raw + rewritten SQL, decision, rows, ms).

## Auth

- **stdio** β€” credentials from the environment (`AEGIS_API_KEY`), per the spec (stdio servers should not use OAuth).
- **HTTP** β€” bearer token validated and mapped to a role/policy; unauthenticated requests get `401` with RFC 9728 Protected Resource Metadata. Stateless (no session id) so it scales behind a round-robin load balancer.

Key map format (`AEGIS_API_KEYS`): `key:role[:attr=val;attr=val]`, e.g. `devkey:analyst, eu-key:regional_analyst:region=EU, admin-key:admin`.

## Configuration

| Env var | Default | Meaning |
|---------|---------|---------|
| `AEGIS_TRANSPORT` / `--http` | `stdio` | Transport selection. |
| `AEGIS_DB_PATH` | `data/warehouse.db` | SQLite warehouse path. |
| `AEGIS_MAX_ROWS` | `200` | Hard row cap (injected LIMIT). |
| `AEGIS_MAX_SCAN_ROWS` | `500000` | Reject queries whose plan exceeds this. |
| `AEGIS_RESPONSE_TOKEN_BUDGET` | `25000` | Result truncation budget. |
| `AEGIS_API_KEYS` | `devkey:analyst` | Principal key map. |
| `AEGIS_HTTP_PORT` | `3333` | HTTP bind port. |

## Architecture

```
transports/{stdio,http}  β†’  server (McpServer)  β†’  tools ─┐
                                    β”‚                      β”œβ”€ pipeline ─ sql/analyzer (AST)
                                auth (principal)           β”‚              policy (governance)
                                catalog (resources)        β”‚              db/adapter β†’ sqlite
                                                           └─ audit
```

The warehouse sits behind a read-only `ReadOnlyAdapter` interface β€” SQLite locally, Postgres in production β€” so governance logic never touches a specific engine.

## Learning-ordered roadmap

This repo is built as a learning vehicle; each milestone is independently demoable and teaches one MCP / AI-engineering competency.

| M | Build | Learn |
|---|-------|-------|
| M0 | stdio server + a tool | JSON-RPC lifecycle, tools/listΒ·call |
| M1 | read-only query + schema resources | resources vs tools, zod I/O, structured content |
| M2 | reproduce the CVE β†’ AST + read-only enforcement | why appsec applies; defense-in-depth |
| M3 | allowlist, PII masking, row policy | data governance |
| M4 | cost guard, truncation, pagination | context engineering, token economics |
| M5 | curated metrics + catalog search | Anthropic tool-authoring principles |
| M6 | audit log, error taxonomy | production ops/observability |
| M7 | Streamable HTTP + OAuth 2.1, stateless | remote MCP, auth |
| M8 | golden-task eval harness | eval β€” the top hiring signal |

## Verification

```bash
npm run typecheck   # strict TS, src + scripts
npm run smoke       # 20 assertions incl. 9 adversarial rejections + integrity check
npm run eval        # dev + held-out scorecard (accuracy, tokens, latency, errors)
```

## Documentation

Full design suite (generated, colorful `.docx`) in [`docs/`](docs/): **SRS**, **HLD**, **LLD**, **Design**. Regenerate with `npm run docs`. All four are driven by a single source of truth in `scripts/docgen/spec.ts`.

## License

MIT.