feat(output-budget): adaptive caps for trace, explore, and node

.changeset/adaptive-output-budgets.md

@@ -0,0 +1,5 @@
+---
+"@stainless-code/codemap": patch
+---
+
+Adaptive snippet budgets for `trace`, `explore`, and `node` scale with indexed file count when `budget_chars` is omitted (explicit override unchanged). `explore` also applies adaptive row limits internally (no transport override).

docs/roadmap.md

@@ -79,7 +79,7 @@ Long-running MCP / HTTP sessions dominate agent workflows; one-shot CLI keeps th
 - [x] **Rich `context` composer** — `start_here` on non-compact `context`: intent-ranked recipe cards, inline index summary, hub leaders with signatures (adaptive caps), debug-biased markers, optional MCP/HTTP `include_snippets`. Shipped [#151](https://github.com/stainless-code/codemap/pull/151).
 - [ ] **Codebase map in bootstrap responses** — hash-stable structural summary (top hubs, CLI entry hints, schema version, index freshness) auto-included in `context` / MCP initialize payload. **Partial:** hubs + `start_here.index_summary` + `index_freshness` ship on `context`; CLI entry hints + hash-stable map id still open. Opt-out via flag. Effort: S–M.
 - [x] **Index staleness surfacing** — `index_freshness.pending_sync` on `context`, MCP tool metadata, and HTTP headers when the watcher debounce queue or in-flight reindex is active. Shipped [#149](https://github.com/stainless-code/codemap/pull/149).
-- [ ] **Adaptive output budgets** — scale explore/trace/node snippet char caps and row limits from indexed file/symbol counts so large trees do not blow token budgets. **Partial:** `context` hub/signature caps via `resolveContextBudget()` (2026-05). Effort: S.
+- [x] **Adaptive output budgets** — scale trace/explore/node snippet char caps (and explore row limits) from indexed file counts via **`resolveOutputBudget(file_count)`** in `output-budget.ts`. Shipped [#152](https://github.com/stainless-code/codemap/pull/152). **`context`** hub/signature caps remain in **`resolveContextBudget()`**.
 - [ ] **MCP session lifecycle hygiene** — idle timeout, client disconnect detection, graceful watcher shutdown on last client; avoid orphaned watchers after agent host crashes. Effort: S–M.
 - [ ] **`agents init` uninstall (teardown)** — symmetric inverse of init for failed pilots, template mistakes, or leaving a repo: remove codemap-managed MCP entries, pointer sections, and IDE symlinks only (same scoped paths as init; never delete user-authored `.agents/` siblings). `--target` filter, `--yes` non-interactive. Not the happy-path docs story — adoption stays `init --mcp --git-hooks` + committed `.agents/`. Effort: S.
 - [x] **HEAD / index freshness warning** — `index_freshness.commit_drift` + `warning` on `context` / tool metadata; boot stderr on `codemap mcp` / `serve` when concerns remain after prime. Shipped [#149](https://github.com/stainless-code/codemap/pull/149).

src/application/mcp-server.test.ts

@@ -7,7 +7,13 @@ import { Client } from "@modelcontextprotocol/sdk/client/index.js";
 import { InMemoryTransport } from "@modelcontextprotocol/sdk/inMemory.js";
 
 import { resolveCodemapConfig } from "../config";
-import { closeDb, createTables, openDb, upsertQueryBaseline } from "../db";
+import {
+  closeDb,
+  createTables,
+  insertFile,
+  openDb,
+  upsertQueryBaseline,
+} from "../db";
 import { initCodemap } from "../runtime";
 import { createMcpServer } from "./mcp-server";
 
@@ -1822,6 +1828,113 @@ describe("MCP server — trace / explore / node tools", () => {
     }
   });
 
+  it("explore sets truncation.rows on large index without row transport override", async () => {
+    seedWideExploreGraph(130);
+    const db = openDb();
+    try {
+      seedBulkExploreFiles(db, 5000);
+    } finally {
+      closeDb(db);
+    }
+    const { client, server } = await makeClient();
+    try {
+      const r = await client.callTool({
+        name: "explore",
+        arguments: { names: ["hub"], budget_chars: 15_000 },
+      });
+      const json = readJson(r) as {
+        rows: unknown[];
+        truncation?: { rows?: boolean; snippets?: boolean };
+      };
+      expect(json.rows.length).toBeLessThanOrEqual(125);
+      expect(json.truncation?.rows).toBe(true);
+      expect(json.truncation?.snippets).toBeUndefined();
+    } finally {
+      await server.close();
+    }
+  });
+
+  it("node sets truncated when budget_chars is tiny", async () => {
+    seedTraceGraph();
+    const { client, server } = await makeClient();
+    try {
+      const r = await client.callTool({
+        name: "node",
+        arguments: { name: "foo", include_snippets: true, budget_chars: 1 },
+      });
+      const json = readJson(r) as {
+        truncated: boolean;
+        truncation?: { snippets?: boolean };
+      };
+      expect(json.truncated).toBe(true);
+      expect(json.truncation?.snippets).toBe(true);
+    } finally {
+      await server.close();
+    }
+  });
+
+  function seedBulkExploreFiles(
+    db: ReturnType<typeof openDb>,
+    count: number,
+  ): void {
+    for (let i = 0; i < count; i++) {
+      insertFile(db, {
+        path: `src/bulk/${i}.ts`,
+        content_hash: `hb${i}`,
+        size: 1,
+        line_count: 1,
+        language: "typescript",
+        last_modified: 1,
+        indexed_at: 1,
+      });
+    }
+  }
+
+  function seedWideExploreGraph(calleeCount: number): void {
+    const bodyLines = [`export function hub() {`];
+    for (let i = 0; i < calleeCount; i++) {
+      bodyLines.push(`  fn${i}();`);
+    }
+    bodyLines.push("}");
+    for (let i = 0; i < calleeCount; i++) {
+      bodyLines.push(`export function fn${i}() { return ${i}; }`);
+    }
+    writeFileSync(
+      join(benchDir, "src", "wide.ts"),
+      `${bodyLines.join("\n")}\n`,
+    );
+    const db = openDb();
+    try {
+      db.run(
+        `INSERT INTO files (path, content_hash, size, line_count, language, last_modified, indexed_at)
+         VALUES ('src/wide.ts', 'hw', 500, ${calleeCount + 3}, 'typescript', 1, 1)`,
+      );
+      const symbolValues: string[] = [
+        "('hub', 'function', 'src/wide.ts', 1, 2, 'hub()', 1, NULL, 'export')",
+      ];
+      for (let i = 0; i < calleeCount; i++) {
+        const line = 3 + i;
+        symbolValues.push(
+          `('fn${i}', 'function', 'src/wide.ts', ${line}, ${line}, 'fn${i}()', 1, NULL, 'export')`,
+        );
+      }
+      db.run(
+        `INSERT INTO symbols (name, kind, file_path, line_start, line_end, signature, is_exported, parent_name, visibility)
+         VALUES ${symbolValues.join(",")}`,
+      );
+      const callValues: string[] = [];
+      for (let i = 0; i < calleeCount; i++) {
+        callValues.push(`('src/wide.ts', 'hub', 'hub', 'fn${i}', 2, 0, 0)`);
+      }
+      db.run(
+        `INSERT INTO calls (file_path, caller_name, caller_scope, callee_name, line_start, column_start, column_end)
+         VALUES ${callValues.join(",")}`,
+      );
+    } finally {
+      closeDb(db);
+    }
+  }
+
   it("records recipe recency after trace and explore", async () => {
     seedTraceGraph();
     const { client, server } = await makeClient();

src/application/mcp-server.ts

@@ -342,7 +342,7 @@ function registerTraceTool(server: McpServer, opts: ServerOpts): void {
     "trace",
     {
       description:
-        "Shortest call path between two symbols plus budget-capped snippets. Composes `call-path` recipe + disk reads (cross-file callee lookup). Args: from, to, max_depth?, via (calls|dependencies|all), budget_chars (default 15000, snippet source text only). Returns {from, to, via?, path, snippets, truncated, truncation?, snippets_skipped_reason?}. `truncated` is true when snippet budget hit (`truncation.snippets`); dependency hops omit auto-snippets (`snippets_skipped_reason`). Fall back to `query_recipe` call-path when unsure.",
+        "Shortest call path between two symbols plus budget-capped snippets. Composes `call-path` recipe + disk reads (cross-file callee lookup). Args: from, to, max_depth?, via (calls|dependencies|all), budget_chars (adaptive default 15k/10k/6k when omitted; snippet source text only). Returns {from, to, via?, path, snippets, truncated, truncation?, snippets_skipped_reason?}. `truncated` is true when snippet budget hit (`truncation.snippets`); dependency hops omit auto-snippets (`snippets_skipped_reason`). Fall back to `query_recipe` call-path when unsure.",
       inputSchema: traceArgsSchema,
     },
     (args) => wrapToolResult(handleTrace(args, opts.root)),
@@ -354,7 +354,7 @@ function registerExploreTool(server: McpServer, opts: ServerOpts): void {
     "explore",
     {
       description:
-        "Multi-symbol neighborhood survey with budget-capped snippets. Composes `symbol-neighborhood` (once per deduped name) + disk reads. Args: names (non-empty array), depth?, kind?, budget_chars (default 15000, snippet source only). Returns {names, rows, snippets, truncated, truncation?} — `truncation.rows` when row cap (500) hit, `truncation.snippets` when budget hit. Fall back to `query_recipe` symbol-neighborhood.",
+        "Multi-symbol neighborhood survey with budget-capped snippets. Composes `symbol-neighborhood` (once per deduped name) + disk reads. Args: names (non-empty array), depth?, kind?, budget_chars (adaptive default 15k/10k/6k snippet chars when omitted). Returns {names, rows, snippets, truncated, truncation?} — `truncation.rows` when adaptive row cap hit (500/250/125 by repo size), `truncation.snippets` when budget hit. Fall back to `query_recipe` symbol-neighborhood.",
       inputSchema: exploreArgsSchema,
     },
     (args) => wrapToolResult(handleExplore(args, opts.root)),
@@ -366,7 +366,7 @@ function registerNodeTool(server: McpServer, opts: ServerOpts): void {
     "node",
     {
       description:
-        "One-hop symbol survey: `show` center + scoped depth-1 `symbol-neighborhood` + optional inline snippets. When center is unique (`in` or single match), neighborhood filters to that instance's connected files. Args: name, kind?, in?, include_snippets (default false), budget_chars? (default 15000 when snippets enabled; snippet source only). Returns {center, neighborhood, snippets, truncated, truncation?}. `truncated` only when `include_snippets: true` and snippet budget hit.",
+        "One-hop symbol survey: `show` center + scoped depth-1 `symbol-neighborhood` + optional inline snippets. When center is unique (`in` or single match), neighborhood filters to that instance's connected files. Args: name, kind?, in?, include_snippets (default false), budget_chars? (adaptive default 15k/10k/6k when snippets enabled and omitted; snippet source only). Returns {center, neighborhood, snippets, truncated, truncation?}. `truncated` only when `include_snippets: true` and snippet budget hit.",
       inputSchema: nodeArgsSchema,
     },
     (args) => wrapToolResult(handleNode(args, opts.root)),

src/application/output-budget.test.ts

@@ -1,10 +1,138 @@
-import { describe, expect, it } from "bun:test";
+import { afterEach, beforeEach, describe, expect, it } from "bun:test";
+import { mkdirSync, mkdtempSync, rmSync } from "node:fs";
+import { tmpdir } from "node:os";
+import { join } from "node:path";
 
+import { resolveCodemapConfig } from "../config";
+import { closeDb, createTables, insertFile, openDb } from "../db";
+import { initCodemap } from "../runtime";
 import {
   applySourceCharBudget,
+  DEFAULT_EXPLORE_ROW_LIMIT,
   DEFAULT_OUTPUT_CHAR_BUDGET,
+  readIndexedFileCount,
+  resolveEffectiveOutputBudget,
+  resolveEffectiveSnippetBudget,
+  resolveOutputBudget,
 } from "./output-budget";
 
+let benchDir: string;
+
+beforeEach(() => {
+  benchDir = mkdtempSync(join(tmpdir(), "output-budget-"));
+  mkdirSync(join(benchDir, "src"), { recursive: true });
+  initCodemap(resolveCodemapConfig(benchDir, undefined));
+});
+
+afterEach(() => {
+  rmSync(benchDir, { recursive: true, force: true });
+});
+
+function seedFileCount(count: number): ReturnType<typeof openDb> {
+  const db = openDb();
+  createTables(db);
+  for (let i = 0; i < count; i++) {
+    insertFile(db, {
+      path: `src/f${i}.ts`,
+      content_hash: `h${i}`,
+      size: 1,
+      line_count: 1,
+      language: "typescript",
+      last_modified: 1,
+      indexed_at: 1,
+    });
+  }
+  return db;
+}
+
+describe("resolveOutputBudget", () => {
+  it("uses full caps on small repos", () => {
+    expect(resolveOutputBudget(100)).toEqual({
+      snippet_char_budget: DEFAULT_OUTPUT_CHAR_BUDGET,
+      explore_row_limit: DEFAULT_EXPLORE_ROW_LIMIT,
+    });
+    expect(resolveOutputBudget(500).snippet_char_budget).toBe(15_000);
+  });
+
+  it("tightens caps on mid-size repos", () => {
+    expect(resolveOutputBudget(501)).toEqual({
+      snippet_char_budget: 10_000,
+      explore_row_limit: 250,
+    });
+    expect(resolveOutputBudget(5000).explore_row_limit).toBe(250);
+  });
+
+  it("tightens caps on large repos", () => {
+    expect(resolveOutputBudget(6000)).toEqual({
+      snippet_char_budget: 6_000,
+      explore_row_limit: 125,
+    });
+  });
+});
+
+describe("resolveEffectiveSnippetBudget", () => {
+  it("honors explicit budget_chars", () => {
+    const db = seedFileCount(6000);
+    try {
+      expect(resolveEffectiveSnippetBudget(db, 42)).toBe(42);
+    } finally {
+      closeDb(db);
+    }
+  });
+
+  it("derives adaptive cap from indexed file count", () => {
+    const smallDir = mkdtempSync(join(tmpdir(), "output-budget-small-"));
+    const largeDir = mkdtempSync(join(tmpdir(), "output-budget-large-"));
+    mkdirSync(join(smallDir, "src"), { recursive: true });
+    mkdirSync(join(largeDir, "src"), { recursive: true });
+    initCodemap(resolveCodemapConfig(smallDir, undefined));
+    const small = seedFileCount(3);
+    const smallBudget = resolveEffectiveSnippetBudget(small);
+    closeDb(small);
+
+    initCodemap(resolveCodemapConfig(largeDir, undefined));
+    const large = seedFileCount(6000);
+    try {
+      expect(smallBudget).toBe(15_000);
+      expect(resolveEffectiveSnippetBudget(large)).toBe(6_000);
+      expect(readIndexedFileCount(large)).toBe(6000);
+    } finally {
+      closeDb(large);
+      rmSync(smallDir, { recursive: true, force: true });
+      rmSync(largeDir, { recursive: true, force: true });
+    }
+  });
+});
+
+describe("resolveEffectiveOutputBudget", () => {
+  it("resolves both caps in one indexed file count read", () => {
+    const db = seedFileCount(501);
+    try {
+      expect(resolveEffectiveOutputBudget(db)).toEqual({
+        snippet_char_budget: 10_000,
+        explore_row_limit: 250,
+      });
+      expect(
+        resolveEffectiveOutputBudget(db, {
+          budgetChars: 42,
+          rowLimit: 7,
+        }),
+      ).toEqual({
+        snippet_char_budget: 42,
+        explore_row_limit: 7,
+      });
+      expect(resolveEffectiveOutputBudget(db, { budgetChars: 99_999 })).toEqual(
+        {
+          snippet_char_budget: 99_999,
+          explore_row_limit: 250,
+        },
+      );
+    } finally {
+      closeDb(db);
+    }
+  });
+});
+
 describe("applySourceCharBudget", () => {
   it("returns all items when under budget", () => {
     const items = [{ source: "abc" }, { source: "de" }];

src/application/output-budget.ts

@@ -1,6 +1,67 @@
+import type { CodemapDatabase } from "../db";
+
 /** Default char budget for trace/explore/node snippet payloads (plan L.3). */
 export const DEFAULT_OUTPUT_CHAR_BUDGET = 15_000;
 
+/** Default row cap for explore before `truncation.rows` (structural payload guard). */
+export const DEFAULT_EXPLORE_ROW_LIMIT = 500;
+
+export interface OutputBudget {
+  snippet_char_budget: number;
+  explore_row_limit: number;
+}
+
+/**
+ * Scale trace/explore/node snippet and row caps from indexed file count — same
+ * tier boundaries as {@link resolveContextBudget} in context-engine.
+ */
+export function resolveOutputBudget(fileCount: number): OutputBudget {
+  if (fileCount <= 500) {
+    return {
+      snippet_char_budget: DEFAULT_OUTPUT_CHAR_BUDGET,
+      explore_row_limit: DEFAULT_EXPLORE_ROW_LIMIT,
+    };
+  }
+  if (fileCount <= 5000) {
+    return {
+      snippet_char_budget: 10_000,
+      explore_row_limit: 250,
+    };
+  }
+  return {
+    snippet_char_budget: 6_000,
+    explore_row_limit: 125,
+  };
+}
+
+export function readIndexedFileCount(db: CodemapDatabase): number {
+  const row = db.query("SELECT COUNT(*) AS n FROM files").get() as {
+    n: number;
+  };
+  return row.n;
+}
+
+/** Resolve snippet + explore caps with one indexed file count read. */
+export function resolveEffectiveOutputBudget(
+  db: CodemapDatabase,
+  opts?: { budgetChars?: number; rowLimit?: number },
+): OutputBudget {
+  const adaptive = resolveOutputBudget(readIndexedFileCount(db));
+  return {
+    snippet_char_budget: opts?.budgetChars ?? adaptive.snippet_char_budget,
+    explore_row_limit: opts?.rowLimit ?? adaptive.explore_row_limit,
+  };
+}
+
+/** Explicit `budget_chars` wins; otherwise adaptive cap from indexed file count. */
+export function resolveEffectiveSnippetBudget(
+  db: CodemapDatabase,
+  budgetChars?: number,
+): number {
+  if (budgetChars !== undefined) return budgetChars;
+  return resolveOutputBudget(readIndexedFileCount(db)).snippet_char_budget;
+}
+
 export interface SourceCharBudgetResult<
   T extends { source?: string | undefined },
 > {