token-pilot 0.19.2 → 0.23.0

package/dist/server.js

@@ -1,49 +1,52 @@
-import { Server } from '@modelcontextprotocol/sdk/server/index.js';
-import { CallToolRequestSchema, ListToolsRequestSchema, } from '@modelcontextprotocol/sdk/types.js';
-import { AstIndexClient } from './ast-index/client.js';
-import { FileCache } from './core/file-cache.js';
-import { ContextRegistry } from './core/context-registry.js';
-import { SymbolResolver } from './core/symbol-resolver.js';
-import { SessionAnalytics } from './core/session-analytics.js';
-import { classifyIntent } from './core/intent-classifier.js';
-import { buildDecisionTrace } from './core/decision-trace.js';
-import { SessionCache } from './core/session-cache.js';
-import { loadConfig } from './config/loader.js';
-import { readFileSync } from 'node:fs';
-import { dirname, resolve } from 'node:path';
-import { execFile } from 'node:child_process';
-import { isDangerousRoot } from './core/validation.js';
-import { promisify } from 'node:util';
-import { GitWatcher } from './git/watcher.js';
+import { Server } from "@modelcontextprotocol/sdk/server/index.js";
+import { CallToolRequestSchema, ListToolsRequestSchema, } from "@modelcontextprotocol/sdk/types.js";
+import { AstIndexClient } from "./ast-index/client.js";
+import { FileCache } from "./core/file-cache.js";
+import { ContextRegistry } from "./core/context-registry.js";
+import { SessionRegistryManager } from "./core/session-registry.js";
+import { SymbolResolver } from "./core/symbol-resolver.js";
+import { SessionAnalytics, } from "./core/session-analytics.js";
+import { classifyIntent } from "./core/intent-classifier.js";
+import { buildDecisionTrace } from "./core/decision-trace.js";
+import { SessionCache } from "./core/session-cache.js";
+import { loadConfig } from "./config/loader.js";
+import { readFileSync } from "node:fs";
+import { dirname, resolve } from "node:path";
+import { execFile } from "node:child_process";
+import { isDangerousRoot } from "./core/validation.js";
+import { promisify } from "node:util";
+import { GitWatcher } from "./git/watcher.js";
 const execFilePromise = promisify(execFile);
