From be54229bba1de057a8a418fe39706c74e21a932c Mon Sep 17 00:00:00 2001 From: agammann <161159040+agammann@users.noreply.github.com> Date: Thu, 1 Oct 2026 17:48:18 -0700 Subject: [PATCH] fix(session): attribute MCP work to caller session --- src/server.ts | 18 ++--- src/session/caller-session.ts | 29 ++++++++ src/session/event-emit.ts | 10 +-- src/session/persist-tool-calls.ts | 11 +-- tests/session/caller-session.test.ts | 105 +++++++++++++++++++++++++++ 5 files changed, 154 insertions(+), 19 deletions(-) create mode 100644 src/session/caller-session.ts create mode 100644 tests/session/caller-session.test.ts diff --git a/src/server.ts b/src/server.ts index ff2767536..1bb03d66a 100644 --- a/src/server.ts +++ b/src/server.ts @@ -56,6 +56,7 @@ import { emitSandboxExecuteEvent, } from "./session/event-emit.js"; import { persistToolCallCounter, restoreSessionStats } from "./session/persist-tool-calls.js"; +import { resolveCallerSessionId } from "./session/caller-session.js"; import { searchAllSources } from "./search/unified.js"; import { buildCtxSearchInputSchema, @@ -440,8 +441,7 @@ let _store: ContentStore | null = null; /** * Build the FK-attribution object passed to every ContentStore.index*() call - * in this process. CLAUDE_SESSION_ID is the only MCP-side handle we have on - * the current session — eventId stays undefined because MCP tool invocations + * in this process. eventId stays undefined because MCP tool invocations * are not paired with PostToolUse event rows at index time (the hook fires * AFTER the tool returns). Empty-string fallback inside #insertChunks keeps * legacy unattributed rows readable. @@ -450,13 +450,13 @@ export function currentAttribution(): { sessionId?: string } | undefined { const override = projectDirOverride.getStore(); if (override?.sessionId) return { sessionId: override.sessionId }; - // CLAUDE_SESSION_ID env var is NOT propagated to MCP servers (only to hooks). - // Cross-adapter resolution: every adapter (15 of them) sets *_PROJECT_DIR env - // and writes session_events via hooks. Read the most-recent session_id from - // THIS project's session DB. Works for claude-code/cursor/gemini-cli/codex/ + // Resolve Claude Code's caller first: the session file tracks /clear and + // resume even when the inherited env is stale. Other adapters set + // *_PROJECT_DIR env and write session_events via hooks, so fall back to the + // most-recent session_id from THIS project's DB. Works for cursor/gemini-cli/codex/ // kiro/opencode/zed/kilo/openclaw/qwen-code/vscode-copilot/jetbrains-copilot/ // omp/pi/antigravity — no adapter-specific transcript path required. - const sessionId = process.env.CLAUDE_SESSION_ID ?? resolveSessionIdFromSessionDB(); + const sessionId = resolveCallerSessionId() ?? resolveSessionIdFromSessionDB(); if (!sessionId) return undefined; return { sessionId }; } @@ -4463,12 +4463,12 @@ server.registerTool( // narrative 5-section "kitap gibi" layout (timeline, ladder, receipt, // example cost, auto-memory). Without these, formatReport falls back // to the legacy active-session header. Best-effort — failures absorbed. - // Resolve session_id: prefer env (CLAUDE_SESSION_ID), else most-recent + // Resolve the caller's session before falling back to the most-recent // UUID session_id from session_events in this DB. let conversation; let realBytes; try { - let sid = process.env.CLAUDE_SESSION_ID; + let sid = resolveCallerSessionId(); if (!sid) { const row = sdb.prepare( "SELECT session_id FROM session_events WHERE session_id LIKE '________-____-____-____-____________' ORDER BY created_at DESC LIMIT 1" diff --git a/src/session/caller-session.ts b/src/session/caller-session.ts new file mode 100644 index 000000000..c23a27a75 --- /dev/null +++ b/src/session/caller-session.ts @@ -0,0 +1,29 @@ +import { readFileSync } from "node:fs"; +import { join } from "node:path"; +import { resolveClaudeConfigDir } from "../util/claude-config.js"; + +/** Resolve the Claude Code session that owns this MCP process, if available. */ +export function resolveCallerSessionId(opts: { + env?: NodeJS.ProcessEnv; + parentPid?: number; + configDir?: string; +} = {}): string | undefined { + const env = opts.env ?? process.env; + if (env.CLAUDE_SESSION_ID) return env.CLAUDE_SESSION_ID; + + // Claude Code updates this file after /clear and resume, while the MCP + // process keeps the environment it inherited when it first started. + try { + const file = join( + opts.configDir ?? resolveClaudeConfigDir(env), + "sessions", + `${opts.parentPid ?? process.ppid}.json`, + ); + const sessionId: unknown = JSON.parse(readFileSync(file, "utf8")).sessionId; + if (typeof sessionId === "string" && sessionId.trim()) return sessionId; + } catch { + // Other hosts do not have Claude Code's per-process session file. + } + + return env.CLAUDE_CODE_SESSION_ID || undefined; +} diff --git a/src/session/event-emit.ts b/src/session/event-emit.ts index a34ea4c66..d2e8e6475 100644 --- a/src/session/event-emit.ts +++ b/src/session/event-emit.ts @@ -14,16 +14,16 @@ * - Best-effort error swallowing matches `persistToolCallCounter` in * `persist-tool-calls.ts`. A stats-side failure must NEVER break the * parent MCP tool call. - * - Resolves the latest `session_id` from `session_meta` so the wiring - * in `server.ts` is `setImmediate(() => emit*({...}))` — no need to - * plumb session ids through every handler. + * - Resolves the caller's session id when Claude Code exposes it, with the + * latest `session_meta` row as a fallback for other hosts. */ import { existsSync } from "node:fs"; import { SessionDB } from "./db.js"; +import { resolveCallerSessionId } from "./caller-session.js"; /** - * Open the SessionDB at `dbPath`, find the latest session_id, and run + * Open the SessionDB at `dbPath`, find the caller's session_id, and run * `fn` with both. Wraps everything in try/catch so callers stay * fire-and-forget. */ @@ -35,7 +35,7 @@ function withLatestSession( if (!existsSync(dbPath)) return; const sdb = new SessionDB({ dbPath }); try { - const sid = sdb.getLatestSessionId(); + const sid = resolveCallerSessionId() ?? sdb.getLatestSessionId(); if (!sid) return; fn(sdb, sid); } finally { diff --git a/src/session/persist-tool-calls.ts b/src/session/persist-tool-calls.ts index 8a0c319bd..1e46d5d5e 100644 --- a/src/session/persist-tool-calls.ts +++ b/src/session/persist-tool-calls.ts @@ -20,6 +20,7 @@ import { existsSync } from "node:fs"; import { SessionDB } from "./db.js"; +import { resolveCallerSessionId } from "./caller-session.js"; /** * Shape returned by {@link restoreSessionStats}. Subset of the in-memory @@ -40,8 +41,8 @@ export interface RestoredSessionStats { } /** - * Increment the persistent tool-call counter for `toolName` under whatever - * session_id `session_meta` currently treats as the most recent. This is + * Increment the persistent tool-call counter for `toolName` under the + * caller's session id, falling back to the most recent session. This is * called from {@link trackResponse} on every tool response and must be * cheap, non-throwing, and best-effort — a stats failure must never break * the MCP tool call. @@ -55,7 +56,7 @@ export function persistToolCallCounter( if (!existsSync(sessionDbPath)) return; const sdb = new SessionDB({ dbPath: sessionDbPath }); try { - const sid = sdb.getLatestSessionId(); + const sid = resolveCallerSessionId() ?? sdb.getLatestSessionId(); if (!sid) return; sdb.incrementToolCall(sid, toolName, bytes); } finally { @@ -67,7 +68,7 @@ export function persistToolCallCounter( } /** - * Read the latest session's tool-call totals back out of SessionDB so the + * Read the caller's tool-call totals back out of SessionDB so the * MCP server can hydrate its in-memory `sessionStats` on startup. Returns * `null` when the DB is missing or empty so the caller can keep the * default zero-state without branching twice. @@ -82,7 +83,7 @@ export function restoreSessionStats( if (!existsSync(sessionDbPath)) return null; const sdb = new SessionDB({ dbPath: sessionDbPath }); try { - const sid = sdb.getLatestSessionId(); + const sid = resolveCallerSessionId() ?? sdb.getLatestSessionId(); if (!sid) return null; const stats = sdb.getToolCallStats(sid); diff --git a/tests/session/caller-session.test.ts b/tests/session/caller-session.test.ts new file mode 100644 index 000000000..48f2abd4a --- /dev/null +++ b/tests/session/caller-session.test.ts @@ -0,0 +1,105 @@ +import { mkdtempSync, mkdirSync, rmSync, writeFileSync } from "node:fs"; +import { tmpdir } from "node:os"; +import { join } from "node:path"; +import { afterEach, describe, expect, test } from "vitest"; +import { loadDatabase } from "../../src/db-base.js"; +import { resolveCallerSessionId } from "../../src/session/caller-session.js"; +import { SessionDB } from "../../src/session/db.js"; +import { emitIndexWriteEvent } from "../../src/session/event-emit.js"; +import { persistToolCallCounter, restoreSessionStats } from "../../src/session/persist-tool-calls.js"; + +const tempDirs: string[] = []; +afterEach(() => { + for (const dir of tempDirs.splice(0)) rmSync(dir, { recursive: true, force: true }); +}); + +function tempDir(): string { + const dir = mkdtempSync(join(tmpdir(), "caller-session-")); + tempDirs.push(dir); + return dir; +} + +describe("resolveCallerSessionId", () => { + test("prefers explicit override over the live file and inherited environment", () => { + const configDir = tempDir(); + mkdirSync(join(configDir, "sessions")); + writeFileSync(join(configDir, "sessions", "123.json"), JSON.stringify({ sessionId: "current" })); + expect(resolveCallerSessionId({ + configDir, + parentPid: 123, + env: { CLAUDE_SESSION_ID: "override", CLAUDE_CODE_SESSION_ID: "stale" }, + })).toBe("override"); + }); + + test("uses the live file after /clear instead of the inherited session id", () => { + const configDir = tempDir(); + mkdirSync(join(configDir, "sessions")); + const file = join(configDir, "sessions", "123.json"); + writeFileSync(file, JSON.stringify({ sessionId: "before-clear" })); + const opts = { configDir, parentPid: 123, env: { CLAUDE_CODE_SESSION_ID: "before-clear" } }; + expect(resolveCallerSessionId(opts)).toBe("before-clear"); + writeFileSync(file, JSON.stringify({ sessionId: "after-clear" })); + expect(resolveCallerSessionId(opts)).toBe("after-clear"); + }); + + test("falls back to the inherited ID when the file is absent or invalid", () => { + const configDir = tempDir(); + const opts = { configDir, parentPid: 123, env: { CLAUDE_CODE_SESSION_ID: "inherited" } }; + expect(resolveCallerSessionId(opts)).toBe("inherited"); + mkdirSync(join(configDir, "sessions")); + writeFileSync(join(configDir, "sessions", "123.json"), "{not json"); + expect(resolveCallerSessionId(opts)).toBe("inherited"); + expect(resolveCallerSessionId({ configDir, parentPid: 123, env: {} })).toBeUndefined(); + }); +}); + +test("event and tool-call stats stay with the caller when a newer session exists", () => { + const configDir = tempDir(); + const dbPath = join(configDir, "sessions.db"); + const db = new SessionDB({ dbPath }); + db.ensureSession("caller", configDir); + db.ensureSession("neighbor", configDir); + db.close(); + + const Database = loadDatabase(); + const raw = new Database(dbPath); + raw.prepare("UPDATE session_meta SET started_at = ? WHERE session_id = ?").run("2026-01-01 00:00:00", "caller"); + raw.prepare("UPDATE session_meta SET started_at = ? WHERE session_id = ?").run("2026-01-02 00:00:00", "neighbor"); + raw.close(); + + mkdirSync(join(configDir, "sessions")); + writeFileSync( + join(configDir, "sessions", `${process.ppid}.json`), + JSON.stringify({ sessionId: "caller" }), + ); + + const previous = { + explicit: process.env.CLAUDE_SESSION_ID, + inherited: process.env.CLAUDE_CODE_SESSION_ID, + configDir: process.env.CLAUDE_CONFIG_DIR, + }; + try { + delete process.env.CLAUDE_SESSION_ID; + process.env.CLAUDE_CODE_SESSION_ID = "neighbor"; + process.env.CLAUDE_CONFIG_DIR = configDir; + persistToolCallCounter(dbPath, "ctx_search", 42); + emitIndexWriteEvent({ sessionDbPath: dbPath, source: "sample", bytesAvoided: 64 }); + + expect(restoreSessionStats(dbPath)?.calls.ctx_search).toBe(1); + const check = new Database(dbPath, { readonly: true }); + try { + expect(check.prepare("SELECT session_id FROM tool_calls").all()).toEqual([{ session_id: "caller" }]); + expect(check.prepare("SELECT session_id FROM session_events WHERE type = 'index-write'").all()) + .toEqual([{ session_id: "caller" }]); + } finally { + check.close(); + } + } finally { + if (previous.explicit === undefined) delete process.env.CLAUDE_SESSION_ID; + else process.env.CLAUDE_SESSION_ID = previous.explicit; + if (previous.inherited === undefined) delete process.env.CLAUDE_CODE_SESSION_ID; + else process.env.CLAUDE_CODE_SESSION_ID = previous.inherited; + if (previous.configDir === undefined) delete process.env.CLAUDE_CONFIG_DIR; + else process.env.CLAUDE_CONFIG_DIR = previous.configDir; + } +});