Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 5 additions & 0 deletions .changeset/mcp-instructions-allowlist.md
Original file line number Diff line number Diff line change
@@ -0,0 +1,5 @@
---
"@stainless-code/codemap": patch
---

Add MCP initialize server instructions (tool-selection playbook) and `CODEMAP_MCP_TOOLS` env for subset tool registration.
2 changes: 1 addition & 1 deletion README.md
Original file line number Diff line number Diff line change
Expand Up @@ -234,7 +234,7 @@ codemap agents init --force
codemap agents init --interactive # -i; IDE wiring + symlink vs copy
```

**Environment / flags:** `--root` overrides **`CODEMAP_ROOT`** / **`CODEMAP_TEST_BENCH`**, then **`process.cwd()`**; **`--state-dir`** overrides **`CODEMAP_STATE_DIR`** (default `.codemap/`); **`CODEMAP_WATCH=0`** opts out of the default-ON watcher on `mcp` / `serve` (mirrors `--no-watch`). Indexing a project outside this clone: [docs/benchmark.md § Indexing another project](docs/benchmark.md#indexing-another-project).
**Environment / flags:** `--root` overrides **`CODEMAP_ROOT`** / **`CODEMAP_TEST_BENCH`**, then **`process.cwd()`**; **`--state-dir`** overrides **`CODEMAP_STATE_DIR`** (default `.codemap/`); **`CODEMAP_WATCH=0`** opts out of the default-ON watcher on `mcp` / `serve` (mirrors `--no-watch`); **`CODEMAP_MCP_TOOLS`** registers a subset of MCP tools (comma-separated snake_case names; see [agents.md § MCP tool allowlist](docs/agents.md#mcp-tool-allowlist)). Indexing a project outside this clone: [docs/benchmark.md § Indexing another project](docs/benchmark.md#indexing-another-project).

**Configuration:** optional **`<state-dir>/config.{ts,js,json}`** (default `.codemap/config.*`; default export object or async factory). Shape: [codemap.config.example.json](codemap.config.example.json). Runtime validation (**Zod**, strict keys) and API surface: [docs/architecture.md § User config](docs/architecture.md#user-config). When developing inside this repo you can use `defineConfig` from `@stainless-code/codemap` or `./src/config`. If you set **`include`**, it **replaces** the default glob list entirely. **Self-healing files (D11):** `<state-dir>/.gitignore` is rewritten to canonical on every codemap boot; JSON config gets unknown-key pruning + key-sort drift; TS/JS configs are validate-only.

Expand Down
20 changes: 19 additions & 1 deletion docs/agents.md
Original file line number Diff line number Diff line change
Expand Up @@ -78,7 +78,25 @@ Once `agents init` has written the pointer templates, the consumer's disk holds
| MCP | resource `codemap://skill` | resource `codemap://rule` |
| HTTP (`codemap serve`) | `GET /resources/{encoded uri}` against `codemap://skill` | `GET /resources/{encoded uri}` against `codemap://rule` |

All three transports resolve to the same `assembleAgentContent(kind)` function in `src/application/agent-content.ts` — there is no MCP-only or HTTP-only path for skill/rule content. The MCP and HTTP paths share a lazy per-process cache via `readResource()` in `src/application/resource-handlers.ts` for schema/skill/rule; recipes, files, and symbols read live every call. The CLI re-assembles every call (cheap — markdown read + concat).
All three transports resolve to the same `assembleAgentContent(kind)` function in `src/application/agent-content.ts` — there is no MCP-only or HTTP-only path for skill/rule content. The MCP and HTTP paths share a lazy per-process cache via `readResource()` in `src/application/resource-handlers.ts` for schema/skill/rule/mcp-instructions; recipes, files, and symbols read live every call. The CLI re-assembles every call (cheap — markdown read + concat).

## MCP server instructions

`codemap mcp` passes a tool-selection playbook in the MCP **`initialize`** response **`instructions`** field. MCP clients (Cursor, Claude Code, etc.) inject this into the agent system prompt — operational guidance only (which tool when, common chains, anti-patterns). Full schema and recipe catalog stay on **`codemap://skill`** / **`codemap://rule`**.