-import { FileWatcher } from './git/file-watcher.js';
-import { handleSmartRead } from './handlers/smart-read.js';
-import { handleReadSymbol } from './handlers/read-symbol.js';
-import { handleReadSymbols } from './handlers/read-symbols.js';
-import { handleReadRange } from './handlers/read-range.js';
-import { handleReadDiff } from './handlers/read-diff.js';
-import { handleFindUsages } from './handlers/find-usages.js';
-import { handleSmartReadMany } from './handlers/smart-read-many.js';
-import { handleProjectOverview } from './handlers/project-overview.js';
-import { handleNonCodeRead, isNonCodeStructured } from './handlers/non-code.js';
-import { handleFindUnused } from './handlers/find-unused.js';
-import { handleReadForEdit } from './handlers/read-for-edit.js';
-import { handleRelatedFiles } from './handlers/related-files.js';
-import { handleOutline } from './handlers/outline.js';
-import { handleCodeAudit } from './handlers/code-audit.js';
-import { handleModuleInfo } from './handlers/module-info.js';
-import { handleSmartDiff } from './handlers/smart-diff.js';
-import { handleExploreArea } from './handlers/explore-area.js';
-import { handleSmartLog } from './handlers/smart-log.js';
-import { handleTestSummary } from './handlers/test-summary.js';
-import { handleSessionSnapshot } from './handlers/session-snapshot.js';
-import { handleReadSection } from './handlers/read-section.js';
-import { detectContextMode } from './integration/context-mode-detector.js';
-import { estimateTokens } from './core/token-estimator.js';
-import { checkPolicy, isFullReadTool } from './core/policy-engine.js';
-import { MCP_INSTRUCTIONS, TOOL_DEFINITIONS } from './server/tool-definitions.js';
-import { createTokenEstimates } from './server/token-estimates.js';
-import { validateSmartReadArgs, validateReadSymbolArgs, validateReadSymbolsArgs, validateReadRangeArgs, validateReadDiffArgs, validateFindUsagesArgs, validateSmartReadManyArgs, validateReadForEditArgs, validateRelatedFilesArgs, validateOutlineArgs, validateFindUnusedArgs, validateCodeAuditArgs, validateProjectOverviewArgs, validateModuleInfoArgs, validateSmartDiffArgs, validateExploreAreaArgs, validateSmartLogArgs, validateTestSummaryArgs, validateReadSectionArgs, } from './core/validation.js';
+import { FileWatcher } from "./git/file-watcher.js";
+import { handleSmartRead } from "./handlers/smart-read.js";
+import { handleReadSymbol } from "./handlers/read-symbol.js";
+import { handleReadSymbols } from "./handlers/read-symbols.js";
+import { handleReadRange } from "./handlers/read-range.js";
+import { handleReadDiff } from "./handlers/read-diff.js";
+import { handleFindUsages } from "./handlers/find-usages.js";
+import { handleSmartReadMany } from "./handlers/smart-read-many.js";
+import { handleProjectOverview } from "./handlers/project-overview.js";
+import { handleNonCodeRead, isNonCodeStructured } from "./handlers/non-code.js";
+import { handleFindUnused } from "./handlers/find-unused.js";
+import { handleReadForEdit } from "./handlers/read-for-edit.js";
+import { handleRelatedFiles } from "./handlers/related-files.js";
+import { handleOutline } from "./handlers/outline.js";
+import { handleCodeAudit } from "./handlers/code-audit.js";
+import { handleModuleInfo } from "./handlers/module-info.js";
+import { handleSmartDiff } from "./handlers/smart-diff.js";
+import { handleExploreArea } from "./handlers/explore-area.js";
+import { handleSmartLog } from "./handlers/smart-log.js";
+import { handleTestSummary } from "./handlers/test-summary.js";
+import { handleSessionSnapshot } from "./handlers/session-snapshot.js";
+import { persistSnapshot } from "./handlers/session-snapshot-persist.js";
+import { handleSessionBudget } from "./handlers/session-budget.js";
+import { handleReadSection } from "./handlers/read-section.js";
+import { detectContextMode } from "./integration/context-mode-detector.js";
+import { estimateTokens } from "./core/token-estimator.js";
+import { checkPolicy, isFullReadTool } from "./core/policy-engine.js";
+import { MCP_INSTRUCTIONS, TOOL_DEFINITIONS, } from "./server/tool-definitions.js";
+import { createTokenEstimates } from "./server/token-estimates.js";
+import { validateSmartReadArgs, validateReadSymbolArgs, validateReadSymbolsArgs, validateReadRangeArgs, validateReadDiffArgs, validateFindUsagesArgs, validateSmartReadManyArgs, validateReadForEditArgs, validateRelatedFilesArgs, validateOutlineArgs, validateFindUnusedArgs, validateCodeAuditArgs, validateProjectOverviewArgs, validateModuleInfoArgs, validateSmartDiffArgs, validateExploreAreaArgs, validateSmartLogArgs, validateTestSummaryArgs, validateReadSectionArgs, } from "./core/validation.js";
 export async function createServer(projectRoot, options) {
     const config = await loadConfig(projectRoot);
     const astIndex = new AstIndexClient(projectRoot, config.astIndex.timeout, {
@@ -52,7 +55,38 @@ export async function createServer(projectRoot, options) {
     });
     const fileCache = new FileCache(config.cache.maxSizeMB, config.smartRead.smallFileThreshold);
     const contextRegistry = new ContextRegistry();
+    const sessionRegistries = new SessionRegistryManager(projectRoot);
+    // Flush persisted session registries on shutdown (best-effort; every hot
+    // tool-call path also flushes immediately, so this is only for registries
+    // whose last access never got a post-call flush). `beforeExit` doesn't
+    // fire on signal-based termination (SIGINT / SIGTERM), so we hook those
+    // too — Node runs every signal listener before the default action, giving
+    // flushAll a fair chance to complete. `process.exit()` bypasses listeners
+    // entirely; callers that care about durability should not use it.
+    const shutdownFlush = () => {
+        void sessionRegistries.flushAll();
+    };
+    process.once("beforeExit", shutdownFlush);
+    process.once("SIGINT", shutdownFlush);
+    process.once("SIGTERM", shutdownFlush);
     const symbolResolver = new SymbolResolver(astIndex);
+    /**
+     * TP-69m — pick the right ContextRegistry for this tool call.
+     *   - force:true  → empty registry (agent wants to bypass dedup)
+     *   - session_id present → per-session, disk-backed registry
+     *   - neither    → process-default (legacy behaviour for callers that
+     *     don't yet know their session_id)
+     */
+    function pickRegistry(rawArgs) {
+        const a = (rawArgs ?? {});
+        const force = a.force === true;
+        const sessionId = typeof a.session_id === "string" ? a.session_id : "";
+        if (force)
+            return { reg: new ContextRegistry(), sessionId, force: true };
+        if (sessionId)
+            return { reg: sessionRegistries.getFor(sessionId), sessionId, force };
+        return { reg: contextRegistry, sessionId, force };
+    }
     // Try to init ast-index (non-fatal if not available)
     const needsAutoDetect = !!options?.skipAstIndex;
     try {
@@ -61,7 +95,7 @@ export async function createServer(projectRoot, options) {
             // Dangerous root (/, home dir) — don't build index yet
             // Will auto-detect real project root from first file path
             astIndex.disableIndex();
-            console.error('[token-pilot] ast-index: waiting for first file path to auto-detect project root');
+            console.error("[token-pilot] ast-index: waiting for first file path to auto-detect project root");
         }
         else if (config.astIndex.buildOnStart) {
             await astIndex.ensureIndex();
@@ -96,10 +130,10 @@ export async function createServer(projectRoot, options) {
             if (caps?.roots) {
                 const { roots } = await server.listRoots();
                 for (const root of roots) {
-                    if (root.uri.startsWith('file://')) {
+                    if (root.uri.startsWith("file://")) {
                         const rootPath = decodeURIComponent(new URL(root.uri).pathname);
                         if (rootPath && !isDangerousRoot(rootPath)) {
-                            await applyDetectedRoot(rootPath, 'MCP roots');
+                            await applyDetectedRoot(rootPath, "MCP roots");
                             return;
                         }
                     }
@@ -113,7 +147,7 @@ export async function createServer(projectRoot, options) {
         if (filePath) {
             const dir = dirname(filePath);
             try {
-                const { stdout } = await execFilePromise('git', ['rev-parse', '--show-toplevel'], {
+                const { stdout } = await execFilePromise("git", ["rev-parse", "--show-toplevel"], {
                     cwd: dir,
                     timeout: 3000,
                 });
@@ -133,16 +167,16 @@ export async function createServer(projectRoot, options) {
      */
     function extractFilePath(toolArgs) {
         const path = toolArgs?.path;
-        if (path && typeof path === 'string' && path.startsWith('/'))
+        if (path && typeof path === "string" && path.startsWith("/"))
             return path;
         const paths = toolArgs?.paths;
-        if (paths?.[0] && typeof paths[0] === 'string' && paths[0].startsWith('/'))
+        if (paths?.[0] && typeof paths[0] === "string" && paths[0].startsWith("/"))
             return paths[0];
         const file = toolArgs?.file;
-        if (file && typeof file === 'string' && file.startsWith('/'))
+        if (file && typeof file === "string" && file.startsWith("/"))
             return file;
         const mod = toolArgs?.module;
-        if (mod && typeof mod === 'string' && mod.startsWith('/'))
+        if (mod && typeof mod === "string" && mod.startsWith("/"))
             return mod;
         return undefined;
     }
@@ -160,7 +194,7 @@ export async function createServer(projectRoot, options) {
     const readForEditCalled = new Set();
     // Detect context-mode companion
     const cmEnabled = config.contextMode.enabled;
-    const contextModeStatus = await detectContextMode(projectRoot, cmEnabled === 'auto' ? undefined : cmEnabled);
+    const contextModeStatus = await detectContextMode(projectRoot, cmEnabled === "auto" ? undefined : cmEnabled);
     if (contextModeStatus.detected) {
         console.error(`[token-pilot] context-mode detected (source: ${contextModeStatus.source})`);
     }
@@ -193,14 +227,16 @@ export async function createServer(projectRoot, options) {
         });
     }
     // Read version from package.json
-    let pkgVersion = '0.1.1';
+    let pkgVersion = "0.1.1";
     try {
-        const pkgPath = new URL('../package.json', import.meta.url).pathname;
-        const pkg = JSON.parse(readFileSync(pkgPath, 'utf-8'));
+        const pkgPath = new URL("../package.json", import.meta.url).pathname;
+        const pkg = JSON.parse(readFileSync(pkgPath, "utf-8"));
         pkgVersion = pkg.version;
     }
-    catch { /* fallback to hardcoded */ }
-    const server = new Server({ name: 'token-pilot', version: pkgVersion }, {
+    catch {
+        /* fallback to hardcoded */
+    }
+    const server = new Server({ name: "token-pilot", version: pkgVersion }, {
         capabilities: { tools: {} },
         instructions: MCP_INSTRUCTIONS,
     });
@@ -232,7 +268,7 @@ export async function createServer(projectRoot, options) {
         if (isFullReadTool(rest.tool)) {
             fullFileReadsCount++;
         }
-        if (rest.tool === 'read_for_edit' && call.path) {
+        if (rest.tool === "read_for_edit" && call.path) {
             readForEditCalled.add(call.path);
         }
         // Policy check
@@ -251,30 +287,32 @@ export async function createServer(projectRoot, options) {
         // Auto-detect project root on first tool call (when startup root was /)
         // Tries: MCP roots → git detect from file path in args
         if (needsAutoDetect && !autoDetectDone) {
-            const detectedPath = extractFilePath((args ?? {}));
+            const detectedPath = extractFilePath(args ?? {});
             await tryAutoDetectRoot(detectedPath);
         }
         try {
             switch (name) {
-                case 'smart_read': {
+                case "smart_read": {
                     const validArgs = validateSmartReadArgs(args);
+                    const picked = pickRegistry(args);
                     // Try non-code handler for JSON/YAML/MD etc.
                     if (isNonCodeStructured(validArgs.path)) {
-                        const nonCodeResult = await handleNonCodeRead(validArgs.path, projectRoot, contextRegistry, {
+                        const nonCodeResult = await handleNonCodeRead(validArgs.path, projectRoot, picked.reg, {
                             contextModeStatus,
                             largeNonCodeThreshold: config.contextMode.largeNonCodeThreshold,
                             adviseDelegation: config.contextMode.adviseDelegation,
                         });
                         if (nonCodeResult) {
-                            const text = nonCodeResult.content[0]?.text ?? '';
+                            const text = nonCodeResult.content[0]?.text ?? "";
                             recordWithTrace({
-                                tool: 'smart_read',
+                                tool: "smart_read",
                                 path: validArgs.path,
                                 tokensReturned: estimateTokens(text),
-                                tokensWouldBe: await fullFileTokens(validArgs.path) || estimateTokens(text),
+                                tokensWouldBe: (await fullFileTokens(validArgs.path)) ||
+                                    estimateTokens(text),
                                 timestamp: Date.now(),
-                                delegatedToContextMode: text.includes('ADVISORY:') && text.includes('context-mode'),
-                                savingsCategory: 'compression',
+                                delegatedToContextMode: text.includes("ADVISORY:") && text.includes("context-mode"),
+                                savingsCategory: "compression",
                                 absPath: resolve(projectRoot, validArgs.path),
                                 args: validArgs,
                             });
@@ -282,78 +320,151 @@ export async function createServer(projectRoot, options) {
                         }
                     }
                     // Dedup is handled inside handleSmartRead (step 5)
-                    const result = await handleSmartRead(validArgs, projectRoot, astIndex, fileCache, contextRegistry, config);
-                    const text = result.content[0]?.text ?? '';
+                    const result = await handleSmartRead(validArgs, projectRoot, astIndex, fileCache, picked.reg, config);
+                    if (picked.sessionId && !picked.force) {
+                        void sessionRegistries.flush(picked.sessionId);
+                    }
+                    const text = result.content[0]?.text ?? "";
                     const fullTokensSR = await fullFileTokens(validArgs.path);
-                    const policyAdv = recordWithTrace({ tool: 'smart_read', path: validArgs.path, tokensReturned: estimateTokens(text), tokensWouldBe: fullTokensSR || estimateTokens(text), timestamp: Date.now(), savingsCategory: detectSavingsCategory(text), absPath: resolve(projectRoot, validArgs.path), args: validArgs });
+                    const policyAdv = recordWithTrace({
+                        tool: "smart_read",
+                        path: validArgs.path,
+                        tokensReturned: estimateTokens(text),
+                        tokensWouldBe: fullTokensSR || estimateTokens(text),
+                        timestamp: Date.now(),
+                        savingsCategory: detectSavingsCategory(text),
+                        absPath: resolve(projectRoot, validArgs.path),
+                        args: validArgs,
+                    });
                     if (policyAdv)
-                        result.content[0] = { type: 'text', text: text + policyAdv };
+                        result.content[0] = { type: "text", text: text + policyAdv };
                     return result;
                 }
-                case 'read_symbol': {
+                case "read_symbol": {
                     const symArgs = validateReadSymbolArgs(args);
+                    const pickedSym = pickRegistry(args);
                     // Dedup is handled inside handleReadSymbol
-                    const symResult = await handleReadSymbol(symArgs, projectRoot, symbolResolver, fileCache, contextRegistry, astIndex, config.smartRead.advisoryReminders);
-                    const symText = symResult.content[0]?.text ?? '';
+                    const symResult = await handleReadSymbol(symArgs, projectRoot, symbolResolver, fileCache, pickedSym.reg, astIndex, config.smartRead.advisoryReminders);
+                    if (pickedSym.sessionId && !pickedSym.force) {
+                        void sessionRegistries.flush(pickedSym.sessionId);
+                    }
+                    const symText = symResult.content[0]?.text ?? "";
                     const symTokens = estimateTokens(symText);
                     const fullTokensSym = await fullFileTokens(symArgs.path);
-                    recordWithTrace({ tool: 'read_symbol', path: symArgs.path, tokensReturned: symTokens, tokensWouldBe: fullTokensSym || symTokens, timestamp: Date.now(), savingsCategory: detectSavingsCategory(symText), absPath: resolve(projectRoot, symArgs.path), args: symArgs });
+                    recordWithTrace({
+                        tool: "read_symbol",
+                        path: symArgs.path,
+                        tokensReturned: symTokens,
+                        tokensWouldBe: fullTokensSym || symTokens,
+                        timestamp: Date.now(),
+                        savingsCategory: detectSavingsCategory(symText),
+                        absPath: resolve(projectRoot, symArgs.path),
+                        args: symArgs,
+                    });
                     return symResult;
                 }
-                case 'read_symbols': {
+                case "read_symbols": {
                     const rsArgs = validateReadSymbolsArgs(args);
                     const rsResult = await handleReadSymbols(rsArgs, projectRoot, symbolResolver, fileCache, contextRegistry, astIndex, config.smartRead.advisoryReminders);
-                    const rsText = rsResult.content[0]?.text ?? '';
+                    const rsText = rsResult.content[0]?.text ?? "";
                     const rsTokens = estimateTokens(rsText);
                     const fullTokensRs = await fullFileTokens(rsArgs.path);
-                    recordWithTrace({ tool: 'read_symbols', path: rsArgs.path, tokensReturned: rsTokens, tokensWouldBe: fullTokensRs || rsTokens, timestamp: Date.now(), savingsCategory: 'compression', absPath: resolve(projectRoot, rsArgs.path), args: rsArgs });
+                    recordWithTrace({
+                        tool: "read_symbols",
+                        path: rsArgs.path,
+                        tokensReturned: rsTokens,
+                        tokensWouldBe: fullTokensRs || rsTokens,
+                        timestamp: Date.now(),
+                        savingsCategory: "compression",
+                        absPath: resolve(projectRoot, rsArgs.path),
+                        args: rsArgs,
+                    });
                     return rsResult;
                 }
-                case 'read_range': {
+                case "read_range": {
                     const rangeArgs = validateReadRangeArgs(args);
-                    const rangeResult = await handleReadRange(rangeArgs, projectRoot, fileCache, contextRegistry, config.smartRead.advisoryReminders);
-                    const rangeText = rangeResult.content[0]?.text ?? '';
+                    const pickedRange = pickRegistry(args);
+                    const rangeResult = await handleReadRange(rangeArgs, projectRoot, fileCache, pickedRange.reg, config.smartRead.advisoryReminders);
+                    if (pickedRange.sessionId && !pickedRange.force) {
+                        void sessionRegistries.flush(pickedRange.sessionId);
+                    }
+                    const rangeText = rangeResult.content[0]?.text ?? "";
                     const rangeTokens = estimateTokens(rangeText);
                     const fullTokensRange = await fullFileTokens(rangeArgs.path);
-                    recordWithTrace({ tool: 'read_range', path: rangeArgs.path, tokensReturned: rangeTokens, tokensWouldBe: fullTokensRange || rangeTokens, timestamp: Date.now(), savingsCategory: detectSavingsCategory(rangeText), absPath: resolve(projectRoot, rangeArgs.path), args: rangeArgs });
+                    recordWithTrace({
+                        tool: "read_range",
+                        path: rangeArgs.path,
+                        tokensReturned: rangeTokens,
+                        tokensWouldBe: fullTokensRange || rangeTokens,
+                        timestamp: Date.now(),
+                        savingsCategory: detectSavingsCategory(rangeText),
+                        absPath: resolve(projectRoot, rangeArgs.path),
+                        args: rangeArgs,
+                    });
                     return rangeResult;
                 }
-                case 'read_section': {
+                case "read_section": {
                     const secArgs = validateReadSectionArgs(args);
                     const secResult = await handleReadSection(secArgs, projectRoot, contextRegistry);
-                    const secText = secResult.content[0]?.text ?? '';
+                    const secText = secResult.content[0]?.text ?? "";
                     const secTokens = estimateTokens(secText);
                     const fullTokensSec = await fullFileTokens(secArgs.path);
                     recordWithTrace({
-                        tool: 'read_section', path: secArgs.path,
-                        tokensReturned: secTokens, tokensWouldBe: fullTokensSec || secTokens,
-                        timestamp: Date.now(), savingsCategory: 'compression',
-                        absPath: resolve(projectRoot, secArgs.path), args: secArgs,
+                        tool: "read_section",
+                        path: secArgs.path,
+                        tokensReturned: secTokens,
+                        tokensWouldBe: fullTokensSec || secTokens,
+                        timestamp: Date.now(),
+                        savingsCategory: "compression",
+                        absPath: resolve(projectRoot, secArgs.path),
+                        args: secArgs,
                     });
                     return secResult;
                 }
-                case 'read_diff': {
+                case "read_diff": {
                     const diffArgs = validateReadDiffArgs(args);
                     const diffResult = await handleReadDiff(diffArgs, projectRoot, fileCache, contextRegistry);
-                    const diffText = diffResult.content[0]?.text ?? '';
+                    const diffText = diffResult.content[0]?.text ?? "";
                     const diffTokens = estimateTokens(diffText);
                     const fullTokensDiff = await fullFileTokens(diffArgs.path);
-                    recordWithTrace({ tool: 'read_diff', path: diffArgs.path, tokensReturned: diffTokens, tokensWouldBe: fullTokensDiff || diffTokens, timestamp: Date.now(), savingsCategory: 'compression', absPath: resolve(projectRoot, diffArgs.path), args: diffArgs });
+                    recordWithTrace({
+                        tool: "read_diff",
+                        path: diffArgs.path,
+                        tokensReturned: diffTokens,
+                        tokensWouldBe: fullTokensDiff || diffTokens,
+                        timestamp: Date.now(),
+                        savingsCategory: "compression",
+                        absPath: resolve(projectRoot, diffArgs.path),
+                        args: diffArgs,
+                    });
                     return diffResult;
                 }
-                case 'read_for_edit': {
+                case "read_for_edit": {
                     const editArgs = validateReadForEditArgs(args);
                     const editResult = await handleReadForEdit(editArgs, projectRoot, symbolResolver, fileCache, contextRegistry, astIndex, { actionableHints: config.display.actionableHints });
-                    const editText = editResult.content[0]?.text ?? '';
+                    const editText = editResult.content[0]?.text ?? "";
                     const editTokens = estimateTokens(editText);
                     const fullTokensEdit = await fullFileTokens(editArgs.path);
-                    recordWithTrace({ tool: 'read_for_edit', path: editArgs.path, tokensReturned: editTokens, tokensWouldBe: fullTokensEdit || editTokens, timestamp: Date.now(), savingsCategory: 'compression', absPath: resolve(projectRoot, editArgs.path), args: editArgs });
+                    recordWithTrace({
+                        tool: "read_for_edit",
+                        path: editArgs.path,
+                        tokensReturned: editTokens,
+                        tokensWouldBe: fullTokensEdit || editTokens,
+                        timestamp: Date.now(),
+                        savingsCategory: "compression",
+                        absPath: resolve(projectRoot, editArgs.path),
+                        args: editArgs,
+                    });
                     return editResult;
                 }
-                case 'smart_read_many': {
+                case "smart_read_many": {
                     const manyArgs = validateSmartReadManyArgs(args);
-                    const manyResult = await handleSmartReadMany(manyArgs, projectRoot, astIndex, fileCache, contextRegistry, config);
-                    const manyText = manyResult.content[0]?.text ?? '';
+                    const pickedMany = pickRegistry(args);
+                    const manyResult = await handleSmartReadMany(manyArgs, projectRoot, astIndex, fileCache, pickedMany.reg, config);
+                    if (pickedMany.sessionId && !pickedMany.force) {
+                        void sessionRegistries.flush(pickedMany.sessionId);
+                    }
+                    const manyText = manyResult.content[0]?.text ?? "";
                     const manyTokens = estimateTokens(manyText);
                     const uniqueManyPaths = Array.from(new Set(manyArgs.paths));
                     let fullTokensMany = 0;
@@ -361,241 +472,407 @@ export async function createServer(projectRoot, options) {
                         fullTokensMany += await fullFileTokens(p);
                     }
                     recordWithTrace({
-                        tool: 'smart_read_many',
-                        path: uniqueManyPaths.join(', '),
+                        tool: "smart_read_many",
+                        path: uniqueManyPaths.join(", "),
                         tokensReturned: manyTokens,
                         tokensWouldBe: fullTokensMany || manyTokens,
                         timestamp: Date.now(),
-                        savingsCategory: 'compression',
+                        savingsCategory: "compression",
                         args: manyArgs,
                     });
                     return manyResult;
                 }
-                case 'find_usages': {
+                case "find_usages": {
                     const usagesArgs = validateFindUsagesArgs(args);
-                    const cachedUsages = sessionCache?.get('find_usages', usagesArgs);
+                    const cachedUsages = sessionCache?.get("find_usages", usagesArgs);
                     if (cachedUsages) {
-                        recordWithTrace({ tool: 'find_usages', path: usagesArgs.symbol, tokensReturned: cachedUsages.tokenEstimate, tokensWouldBe: cachedUsages.tokensWouldBe ?? cachedUsages.tokenEstimate, timestamp: Date.now(), sessionCacheHit: true, savingsCategory: 'cache', args: usagesArgs });
+                        recordWithTrace({
+                            tool: "find_usages",
+                            path: usagesArgs.symbol,
+                            tokensReturned: cachedUsages.tokenEstimate,
+                            tokensWouldBe: cachedUsages.tokensWouldBe ?? cachedUsages.tokenEstimate,
+                            timestamp: Date.now(),
+                            sessionCacheHit: true,
+                            savingsCategory: "cache",
+                            args: usagesArgs,
+                        });
                         return cachedUsages.result;
                     }
                     const usagesResult = await handleFindUsages(usagesArgs, astIndex, projectRoot);
-                    const usagesText = usagesResult.content[0]?.text ?? '';
+                    const usagesText = usagesResult.content[0]?.text ?? "";
                     const usagesTokens = estimateTokens(usagesText);
                     const usagesWouldBe = await estimateFindUsagesWorkflowTokens(usagesResult.meta.files);
-                    sessionCache?.set('find_usages', usagesArgs, usagesResult, {
-                        files: usagesResult.meta.files.map(f => resolve(projectRoot, f)),
+                    sessionCache?.set("find_usages", usagesArgs, usagesResult, {
+                        files: usagesResult.meta.files.map((f) => resolve(projectRoot, f)),
                         dependsOnAst: true,
                     }, usagesTokens, usagesWouldBe || usagesTokens);
                     recordWithTrace({
-                        tool: 'find_usages',
+                        tool: "find_usages",
                         path: usagesArgs.symbol,
                         tokensReturned: usagesTokens,
                         tokensWouldBe: usagesWouldBe || usagesTokens,
                         timestamp: Date.now(),
-                        savingsCategory: 'compression',
+                        savingsCategory: "compression",
                         args: usagesArgs,
                     });
                     return usagesResult;
                 }
-                case 'project_overview': {
+                case "project_overview": {
                     const overviewArgs = validateProjectOverviewArgs(args);
-                    const cachedOverview = sessionCache?.get('project_overview', overviewArgs);
+                    const cachedOverview = sessionCache?.get("project_overview", overviewArgs);
                     if (cachedOverview) {
-                        recordWithTrace({ tool: 'project_overview', path: projectRoot, tokensReturned: cachedOverview.tokenEstimate, tokensWouldBe: cachedOverview.tokensWouldBe ?? cachedOverview.tokenEstimate, timestamp: Date.now(), sessionCacheHit: true, savingsCategory: 'cache', args: overviewArgs });
+                        recordWithTrace({
+                            tool: "project_overview",
+                            path: projectRoot,
+                            tokensReturned: cachedOverview.tokenEstimate,
+                            tokensWouldBe: cachedOverview.tokensWouldBe ?? cachedOverview.tokenEstimate,
+                            timestamp: Date.now(),
+                            sessionCacheHit: true,
+                            savingsCategory: "cache",
+                            args: overviewArgs,
+                        });
                         return cachedOverview.result;
                     }
                     const overviewResult = await handleProjectOverview(overviewArgs, projectRoot, astIndex, pkgVersion);
-                    const overviewText = overviewResult.content[0]?.text ?? '';
-                    overviewResult.content[0] = { type: 'text', text: `TOKEN PILOT v${pkgVersion}\n\n${overviewText}` };
+                    const overviewText = overviewResult.content[0]?.text ?? "";
+                    overviewResult.content[0] = {
+                        type: "text",
+                        text: `TOKEN PILOT v${pkgVersion}\n\n${overviewText}`,
+                    };
                     const ovTokens = estimateTokens(overviewResult.content[0].text);
-                    const overviewWouldBe = await estimateProjectOverviewWorkflowTokens(overviewArgs.include ?? ['stack', 'ci', 'quality', 'architecture']);
-                    sessionCache?.set('project_overview', overviewArgs, overviewResult, {
+                    const overviewWouldBe = await estimateProjectOverviewWorkflowTokens(overviewArgs.include ?? ["stack", "ci", "quality", "architecture"]);
+                    sessionCache?.set("project_overview", overviewArgs, overviewResult, {
                         dependsOnAst: true,
                     }, ovTokens, overviewWouldBe || ovTokens);
                     recordWithTrace({
-                        tool: 'project_overview',
+                        tool: "project_overview",
                         path: projectRoot,
                         tokensReturned: ovTokens,
                         tokensWouldBe: overviewWouldBe || ovTokens,
                         timestamp: Date.now(),
-                        savingsCategory: 'compression',
+                        savingsCategory: "compression",
                         args: overviewArgs,
                     });
                     return overviewResult;
                 }
-                case 'related_files': {
+                case "related_files": {
                     const relArgs = validateRelatedFilesArgs(args);
-                    const cachedRel = sessionCache?.get('related_files', relArgs);
+                    const cachedRel = sessionCache?.get("related_files", relArgs);
                     if (cachedRel) {
-                        recordWithTrace({ tool: 'related_files', path: relArgs.path, tokensReturned: cachedRel.tokenEstimate, tokensWouldBe: cachedRel.tokensWouldBe ?? cachedRel.tokenEstimate, timestamp: Date.now(), sessionCacheHit: true, savingsCategory: 'cache', absPath: resolve(projectRoot, relArgs.path), args: relArgs });
+                        recordWithTrace({
+                            tool: "related_files",
+                            path: relArgs.path,
+                            tokensReturned: cachedRel.tokenEstimate,
+                            tokensWouldBe: cachedRel.tokensWouldBe ?? cachedRel.tokenEstimate,
+                            timestamp: Date.now(),
+                            sessionCacheHit: true,
+                            savingsCategory: "cache",
+                            absPath: resolve(projectRoot, relArgs.path),
+                            args: relArgs,
+                        });
                         return cachedRel.result;
                     }
                     const relResult = await handleRelatedFiles(relArgs, projectRoot, astIndex);
-                    const relText = relResult.content[0]?.text ?? '';
+                    const relText = relResult.content[0]?.text ?? "";
                     const relTokens = estimateTokens(relText);
                     const relWouldBe = await estimateRelatedFilesWorkflowTokens(relArgs.path, relResult.meta);
                     const relDeps = [
                         resolve(projectRoot, relArgs.path),
-                        ...relResult.meta.imports.map(f => resolve(projectRoot, f)),
-                        ...relResult.meta.importedBy.map(f => resolve(projectRoot, f)),
-                        ...relResult.meta.tests.map(f => resolve(projectRoot, f)),
+                        ...relResult.meta.imports.map((f) => resolve(projectRoot, f)),
+                        ...relResult.meta.importedBy.map((f) => resolve(projectRoot, f)),
+                        ...relResult.meta.tests.map((f) => resolve(projectRoot, f)),
                     ];
-                    sessionCache?.set('related_files', relArgs, relResult, {
+                    sessionCache?.set("related_files", relArgs, relResult, {
                         files: relDeps,
                         dependsOnAst: true,
                     }, relTokens, relWouldBe || relTokens);
                     recordWithTrace({
-                        tool: 'related_files',
+                        tool: "related_files",
                         path: relArgs.path,
                         tokensReturned: relTokens,
                         tokensWouldBe: relWouldBe || relTokens,
                         timestamp: Date.now(),
-                        savingsCategory: 'compression',
+                        savingsCategory: "compression",
                         absPath: resolve(projectRoot, relArgs.path),
                         args: relArgs,
                     });
                     return relResult;
                 }
-                case 'outline': {
+                case "outline": {
                     const outlineArgs = validateOutlineArgs(args);
-                    const cachedOutline = sessionCache?.get('outline', outlineArgs);
+                    const cachedOutline = sessionCache?.get("outline", outlineArgs);
                     if (cachedOutline) {
-                        recordWithTrace({ tool: 'outline', path: outlineArgs.path, tokensReturned: cachedOutline.tokenEstimate, tokensWouldBe: cachedOutline.tokensWouldBe ?? cachedOutline.tokenEstimate, timestamp: Date.now(), sessionCacheHit: true, savingsCategory: 'cache', args: outlineArgs });
+                        recordWithTrace({
+                            tool: "outline",
+                            path: outlineArgs.path,
+                            tokensReturned: cachedOutline.tokenEstimate,
+                            tokensWouldBe: cachedOutline.tokensWouldBe ?? cachedOutline.tokenEstimate,
+                            timestamp: Date.now(),
+                            sessionCacheHit: true,
+                            savingsCategory: "cache",
+                            args: outlineArgs,
+                        });
                         return cachedOutline.result;
                     }
                     const outlineResult = await handleOutline(outlineArgs, projectRoot, astIndex);
-                    const outlineText = outlineResult.content[0]?.text ?? '';
+                    const outlineText = outlineResult.content[0]?.text ?? "";
                     const outlineTokens = estimateTokens(outlineText);
                     const outlineWouldBe = await estimateOutlineWorkflowTokens(outlineArgs.path, outlineArgs.recursive ?? false, outlineArgs.max_depth ?? 2);
-                    sessionCache?.set('outline', outlineArgs, outlineResult, {
-                        files: [resolve(projectRoot, outlineArgs.path) + '/'],
+                    sessionCache?.set("outline", outlineArgs, outlineResult, {
+                        files: [resolve(projectRoot, outlineArgs.path) + "/"],
                         dependsOnAst: true,
                     }, outlineTokens, outlineWouldBe || outlineTokens);
                     recordWithTrace({
-                        tool: 'outline',
+                        tool: "outline",
                         path: outlineArgs.path,
                         tokensReturned: outlineTokens,
                         tokensWouldBe: outlineWouldBe || outlineTokens,
                         timestamp: Date.now(),
-                        savingsCategory: 'compression',
+                        savingsCategory: "compression",
                         args: outlineArgs,
                     });
                     return outlineResult;
                 }
-                case 'session_analytics': {
+                case "session_analytics": {
                     const verbose = args?.verbose === true;
-                    return { content: [{ type: 'text', text: `TOKEN PILOT v${pkgVersion}\n\n${analytics.report(verbose)}` }] };
+                    return {
+                        content: [
+                            {
+                                type: "text",
+                                text: `TOKEN PILOT v${pkgVersion}\n\n${analytics.report(verbose)}`,
+                            },
+                        ],
+                    };
                 }
-                case 'find_unused': {
+                case "find_unused": {
                     const unusedArgs = validateFindUnusedArgs(args);
-                    const cachedUnused = sessionCache?.get('find_unused', unusedArgs);
+                    const cachedUnused = sessionCache?.get("find_unused", unusedArgs);
                     if (cachedUnused) {
-                        recordWithTrace({ tool: 'find_unused', path: unusedArgs.module ?? 'all', tokensReturned: cachedUnused.tokenEstimate, tokensWouldBe: cachedUnused.tokensWouldBe ?? cachedUnused.tokenEstimate, timestamp: Date.now(), sessionCacheHit: true, savingsCategory: 'cache', args: unusedArgs });
+                        recordWithTrace({
+                            tool: "find_unused",
+                            path: unusedArgs.module ?? "all",
+                            tokensReturned: cachedUnused.tokenEstimate,
+                            tokensWouldBe: cachedUnused.tokensWouldBe ?? cachedUnused.tokenEstimate,
+                            timestamp: Date.now(),
+                            sessionCacheHit: true,
+                            savingsCategory: "cache",
+                            args: unusedArgs,
+                        });
                         return cachedUnused.result;
                     }
                     const unusedResult = await handleFindUnused(unusedArgs, astIndex);
-                    const unusedText = unusedResult.content[0]?.text ?? '';
+                    const unusedText = unusedResult.content[0]?.text ?? "";
                     const unusedTokens = estimateTokens(unusedText);
                     const unusedWouldBe = await estimateFindUsagesWorkflowTokens(unusedResult.meta.files);
-                    sessionCache?.set('find_unused', unusedArgs, unusedResult, { dependsOnAst: true }, unusedTokens, unusedWouldBe || unusedTokens);
-                    recordWithTrace({ tool: 'find_unused', path: unusedArgs.module ?? 'all', tokensReturned: unusedTokens, tokensWouldBe: unusedWouldBe || unusedTokens, timestamp: Date.now(), savingsCategory: 'compression', args: unusedArgs });
+                    sessionCache?.set("find_unused", unusedArgs, unusedResult, { dependsOnAst: true }, unusedTokens, unusedWouldBe || unusedTokens);
+                    recordWithTrace({
+                        tool: "find_unused",
+                        path: unusedArgs.module ?? "all",
+                        tokensReturned: unusedTokens,
+                        tokensWouldBe: unusedWouldBe || unusedTokens,
+                        timestamp: Date.now(),
+                        savingsCategory: "compression",
+                        args: unusedArgs,
+                    });
                     return unusedResult;
                 }
-                case 'code_audit': {
+                case "code_audit": {
                     const auditArgs = validateCodeAuditArgs(args);
-                    const cachedAudit = sessionCache?.get('code_audit', auditArgs);
+                    const cachedAudit = sessionCache?.get("code_audit", auditArgs);
                     if (cachedAudit) {
-                        recordWithTrace({ tool: 'code_audit', path: auditArgs.check, tokensReturned: cachedAudit.tokenEstimate, tokensWouldBe: cachedAudit.tokensWouldBe ?? cachedAudit.tokenEstimate, timestamp: Date.now(), sessionCacheHit: true, savingsCategory: 'cache', args: auditArgs });
+                        recordWithTrace({
+                            tool: "code_audit",
+                            path: auditArgs.check,
+                            tokensReturned: cachedAudit.tokenEstimate,
+                            tokensWouldBe: cachedAudit.tokensWouldBe ?? cachedAudit.tokenEstimate,
+                            timestamp: Date.now(),
+                            sessionCacheHit: true,
+                            savingsCategory: "cache",
+                            args: auditArgs,
+                        });
                         return cachedAudit.result;
                     }
                     const auditResult = await handleCodeAudit(auditArgs, projectRoot, astIndex);
-                    const auditText = auditResult.content[0]?.text ?? '';
+                    const auditText = auditResult.content[0]?.text ?? "";
                     const auditTokens = estimateTokens(auditText);
                     const auditWouldBe = await estimateFindUsagesWorkflowTokens(auditResult.meta.files);
-                    sessionCache?.set('code_audit', auditArgs, auditResult, { dependsOnAst: true }, auditTokens, auditWouldBe || auditTokens);
-                    recordWithTrace({ tool: 'code_audit', path: auditArgs.check, tokensReturned: auditTokens, tokensWouldBe: auditWouldBe || auditTokens, timestamp: Date.now(), savingsCategory: 'compression', args: auditArgs });
+                    sessionCache?.set("code_audit", auditArgs, auditResult, { dependsOnAst: true }, auditTokens, auditWouldBe || auditTokens);
+                    recordWithTrace({
+                        tool: "code_audit",
+                        path: auditArgs.check,
+                        tokensReturned: auditTokens,
+                        tokensWouldBe: auditWouldBe || auditTokens,
+                        timestamp: Date.now(),
+                        savingsCategory: "compression",
+                        args: auditArgs,
+                    });
                     return auditResult;
                 }
-                case 'module_info': {
+                case "module_info": {
                     const moduleArgs = validateModuleInfoArgs(args);
-                    const cachedModule = sessionCache?.get('module_info', moduleArgs);
+                    const cachedModule = sessionCache?.get("module_info", moduleArgs);
                     if (cachedModule) {
-                        recordWithTrace({ tool: 'module_info', path: moduleArgs.module, tokensReturned: cachedModule.tokenEstimate, tokensWouldBe: cachedModule.tokensWouldBe ?? cachedModule.tokenEstimate, timestamp: Date.now(), sessionCacheHit: true, savingsCategory: 'cache', args: moduleArgs });
+                        recordWithTrace({
+                            tool: "module_info",
+                            path: moduleArgs.module,
+                            tokensReturned: cachedModule.tokenEstimate,
+                            tokensWouldBe: cachedModule.tokensWouldBe ?? cachedModule.tokenEstimate,
+                            timestamp: Date.now(),
+                            sessionCacheHit: true,
+                            savingsCategory: "cache",
+                            args: moduleArgs,
+                        });
                         return cachedModule.result;
                     }
                     const moduleResult = await handleModuleInfo(moduleArgs, projectRoot, astIndex);
-                    const moduleText = moduleResult.content[0]?.text ?? '';
+                    const moduleText = moduleResult.content[0]?.text ?? "";
                     const moduleTokens = estimateTokens(moduleText);
                     const moduleWouldBe = await estimateFindUsagesWorkflowTokens(moduleResult.meta.files);
-                    sessionCache?.set('module_info', moduleArgs, moduleResult, { dependsOnAst: true }, moduleTokens, moduleWouldBe || moduleTokens);
-                    recordWithTrace({ tool: 'module_info', path: moduleArgs.module, tokensReturned: moduleTokens, tokensWouldBe: moduleWouldBe || moduleTokens, timestamp: Date.now(), savingsCategory: 'compression', args: moduleArgs });
+                    sessionCache?.set("module_info", moduleArgs, moduleResult, { dependsOnAst: true }, moduleTokens, moduleWouldBe || moduleTokens);
+                    recordWithTrace({
+                        tool: "module_info",
+                        path: moduleArgs.module,
+                        tokensReturned: moduleTokens,
+                        tokensWouldBe: moduleWouldBe || moduleTokens,
+                        timestamp: Date.now(),
+                        savingsCategory: "compression",
+                        args: moduleArgs,
+                    });
                     return moduleResult;
                 }
-                case 'smart_diff': {
+                case "smart_diff": {
                     const sdArgs = validateSmartDiffArgs(args);
                     const sdResult = await handleSmartDiff(sdArgs, projectRoot, astIndex);
-                    const sdText = sdResult.content[0]?.text ?? '';
+                    const sdText = sdResult.content[0]?.text ?? "";
                     const sdTokens = estimateTokens(sdText);
-                    recordWithTrace({ tool: 'smart_diff', path: sdArgs.path ?? sdArgs.scope ?? 'unstaged', tokensReturned: sdTokens, tokensWouldBe: sdResult.rawTokens || sdTokens, timestamp: Date.now(), savingsCategory: 'compression', args: sdArgs });
+                    recordWithTrace({
+                        tool: "smart_diff",
+                        path: sdArgs.path ?? sdArgs.scope ?? "unstaged",
+                        tokensReturned: sdTokens,
+                        tokensWouldBe: sdResult.rawTokens || sdTokens,
+                        timestamp: Date.now(),
+                        savingsCategory: "compression",
+                        args: sdArgs,
+                    });
                     return { content: sdResult.content };
                 }
-                case 'explore_area': {
+                case "explore_area": {
                     const eaArgs = validateExploreAreaArgs(args);
-                    const cachedEa = sessionCache?.get('explore_area', eaArgs);
+                    const cachedEa = sessionCache?.get("explore_area", eaArgs);
                     if (cachedEa) {
-                        recordWithTrace({ tool: 'explore_area', path: eaArgs.path, tokensReturned: cachedEa.tokenEstimate, tokensWouldBe: cachedEa.tokensWouldBe ?? cachedEa.tokenEstimate, timestamp: Date.now(), sessionCacheHit: true, savingsCategory: 'cache', args: eaArgs });
+                        recordWithTrace({
+                            tool: "explore_area",
+                            path: eaArgs.path,
+                            tokensReturned: cachedEa.tokenEstimate,
+                            tokensWouldBe: cachedEa.tokensWouldBe ?? cachedEa.tokenEstimate,
+                            timestamp: Date.now(),
+                            sessionCacheHit: true,
+                            savingsCategory: "cache",
+                            args: eaArgs,
+                        });
                         return cachedEa.result;
                     }
                     const eaResult = await handleExploreArea(eaArgs, projectRoot, astIndex);
-                    const eaText = eaResult.content[0]?.text ?? '';
+                    const eaText = eaResult.content[0]?.text ?? "";
                     const eaTokens = estimateTokens(eaText);
                     const eaWouldBe = await estimateExploreAreaWorkflowTokens(eaResult.meta);
-                    sessionCache?.set('explore_area', eaArgs, eaResult, {
-                        files: [resolve(projectRoot, eaArgs.path) + '/'],
+                    sessionCache?.set("explore_area", eaArgs, eaResult, {
+                        files: [resolve(projectRoot, eaArgs.path) + "/"],
                         dependsOnAst: true,
                         dependsOnGit: true,
                     }, eaTokens, eaWouldBe || eaTokens);
                     recordWithTrace({
-                        tool: 'explore_area',
+                        tool: "explore_area",
                         path: eaArgs.path,
                         tokensReturned: eaTokens,
                         tokensWouldBe: eaWouldBe || eaTokens,
                         timestamp: Date.now(),
-                        savingsCategory: 'compression',
+                        savingsCategory: "compression",
                         args: eaArgs,
                     });
                     return eaResult;
                 }
-                case 'smart_log': {
+                case "smart_log": {
                     const slArgs = validateSmartLogArgs(args);
                     const slResult = await handleSmartLog(slArgs, projectRoot);
-                    const slText = slResult.content[0]?.text ?? '';
+                    const slText = slResult.content[0]?.text ?? "";
                     const slTokens = estimateTokens(slText);
-                    recordWithTrace({ tool: 'smart_log', path: slArgs.path ?? 'all', tokensReturned: slTokens, tokensWouldBe: slResult.rawTokens || slTokens, timestamp: Date.now(), savingsCategory: 'compression', args: slArgs });
+                    recordWithTrace({
+                        tool: "smart_log",
+                        path: slArgs.path ?? "all",
+                        tokensReturned: slTokens,
+                        tokensWouldBe: slResult.rawTokens || slTokens,
+                        timestamp: Date.now(),
+                        savingsCategory: "compression",
+                        args: slArgs,
+                    });
                     return { content: slResult.content };
                 }
-                case 'test_summary': {
+                case "test_summary": {
                     const tsArgs = validateTestSummaryArgs(args);
                     const tsResult = await handleTestSummary(tsArgs, projectRoot);
-                    const tsText = tsResult.content[0]?.text ?? '';
+                    const tsText = tsResult.content[0]?.text ?? "";
                     const tsTokens = estimateTokens(tsText);
-                    recordWithTrace({ tool: 'test_summary', path: tsArgs.command, tokensReturned: tsTokens, tokensWouldBe: tsResult.rawTokens || tsTokens, timestamp: Date.now(), savingsCategory: 'compression', args: tsArgs });
+                    recordWithTrace({
+                        tool: "test_summary",
+                        path: tsArgs.command,
+                        tokensReturned: tsTokens,
+                        tokensWouldBe: tsResult.rawTokens || tsTokens,
+                        timestamp: Date.now(),
+                        savingsCategory: "compression",
+                        args: tsArgs,
+                    });
                     return { content: tsResult.content };
                 }
-                case 'session_snapshot': {
+                case "session_snapshot": {
                     const snapshotArgs = args;
                     if (!snapshotArgs.goal) {
-                        return { content: [{ type: 'text', text: 'Error: goal is required' }], isError: true };
+                        return {
+                            content: [{ type: "text", text: "Error: goal is required" }],
+                            isError: true,
+                        };
                     }
                     const snapshotResult = handleSessionSnapshot(snapshotArgs);
-                    const snapshotText = snapshotResult.content[0]?.text ?? '';
+                    const snapshotText = snapshotResult.content[0]?.text ?? "";
                     const snapshotTokens = estimateTokens(snapshotText);
-                    recordWithTrace({ tool: 'session_snapshot', tokensReturned: snapshotTokens, tokensWouldBe: snapshotTokens, timestamp: Date.now(), savingsCategory: 'compression' });
+                    // TP-340: persist to .token-pilot/snapshots/ unless caller opts out.
+                    if (snapshotArgs.persist !== false) {
+                        try {
+                            await persistSnapshot({ projectRoot, body: snapshotText });
+                        }
+                        catch {
+                            /* best-effort — never fail the tool call */
+                        }
+                    }
+                    recordWithTrace({
+                        tool: "session_snapshot",
+                        tokensReturned: snapshotTokens,
+                        tokensWouldBe: snapshotTokens,
+                        timestamp: Date.now(),
+                        savingsCategory: "compression",
+                    });
                     return { content: snapshotResult.content };
                 }
+                case "session_budget": {
+                    const budgetArgs = args;
+                    const budgetResult = await handleSessionBudget({ sessionId: budgetArgs.sessionId ?? "" }, projectRoot, {
+                        baseThreshold: config.hooks.denyThreshold,
+                        adaptiveThreshold: config.hooks.adaptiveThreshold,
+                        adaptiveBudgetTokens: config.hooks.adaptiveBudgetTokens,
+                    });
+                    const budgetTokens = estimateTokens(budgetResult.content[0]?.text ?? "");
+                    recordWithTrace({
+                        tool: "session_budget",
+                        tokensReturned: budgetTokens,
+                        tokensWouldBe: budgetTokens,
+                        timestamp: Date.now(),
+                        savingsCategory: "compression",
+                    });
+                    return { content: budgetResult.content };
+                }
                 default:
                     return {
-                        content: [{ type: 'text', text: `Unknown tool: ${name}` }],
+                        content: [{ type: "text", text: `Unknown tool: ${name}` }],
                         isError: true,
                     };
             }
@@ -603,7 +880,7 @@ export async function createServer(projectRoot, options) {
         catch (err) {
             const message = err instanceof Error ? err.message : String(err);
             return {
-                content: [{ type: 'text', text: `Error: ${message}` }],
+                content: [{ type: "text", text: `Error: ${message}` }],
                 isError: true,
             };
         }