diff --git a/src/hooks/session-end.ts b/src/hooks/session-end.ts index e32edc0..99c4084 100644 --- a/src/hooks/session-end.ts +++ b/src/hooks/session-end.ts @@ -1,27 +1,23 @@ /** - * SessionEnd hook - runs when Claude session closes. + * SessionEnd hook — runs when Claude Code fires the SessionEnd lifecycle event. * - * Full session audit: - * 1. Reads session worklog + filesChanged - * 2. Runs session-auditor LLM (Sonnet, ~$0.30) - extracts ALL: - * memories, decisions, safety rules, oracle change detection - * 3. Saves everything to storage modules - * 4. If oracle needs re-scan: runs full Oracle Scanner - * 5. Closes session + * In practice this hook fires rarely and unreliably (especially in the VS Code + * extension, where the MCP server is killed without the extension first running + * SessionEnd — see anthropics/claude-code#1935 and #14760). The authoritative + * cleanup path is the MCP server's own transport.onclose handler, which calls + * the same `runSessionCleanup` function this hook does. + * + * The `auditedAt` dedup field ensures that whichever path runs first wins, and + * the others become no-ops. * * Workspace path: passed via --workspace flag (hardcoded at setup time). - * Session ID: read from .axme-code/active-session. + * Session ID: read from .axme-code/active-session (if present). */ -import { readWorklog, logSessionEnd } from "../storage/worklog.js"; -import { saveMemories } from "../storage/memory.js"; -import { addDecision, listDecisions } from "../storage/decisions.js"; -import { updateSafetyRule } from "../storage/safety.js"; -import { writeOracleFiles, oracleContext } from "../storage/oracle.js"; -import { writeHandoff } from "../storage/plans.js"; -import { closeSession, loadSession, readActiveSession, clearActiveSession } from "../storage/sessions.js"; -import { pathExists } from "../storage/engine.js"; import { join } from "node:path"; +import { readActiveSession } from "../storage/sessions.js"; +import { runSessionCleanup } from "../session-cleanup.js"; +import { pathExists } from "../storage/engine.js"; import { AXME_CODE_DIR } from "../types.js"; async function handleSessionEnd(workspacePath: string): Promise { @@ -30,65 +26,11 @@ async function handleSessionEnd(workspacePath: string): Promise { const sessionId = readActiveSession(workspacePath); if (!sessionId) return; - const session = loadSession(workspacePath, sessionId); - const filesChanged = session?.filesChanged ?? []; - const events = readWorklog(workspacePath, { limit: 200 }); - const sessionEvents = events - .filter(e => e.sessionId === sessionId) - .reverse() - .map(e => `[${e.timestamp}] ${e.type}: ${JSON.stringify(e.data)}`) - .join("\n"); - - if (sessionEvents.length > 50) { - try { - const { runSessionAudit } = await import("../agents/session-auditor.js"); - - const oracleSummary = oracleContext(workspacePath).slice(0, 500); - const decisionsCount = listDecisions(workspacePath).length; - - const audit = await runSessionAudit({ - sessionId, - sessionEvents, - filesChanged, - projectPath: workspacePath, - oracleSummary, - decisionsCount, - }); - - if (audit.memories.length > 0) saveMemories(workspacePath, audit.memories); - - for (const d of audit.decisions) addDecision(workspacePath, d); - - for (const r of audit.safetyRules) { - const validTypes = ["bash_deny", "bash_allow", "fs_deny", "git_protected_branch"] as const; - if (validTypes.includes(r.ruleType as any)) { - updateSafetyRule(workspacePath, r.ruleType as any, r.value); - } - } - - if (audit.handoff) writeHandoff(workspacePath, audit.handoff); - - if (audit.oracleNeedsRescan && filesChanged.length > 0) { - try { - const { runOracleScan } = await import("../agents/scanners/oracle.js"); - const oracleResult = await runOracleScan({ projectPath: workspacePath }); - writeOracleFiles(workspacePath, oracleResult.files); - } catch {} - } - } catch {} - } - - logSessionEnd(workspacePath, sessionId, { - turns: session?.turns ?? 0, - filesChanged, - }); - - closeSession(workspacePath, sessionId); - clearActiveSession(workspacePath); + await runSessionCleanup(workspacePath, sessionId, { clearActive: true }); } /** - * CLI entry point - reads JSON from stdin. + * CLI entry point — reads JSON from stdin. * @param workspacePath - from --workspace CLI flag */ export async function runSessionEndHook(workspacePath?: string): Promise { diff --git a/src/server.ts b/src/server.ts index f5d54b1..d44961e 100644 --- a/src/server.ts +++ b/src/server.ts @@ -17,8 +17,14 @@ import { saveDecisionTool } from "./tools/decision-tools.js"; import { updateSafetyTool, showSafetyTool } from "./tools/safety-tools.js"; import { statusTool, worklogTool } from "./tools/status.js"; import { detectWorkspace } from "./utils/workspace-detector.js"; -import { createSession, writeActiveSession, clearActiveSession, closeSession, loadSession, incrementTurns } from "./storage/sessions.js"; -import { logSessionStart, logSessionEnd } from "./storage/worklog.js"; +import { + createSession, + writeActiveSession, + incrementTurns, + findOrphanSessions, +} from "./storage/sessions.js"; +import { logSessionStart, logEvent } from "./storage/worklog.js"; +import { runSessionCleanup } from "./session-cleanup.js"; // --- Server state (detected at startup from cwd) --- @@ -49,24 +55,22 @@ function bumpTurn(): void { lastTurnBumpAt = now; } -// Clean up session on process exit -function onExit() { +// Session cleanup is triggered by transport.onclose (see main() below) rather +// than process.on("exit") — the async runSessionCleanup must finish before +// the MCP process exits, which "exit" handlers cannot await. +// SIGINT/SIGTERM are handled in main() to call the same cleanup path. + +let cleanupRunning = false; +async function cleanupAndExit(reason: string): Promise { + if (cleanupRunning) return; + cleanupRunning = true; try { - // Read latest session data from disk (hooks may have updated it) - const latest = loadSession(defaultProjectPath, currentSession.id); - closeSession(defaultProjectPath, currentSession.id); - clearActiveSession(defaultProjectPath); - logSessionEnd(defaultProjectPath, currentSession.id, { - turns: latest?.turns ?? 0, - filesChanged: latest?.filesChanged ?? [], - }); - } catch { - // Best-effort cleanup + await runSessionCleanup(defaultProjectPath, currentSession.id, { clearActive: true }); + } catch (err) { + process.stderr.write(`AXME session cleanup failed (${reason}): ${err}\n`); } + process.exit(0); } -process.on("exit", onExit); -process.on("SIGINT", () => { onExit(); process.exit(0); }); -process.on("SIGTERM", () => { onExit(); process.exit(0); }); // --- Build instructions for Claude Code --- @@ -284,7 +288,57 @@ server.tool( // --- Start server --- async function main() { const transport = new StdioServerTransport(); + + // Auto-audit on disconnect: when Claude Code closes the stdio pipe, stdin + // receives EOF. The MCP server process survives (Claude Code is known to + // not kill child MCP servers on exit, issue #1935), giving us time to run + // the full LLM audit before we exit ourselves. + // + // Note: we listen on process.stdin directly because MCP SDK's + // StdioServerTransport only handles 'data' and 'error' events — it does + // not react to stdin 'end'/'close', so transport.onclose never fires on + // its own. This bypasses that limitation. + process.stdin.on("end", () => { void cleanupAndExit("stdin-end"); }); + process.stdin.on("close", () => { void cleanupAndExit("stdin-close"); }); + + process.on("SIGINT", () => { void cleanupAndExit("sigint"); }); + process.on("SIGTERM", () => { void cleanupAndExit("sigterm"); }); + process.on("SIGHUP", () => { void cleanupAndExit("sighup"); }); + await server.connect(transport); + + // Startup fallback: audit any orphaned sessions left behind by previous + // MCP servers that were killed before auto-audit could run (e.g. SIGKILL + // from VS Code force-close). Runs in the background so it does not block + // server startup. + setTimeout(() => { + void auditOrphansInBackground(); + }, 3000); +} + +async function auditOrphansInBackground(): Promise { + try { + const orphans = findOrphanSessions(defaultProjectPath); + for (const orphan of orphans) { + if (orphan.id === currentSession.id) continue; // never touch self + try { + const result = await runSessionCleanup(defaultProjectPath, orphan.id); + logEvent(defaultProjectPath, "session_orphan_closed", orphan.id, { + turns: orphan.turns, + filesChanged: orphan.filesChanged.length, + auditRan: result.auditRan, + memories: result.memories, + decisions: result.decisions, + costUsd: result.costUsd, + closedBy: currentSession.id, + }); + } catch (err) { + process.stderr.write(`AXME orphan audit failed for ${orphan.id}: ${err}\n`); + } + } + } catch (err) { + process.stderr.write(`AXME orphan scan failed: ${err}\n`); + } } main().catch((err) => { diff --git a/src/session-cleanup.ts b/src/session-cleanup.ts new file mode 100644 index 0000000..5905364 --- /dev/null +++ b/src/session-cleanup.ts @@ -0,0 +1,180 @@ +/** + * Shared session cleanup logic — runs LLM audit and closes a session. + * + * Used by three entry points: + * 1. MCP server transport close handler (auto-audit on disconnect) + * 2. MCP server startup fallback (orphaned sessions from killed processes) + * 3. SessionEnd hook (if Claude Code fires it — rare but supported) + * + * All three paths set `auditedAt` on success so subsequent paths skip the + * same session. This dedup is the single source of truth for "has audit run?". + */ + +import { join } from "node:path"; +import { readWorklog, logSessionEnd } from "./storage/worklog.js"; +import { saveMemories } from "./storage/memory.js"; +import { addDecision, listDecisions } from "./storage/decisions.js"; +import { updateSafetyRule } from "./storage/safety.js"; +import { writeOracleFiles, oracleContext } from "./storage/oracle.js"; +import { writeHandoff } from "./storage/plans.js"; +import { + closeSession, + loadSession, + markAudited, + clearActiveSession, + readActiveSession, +} from "./storage/sessions.js"; +import { pathExists } from "./storage/engine.js"; +import { AXME_CODE_DIR } from "./types.js"; + +export interface SessionCleanupResult { + sessionId: string; + auditRan: boolean; + memories: number; + decisions: number; + safetyRules: number; + handoffSaved: boolean; + oracleRescanned: boolean; + costUsd: number; + skipped?: "already-audited" | "not-found" | "no-storage"; +} + +/** + * Run full session cleanup: LLM audit + save artifacts + close session. + * + * Idempotent: if the session is already audited (`auditedAt` set), returns + * early with `skipped: "already-audited"` and does nothing. + * + * If the session has insufficient activity (worklog events shorter than 50 + * chars), skips the LLM audit but still closes the session cleanly. + * + * @param workspacePath - Absolute path to the AXME workspace/project root + * @param sessionId - UUID of the session to audit and close + * @param opts.clearActive - If true, also clears .axme-code/active-session (only the + * currently-active MCP server should do this) + */ +export async function runSessionCleanup( + workspacePath: string, + sessionId: string, + opts: { clearActive?: boolean } = {}, +): Promise { + const base: SessionCleanupResult = { + sessionId, + auditRan: false, + memories: 0, + decisions: 0, + safetyRules: 0, + handoffSaved: false, + oracleRescanned: false, + costUsd: 0, + }; + + if (!pathExists(join(workspacePath, AXME_CODE_DIR))) { + return { ...base, skipped: "no-storage" }; + } + + const session = loadSession(workspacePath, sessionId); + if (!session) { + return { ...base, skipped: "not-found" }; + } + + // Dedup: if audit already ran, don't repeat. Just ensure session is closed. + if (session.auditedAt) { + if (!session.closedAt) closeSession(workspacePath, sessionId); + if (opts.clearActive && readActiveSession(workspacePath) === sessionId) { + clearActiveSession(workspacePath); + } + return { ...base, skipped: "already-audited" }; + } + + const filesChanged = session.filesChanged ?? []; + + const events = readWorklog(workspacePath, { limit: 500 }); + const sessionEvents = events + .filter(e => e.sessionId === sessionId) + .reverse() + .map(e => `[${e.timestamp}] ${e.type}: ${JSON.stringify(e.data)}`) + .join("\n"); + + const result: SessionCleanupResult = { ...base }; + let auditSucceeded = false; + const hasActivity = sessionEvents.length > 50; + + // Run LLM audit only if there's meaningful activity to analyze + if (hasActivity) { + try { + const { runSessionAudit } = await import("./agents/session-auditor.js"); + + const oracleSummary = oracleContext(workspacePath).slice(0, 500); + const decisionsCount = listDecisions(workspacePath).length; + + const audit = await runSessionAudit({ + sessionId, + sessionEvents, + filesChanged, + projectPath: workspacePath, + oracleSummary, + decisionsCount, + }); + + if (audit.memories.length > 0) saveMemories(workspacePath, audit.memories); + for (const d of audit.decisions) addDecision(workspacePath, d); + + for (const r of audit.safetyRules) { + const validTypes = ["bash_deny", "bash_allow", "fs_deny", "git_protected_branch"] as const; + if (validTypes.includes(r.ruleType as any)) { + updateSafetyRule(workspacePath, r.ruleType as any, r.value); + } + } + + if (audit.handoff) { + writeHandoff(workspacePath, audit.handoff); + result.handoffSaved = true; + } + + if (audit.oracleNeedsRescan && filesChanged.length > 0) { + try { + const { runOracleScan } = await import("./agents/scanners/oracle.js"); + const oracleResult = await runOracleScan({ projectPath: workspacePath }); + writeOracleFiles(workspacePath, oracleResult.files); + result.oracleRescanned = true; + } catch { + // Oracle rescan failure is non-fatal — audit artifacts already saved + } + } + + result.auditRan = true; + result.memories = audit.memories.length; + result.decisions = audit.decisions.length; + result.safetyRules = audit.safetyRules.length; + result.costUsd = audit.cost?.costUsd ?? 0; + auditSucceeded = true; + } catch { + // Audit failure is non-fatal. We deliberately leave auditedAt null so + // the startup fallback can retry on the next MCP server start. + } + } + + // Mark audited only on success (or when there was nothing to audit). + // Failed audits leave auditedAt null so startup fallback retries them. + if (auditSucceeded || !hasActivity) { + markAudited(workspacePath, sessionId); + } + + // Always close the session — leaving closedAt=null would cause the MCP + // server to appear mid-run on exit, or the startup fallback to repeatedly + // see the session as a pending candidate even if audit already succeeded. + closeSession(workspacePath, sessionId); + + logSessionEnd(workspacePath, sessionId, { + turns: session.turns, + filesChanged, + auditRan: result.auditRan, + }); + + if (opts.clearActive && readActiveSession(workspacePath) === sessionId) { + clearActiveSession(workspacePath); + } + + return result; +} diff --git a/src/storage/sessions.ts b/src/storage/sessions.ts index 19d3c64..c74ca09 100644 --- a/src/storage/sessions.ts +++ b/src/storage/sessions.ts @@ -63,6 +63,7 @@ export function createSession(projectPath: string): SessionMeta { closedAt: null, turns: 0, filesChanged: [], + pid: process.pid, }; writeSession(projectPath, session); return session; @@ -85,6 +86,59 @@ export function closeSession(projectPath: string, id: string): void { writeSession(projectPath, session); } +/** + * Mark a session as audited by the LLM session auditor. + * Used by both auto-audit (transport close) and startup fallback + * to prevent duplicate audits on the same session. + */ +export function markAudited(projectPath: string, id: string): void { + const session = loadSession(projectPath, id); + if (!session) return; + session.auditedAt = new Date().toISOString(); + writeSession(projectPath, session); +} + +/** + * Check if a process with the given PID is currently running. + * Uses signal 0 (no-op signal) to probe process existence. + * + * Returns true if process is alive, false if dead. + * EPERM (permission denied) is treated as alive — the process exists + * but belongs to another user, which is extremely rare in our context + * and safer to treat as alive than as dead. + */ +export function isPidAlive(pid: number): boolean { + try { + process.kill(pid, 0); + return true; + } catch (err: any) { + if (err?.code === "EPERM") return true; + return false; + } +} + +/** + * Find sessions that still need an LLM audit: auditedAt is null AND + * their owning MCP server process is no longer running. + * + * Sessions without a pid field (pre-auto-audit format) are skipped — + * we cannot determine if they are orphans without a PID to probe. + * + * Note: closedAt is intentionally NOT checked. A session may have + * closedAt set but auditedAt null if its auto-audit failed — these + * need a retry on the next startup. + */ +export function findOrphanSessions(projectPath: string): SessionMeta[] { + const orphans: SessionMeta[] = []; + for (const session of listSessions(projectPath)) { + if (session.auditedAt) continue; + if (session.pid == null) continue; + if (isPidAlive(session.pid)) continue; + orphans.push(session); + } + return orphans; +} + export function listSessions(projectPath: string, opts?: { limit?: number }): SessionMeta[] { const root = sessionsRoot(projectPath); const sessions: SessionMeta[] = []; @@ -108,27 +162,18 @@ export function getLastSession(projectPath: string): SessionMeta | null { } /** - * Ensure a session exists - create it if not found. - * Used by hooks that receive session_id from Claude Code harness. + * Track a file change in an existing session. + * + * Called from the PostToolUse hook, which runs in a separate process from + * the MCP server. If the session file cannot be read (transient I/O error + * or missing), this function silently returns instead of creating a new + * session — recreating would destroy the real session's turns counter and + * filesChanged list. Losing a single filesChanged entry is a far better + * tradeoff than resetting the session to turns=0. */ -export function ensureSession(projectPath: string, id: string): SessionMeta { - const existing = loadSession(projectPath, id); - if (existing) return existing; - - initSessionStore(projectPath); - const session: SessionMeta = { - id, - createdAt: new Date().toISOString(), - closedAt: null, - turns: 0, - filesChanged: [], - }; - writeSession(projectPath, session); - return session; -} - export function trackFileChanged(projectPath: string, sessionId: string, filePath: string): void { - const session = ensureSession(projectPath, sessionId); + const session = loadSession(projectPath, sessionId); + if (!session) return; if (!session.filesChanged.includes(filePath)) { session.filesChanged.push(filePath); writeSession(projectPath, session); diff --git a/src/types.ts b/src/types.ts index ef4d969..f7a8395 100644 --- a/src/types.ts +++ b/src/types.ts @@ -140,6 +140,7 @@ export interface FilesystemRules { export type WorklogEventType = | "session_start" | "session_end" + | "session_orphan_closed" | "agent_turn" | "check_result" | "memory_saved" @@ -160,6 +161,10 @@ export interface SessionMeta { closedAt: string | null; turns: number; filesChanged: string[]; + /** PID of MCP server process. Used to detect orphaned sessions after crashes. Optional for backward compat. */ + pid?: number; + /** ISO timestamp when LLM session audit completed. Used to dedupe auto-audit vs startup fallback. */ + auditedAt?: string; } // --- Config ---