| Surface | URI / field |
| ------------------- | -------------------------------------------------------------------------------------------------------------- |
| MCP initialize | `instructions` on handshake |
| MCP / HTTP resource | `codemap://mcp-instructions` |
| Source file | `templates/agent-content/mcp-instructions.md` (assembled by `assembleMcpInstructions()` in `agent-content.ts`) |

Recipe ids cited in the playbook are machine-validated in tests against the live catalog (`extractMcpInstructionRecipeIds`).

## MCP tool allowlist

**`CODEMAP_MCP_TOOLS`** — comma-separated snake_case MCP tool names. When set, only listed tools register (stderr lists the active set). Unknown names are ignored with a warning. Unset = all tools (default). **`query_batch`** registers only when listed or when unset (eval ablation).

Example: `CODEMAP_MCP_TOOLS=query,context,show codemap mcp --no-watch`

## Section assembler and `*.gen.md`

Expand Down
2 changes: 1 addition & 1 deletion docs/packaging.md
Original file line number Diff line number Diff line change
Expand Up @@ -10,7 +10,7 @@ How **@stainless-code/codemap** is built and published. **Doc index:** [README.m
The `templates/` directory ships two parallel subtrees:

- **`templates/agents/`** — consumer-disk targets copied by `codemap agents init` (thin pointer files: ~16-line SKILL.md + ~22-line rule).
- **`templates/agent-content/`** — server-side source assembled live by `codemap skill` / `codemap rule` / `codemap://skill` / `codemap://rule`. Section files in `agent-content/skill/` concatenate in lexical order; `*.gen.md` files are replaced at fetch time by renderers in `src/application/agent-content.ts`. See [agents.md](./agents.md#section-assembler-and-genmd) for the split rationale.
- **`templates/agent-content/`** — server-side source assembled live by `codemap skill` / `codemap rule` / `codemap://skill` / `codemap://rule` / `codemap://mcp-instructions`. Section files in `agent-content/skill/` concatenate in lexical order; `*.gen.md` files are replaced at fetch time by renderers in `src/application/agent-content.ts`. Root-level `mcp-instructions.md` feeds MCP initialize `instructions`. See [agents.md](./agents.md#section-assembler-and-genmd) for the split rationale.
- **`templates/recipes/`** — bundled SQL recipe `.sql` + `.md` pairs (every recipe in `--recipes-json`).

## Consuming locally
Expand Down
19 changes: 19 additions & 0 deletions src/application/agent-content.ts
Original file line number Diff line number Diff line change
Expand Up @@ -39,6 +39,25 @@ export function resolveAgentContentDir(): string {
return join(resolveAgentsTemplateDir(), "..", "agent-content");
}

const MCP_INSTRUCTIONS_FILE = "mcp-instructions.md";
const MCP_RECIPE_REFS_RE = /<!--\s*codemap-mcp-recipe-refs:\s*([^>]+?)-->/;

/** MCP initialize playbook — `templates/agent-content/mcp-instructions.md`. */
export function assembleMcpInstructions(): string {
const path = join(resolveAgentContentDir(), MCP_INSTRUCTIONS_FILE);
return readFileSync(path, "utf8").trimEnd() + "\n";
}

/** Recipe ids declared in the MCP instructions machine-ref comment. */
export function extractMcpInstructionRecipeIds(content: string): string[] {
const match = content.match(MCP_RECIPE_REFS_RE);
if (match === null) return [];
return match[1]!
.split(",")
.map((id) => id.trim())
.filter(Boolean);
}

/**
* Renderer registry — keyed by `<kind>/<filename>`. Files ending in
* `.gen.md` are treated as generated content: if a renderer is
Expand Down
81 changes: 79 additions & 2 deletions src/application/mcp-server.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -33,8 +33,12 @@ afterEach(() => {
rmSync(benchDir, { recursive: true, force: true });
});

async function makeClient() {
const server = createMcpServer({ version: "0.0.0-test", root: benchDir });
async function makeClient(env?: NodeJS.ProcessEnv) {
const server = createMcpServer({
version: "0.0.0-test",
root: benchDir,
env,
});
const client = new Client({ name: "test-client", version: "0.0.0" });
const [clientTransport, serverTransport] =
InMemoryTransport.createLinkedPair();
Expand All @@ -54,6 +58,62 @@ function readJson(result: unknown): any {
return JSON.parse(first.text) as unknown;
}

describe("MCP server — initialize instructions", () => {
it("includes tool-selection playbook in initialize handshake", async () => {
const { client, server } = await makeClient();
try {
const instructions = client.getInstructions();
expect(instructions).toBeDefined();
expect(instructions!.length).toBeGreaterThan(500);
expect(instructions).toContain("Session start");
expect(instructions).toContain("codemap://rule");
} finally {
await server.close();
}
});

it("cites only shipped recipe ids", async () => {
const { assembleMcpInstructions, extractMcpInstructionRecipeIds } =
await import("./agent-content");
const { listQueryRecipeCatalog } = await import("./query-recipes");
const cited = extractMcpInstructionRecipeIds(assembleMcpInstructions());
expect(cited.length).toBeGreaterThan(0);
const catalog = new Set(listQueryRecipeCatalog().map((e) => e.id));
for (const id of cited) {
expect(catalog.has(id)).toBe(true);
}
});
});

describe("MCP server — tool allowlist", () => {
it("registers only listed tools when CODEMAP_MCP_TOOLS is set", async () => {
const { client, server } = await makeClient({
CODEMAP_MCP_TOOLS: "query,show",
});
try {
const tools = await client.listTools();
const names = tools.tools.map((t) => t.name).sort();
expect(names).toEqual(["query", "show"]);
} finally {
await server.close();
}
});

it("excludes query_batch unless explicitly listed", async () => {
const { client, server } = await makeClient({
CODEMAP_MCP_TOOLS: "query",
});
try {
const tools = await client.listTools();
const names = tools.tools.map((t) => t.name);
expect(names).toContain("query");
expect(names).not.toContain("query_batch");
} finally {
await server.close();
}
});
});

describe("MCP server — query tool", () => {
it("lists query and query_batch in tools/list", async () => {
const { client, server } = await makeClient();
Expand Down Expand Up @@ -717,6 +777,7 @@ describe("MCP server — resources", () => {
expect(uris).toContain("codemap://schema");
expect(uris).toContain("codemap://skill");
expect(uris).toContain("codemap://rule");
expect(uris).toContain("codemap://mcp-instructions");
// The recipe-by-id resource is a template — surfaced via list-template
// callback as one entry per recipe id.
const recipeUris = uris.filter((u) => u.startsWith("codemap://recipes/"));
Expand Down Expand Up @@ -815,6 +876,22 @@ describe("MCP server — resources", () => {
}
});

it("codemap://mcp-instructions returns the MCP playbook", async () => {
const { client, server } = await makeClient();
try {
const r = await client.readResource({
uri: "codemap://mcp-instructions",
});
const first = r.contents[0] as { mimeType?: string };
expect(first.mimeType).toBe("text/markdown");
const text = readResourceText(r);
expect(text).toContain("tool selection");
expect(text).toContain("Session start");
} finally {
await server.close();
}
});

it("codemap://files/{path} returns a per-file roll-up", async () => {
const { client, server } = await makeClient();
try {
Expand Down
63 changes: 46 additions & 17 deletions src/application/mcp-server.ts
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,13 @@ import {
getTsconfigPath,
initCodemap,
} from "../runtime";
import { assembleMcpInstructions } from "./agent-content";
import {
isMcpToolEnabled,
logMcpToolAllowlist,
resolveMcpToolAllowlist,
} from "./mcp-tool-allowlist";
import type { McpToolName } from "./mcp-tool-allowlist";
import { listQueryRecipeCatalog } from "./query-recipes";
import { readResource } from "./resource-handlers";
import type { ResourcePayload } from "./resource-handlers";
Expand Down Expand Up @@ -67,6 +74,8 @@ interface ServerOpts {
root: string;
configFile?: string | undefined;
stateDir?: string | undefined;
/** Test hook — defaults to `process.env`. */
env?: NodeJS.ProcessEnv | undefined;
/**
* If true, boot a co-process file watcher (chokidar via
* `runWatchLoop`) so the server's tools always read live data without
Expand Down Expand Up @@ -107,25 +116,39 @@ function wrapToolResult(r: ToolResult) {
* `InMemoryTransport.createLinkedPair()` for in-process driving).
*/
export function createMcpServer(opts: ServerOpts): McpServer {
const server = new McpServer({
name: "codemap",
version: opts.version,
});
const allowlistResolved = resolveMcpToolAllowlist(opts.env ?? process.env);
const server = new McpServer(
{
name: "codemap",
version: opts.version,
},
{
instructions: assembleMcpInstructions(),
},
);

const registered: McpToolName[] = [];
const maybeRegister = (name: McpToolName, register: () => void): void => {
if (!isMcpToolEnabled(name, allowlistResolved.allowlist)) return;
register();
registered.push(name);
};

registerQueryTool(server, opts);
registerQueryBatchTool(server, opts);
registerQueryRecipeTool(server, opts);
registerAuditTool(server);
registerContextTool(server);
registerValidateTool(server);
registerSaveBaselineTool(server, opts);
registerListBaselinesTool(server);
registerDropBaselineTool(server);
registerShowTool(server, opts);
registerSnippetTool(server, opts);
registerImpactTool(server);
registerApplyTool(server, opts);
maybeRegister("query", () => registerQueryTool(server, opts));
maybeRegister("query_batch", () => registerQueryBatchTool(server, opts));
maybeRegister("query_recipe", () => registerQueryRecipeTool(server, opts));
maybeRegister("audit", () => registerAuditTool(server));
maybeRegister("context", () => registerContextTool(server));
maybeRegister("validate", () => registerValidateTool(server));
maybeRegister("save_baseline", () => registerSaveBaselineTool(server, opts));
maybeRegister("list_baselines", () => registerListBaselinesTool(server));
maybeRegister("drop_baseline", () => registerDropBaselineTool(server));
maybeRegister("show", () => registerShowTool(server, opts));
maybeRegister("snippet", () => registerSnippetTool(server, opts));
maybeRegister("impact", () => registerImpactTool(server));
maybeRegister("apply", () => registerApplyTool(server, opts));
registerResources(server);
logMcpToolAllowlist(allowlistResolved, registered);

return server;
}
Expand Down Expand Up @@ -320,6 +343,12 @@ function registerResources(server: McpServer): void {
"codemap://rule",
"Full text of the bundled `templates/agents/rules/codemap.md` (always-on priming for agents working in this repo).",
);
registerStaticResource(
server,
"mcp-instructions",
"codemap://mcp-instructions",
"MCP initialize tool-selection playbook (operational guidance only; full catalog in codemap://skill).",
);

// codemap://recipes/{id} — one recipe (template form). Payload includes
// `body` / `source` / `shadows` from the catalog entry — session-start
Expand Down
55 changes: 55 additions & 0 deletions src/application/mcp-tool-allowlist.test.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,55 @@
import { describe, expect, it } from "bun:test";

import {
isMcpToolEnabled,
resolveMcpToolAllowlist,
} from "./mcp-tool-allowlist";

describe("mcp-tool-allowlist", () => {
it("returns null allowlist when env unset", () => {
expect(resolveMcpToolAllowlist({})).toEqual({
allowlist: null,
unknown: [],
});
});

it("returns null allowlist when env is whitespace", () => {
expect(resolveMcpToolAllowlist({ CODEMAP_MCP_TOOLS: " " })).toEqual({
allowlist: null,
unknown: [],
});
});

it("parses comma-separated tool names", () => {
const { allowlist, unknown } = resolveMcpToolAllowlist({
CODEMAP_MCP_TOOLS: "query, show",
});
expect(unknown).toEqual([]);
expect(allowlist).toEqual(new Set(["query", "show"]));
});

it("ignores unknown names without failing", () => {
const { allowlist, unknown } = resolveMcpToolAllowlist({
CODEMAP_MCP_TOOLS: "query,not_a_tool,show",
});
expect(unknown).toEqual(["not_a_tool"]);
expect(allowlist).toEqual(new Set(["query", "show"]));
});

it("query_batch is excluded unless explicitly listed", () => {
const { allowlist } = resolveMcpToolAllowlist({
CODEMAP_MCP_TOOLS: "query,show",
});
expect(isMcpToolEnabled("query_batch", allowlist)).toBe(false);
expect(
isMcpToolEnabled(
"query_batch",
resolveMcpToolAllowlist({ CODEMAP_MCP_TOOLS: "query_batch" }).allowlist,
),
).toBe(true);
});

it("isMcpToolEnabled allows all when allowlist is null", () => {
expect(isMcpToolEnabled("query_batch", null)).toBe(true);
});
});
Loading
Loading