Merge pull request #40 from AxmeAI/feat/audit-dedup-20260407

fix: prevent duplicate sessions, audits, and stuck audit logs

Author: George-iam

src/cli.ts

@@ -403,6 +403,25 @@ async function main() {
       process.stderr.write(
         `axme-code audit-session: workspace=${workspacePath} session=${sessionId} pid=${process.pid}\n`,
       );
+      // Register signal handlers so SIGTERM/SIGINT (OOM killer, manual kill)
+      // updates audit status before exit. SIGKILL is uncatchable - handled
+      // by 15-minute stale timeout in runSessionCleanup.
+      const signalCleanup = (signal: string) => {
+        process.stderr.write(`axme-code audit-session: received ${signal}, cleaning up\n`);
+        try {
+          const { loadSession, writeSession } = require("./storage/sessions.js");
+          const s = loadSession(workspacePath, sessionId);
+          if (s && s.auditStatus === "pending") {
+            s.auditStatus = "failed";
+            s.lastAuditError = `killed by ${signal}`;
+            s.auditFinishedAt = new Date().toISOString();
+            writeSession(workspacePath, s);
+          }
+        } catch {}
+        process.exit(1);
+      };
+      process.on("SIGTERM", () => signalCleanup("SIGTERM"));
+      process.on("SIGINT", () => signalCleanup("SIGINT"));
       try {
         const { runSessionCleanup } = await import("./session-cleanup.js");
         const result = await runSessionCleanup(workspacePath, sessionId);

src/server.ts

@@ -31,6 +31,8 @@ import {
   clearLegacyPendingAuditsDir,
   readClaudeSessionMapping,
   isPidAlive,
+  loadSession,
+  closeSession,
 } from "./storage/sessions.js";
 import { logEvent } from "./storage/worklog.js";
 import { spawnDetachedAuditWorker } from "./audit-spawner.js";
@@ -138,10 +140,40 @@ async function cleanupAndExit(reason: string): Promise<void> {
   try {
     const mappings = listClaudeSessionMappings(defaultProjectPath);
     const owned = mappings.filter(m => m.ownerPpid === OWN_PPID);
+
+    // Deduplicate: group AXME sessions by Claude session ID.
+    // Multiple AXME sessions can share the same Claude session (race condition
+    // from parallel hooks). Only audit one per Claude session — the newest.
+    const claudeToAxme = new Map<string, { axmeId: string; createdAt: number }[]>();
+    for (const m of owned) {
+      const session = loadSession(defaultProjectPath, m.axmeSessionId);
+      if (!session) continue;
+      for (const ref of session.claudeSessions ?? []) {
+        const list = claudeToAxme.get(ref.id) ?? [];
+        list.push({ axmeId: m.axmeSessionId, createdAt: Date.parse(session.createdAt) || 0 });
+        claudeToAxme.set(ref.id, list);
+      }
+    }
+    const toAudit = new Set<string>();
+    const toSkip = new Set<string>();
+    for (const [, entries] of claudeToAxme) {
+      entries.sort((a, b) => b.createdAt - a.createdAt); // newest first
+      toAudit.add(entries[0].axmeId);
+      for (let i = 1; i < entries.length; i++) toSkip.add(entries[i].axmeId);
+    }
+    // Mark duplicates as done so they don't linger
+    for (const skipId of toSkip) {
+      try { closeSession(defaultProjectPath, skipId); } catch {}
+    }
+
     process.stderr.write(
-      `AXME cleanup (${reason}): ${owned.length} owned session(s) of ${mappings.length} total — spawning detached audit workers\n`,
+      `AXME cleanup (${reason}): ${owned.length} owned, ${toAudit.size} to audit, ${toSkip.size} deduped\n`,
     );
     for (const m of owned) {
+      if (!toAudit.has(m.axmeSessionId)) {
+        try { clearClaudeSessionMapping(defaultProjectPath, m.claudeSessionId); } catch {}
+        continue;
+      }
       try {
         spawnDetachedAuditWorker(defaultProjectPath, m.axmeSessionId);
       } catch (err) {

src/session-cleanup.ts

@@ -31,6 +31,7 @@ import {
   MAX_AUDIT_ATTEMPTS,
   isRetryableError,
   RETRYABLE_MAX_ATTEMPTS,
+  listSessions,
   type AuditLog,
   type AuditLogExtraction,
   type AuditLogResumeInfo,
@@ -230,6 +231,26 @@ export async function runSessionCleanup(
     session.auditAttempts = 0;
   }
 
+  // Dedup 2b: cross-session concurrent-audit protection.
+  // Multiple AXME sessions can share the same Claude session ID (parallel hook
+  // race condition). If another AXME session with the same Claude transcript
+  // is already pending audit (not stale), skip to avoid duplicate LLM calls.
+  const myClaudeIds = new Set((session.claudeSessions ?? []).map((c: any) => c.id));
+  if (myClaudeIds.size > 0) {
+    const allSessions = listSessions(workspacePath);
+    for (const other of allSessions) {
+      if (other.id === sessionId) continue;
+      if (other.auditStatus !== "pending" || !other.auditStartedAt) continue;
+      const startedMs = Date.parse(other.auditStartedAt);
+      if (!Number.isFinite(startedMs) || Date.now() - startedMs > AUDIT_STALE_TIMEOUT_MS) continue;
+      const otherClaudeIds = new Set((other.claudeSessions ?? []).map((c: any) => c.id));
+      const overlap = [...myClaudeIds].some(id => otherClaudeIds.has(id));
+      if (overlap) {
+        return { ...base, skipped: "concurrent-audit" };
+      }
+    }
+  }
+
   // Dedup 3: retry cap. If the session already used up its audit attempts
   // and still has no auditedAt, do NOT retry — it either hit a deterministic
   // failure (too-large prompt, parser rejection) or a bug that needs manual
@@ -387,6 +408,7 @@ export async function runSessionCleanup(
       // Audit log failure is non-fatal.
     }
 
+    let auditLogFinalized = false;
     try {
       const { runSessionAudit } = await import("./agents/session-auditor.js");
 
@@ -707,6 +729,7 @@ export async function runSessionCleanup(
             safetyDeduped: sDeduped,
           },
         });
+        auditLogFinalized = true;
       }
     } catch (err) {
       // Audit failure is non-fatal for the caller (we still close the session),
@@ -721,6 +744,20 @@ export async function runSessionCleanup(
           durationMs: Date.now() - auditStartMs,
           error: err instanceof Error ? err.message : String(err),
         });
+        auditLogFinalized = true;
+      }
+    } finally {
+      // Safety net: if neither success nor catch finalized the audit log
+      // (e.g., process received SIGTERM mid-LLM-call), mark it failed.
+      if (!auditLogFinalized && auditLogPath) {
+        try {
+          updateAuditLog(auditLogPath, {
+            phase: "failed",
+            finishedAt: new Date().toISOString(),
+            durationMs: Date.now() - auditStartMs,
+            error: "audit worker terminated unexpectedly",
+          });
+        } catch {}
       }
     }
   }

src/storage/sessions.ts

@@ -17,7 +17,7 @@
  */
 
 import { join, resolve } from "node:path";
-import { readdirSync, readFileSync, rmSync } from "node:fs";
+import { readdirSync, readFileSync, rmSync, openSync, closeSync, unlinkSync, statSync } from "node:fs";
 import { randomUUID } from "node:crypto";
 import { ensureDir, writeJson, readJson, pathExists, atomicWrite, removeFile, readSafe } from "./engine.js";
 import { logSessionStart } from "./worklog.js";
@@ -478,6 +478,54 @@ export function clearLegacyActiveSession(projectPath: string): void {
   removeFile(legacyActiveSessionPath(projectPath));
 }
 
+/**
+ * Filesystem lock for session creation to prevent parallel hooks from
+ * creating duplicate AXME sessions for the same Claude session.
+ *
+ * Uses O_EXCL (atomic create-or-fail) as a cross-process mutex.
+ * On contention: spin-wait up to 500ms, then re-read the mapping.
+ * Stale lock (>5s) auto-cleaned to handle crashed lock holders.
+ */
+const LOCK_STALE_MS = 5_000;
+const LOCK_WAIT_MS = 500;
+const LOCK_POLL_MS = 50;
+
+// Exported for testing
+export function sessionLockPath(projectPath: string, claudeSessionId: string): string {
+  return join(activeSessionsDir(projectPath), `${claudeSessionId}.lock`);
+}
+
+export function acquireLock(projectPath: string, claudeSessionId: string): boolean {
+  const lp = sessionLockPath(projectPath, claudeSessionId);
+  ensureDir(activeSessionsDir(projectPath));
+  try {
+    // Clean stale lock from crashed process
+    try {
+      const st = statSync(lp);
+      if (Date.now() - st.mtimeMs > LOCK_STALE_MS) unlinkSync(lp);
+    } catch {}
+    const fd = openSync(lp, "wx");
+    closeSync(fd);
+    return true;
+  } catch {
+    return false; // EEXIST or other - graceful degradation
+  }
+}
+
+export function releaseLock(projectPath: string, claudeSessionId: string): void {
+  try { unlinkSync(sessionLockPath(projectPath, claudeSessionId)); } catch {}
+}
+
+function waitForLock(projectPath: string, claudeSessionId: string): boolean {
+  const deadline = Date.now() + LOCK_WAIT_MS;
+  while (Date.now() < deadline) {
+    if (acquireLock(projectPath, claudeSessionId)) return true;
+    const start = Date.now();
+    while (Date.now() - start < LOCK_POLL_MS) { /* spin */ }
+  }
+  return false;
+}
+
 /**
  * Ensure an AXME session exists for the given Claude session. Lazy-created
  * on the first hook call that knows its Claude session_id.
@@ -507,6 +555,7 @@ export function ensureAxmeSessionForClaude(
    *  the existing session id instead of creating a fresh empty-tail session. */
   toolName?: string,
 ): string {
+  // Fast path: live mapping exists, just reuse it (no lock needed).
   const existing = readClaudeSessionMapping(projectPath, claudeSessionId);
   if (existing) {
     const existingSession = loadSession(projectPath, existing);
@@ -515,43 +564,56 @@ export function ensureAxmeSessionForClaude(
       existingSession.auditedAt != null ||
       (existingSession.pid != null && !isPidAlive(existingSession.pid));
     if (!isStale) {
-      // Live mapping — attach transcript and reuse.
       attachClaudeSession(projectPath, existing, {
         id: claudeSessionId,
         transcriptPath,
         role: "main",
       });
-      // Always refresh ownerPpid — after VS Code reload the Claude Code
-      // PID changes but the old process may still be alive (different
-      // window or zombie). The MCP server matches by ownerPpid === OWN_PPID,
-      // so the mapping must point to the current Claude Code instance.
       writeClaudeSessionMapping(projectPath, claudeSessionId, existing);
       return existing;
     }
-    // Read-only tools (Read/Glob/Grep) should not create fresh sessions from
-    // stale mappings — that produces empty "tail" sessions with 0 extractions.
-    // Return the stale id instead; the next mutation tool will create a fresh one.
+    // Read-only tools should not create fresh sessions from stale mappings.
     const READ_ONLY_TOOLS = ["Read", "Glob", "Grep"];
-    if (toolName && READ_ONLY_TOOLS.includes(toolName) && existing) {
+    if (toolName && READ_ONLY_TOOLS.includes(toolName)) {
       return existing;
     }
-    // Stale mapping: log once and fall through to create a fresh session,
-    // which will overwrite the mapping file below.
-    process.stderr.write(
-      `AXME: stale mapping for Claude session ${claudeSessionId} → ` +
-        `AXME ${existing} (audited=${existingSession?.auditedAt ?? "no"}, ` +
-        `pid=${existingSession?.pid ?? "?"}). Creating fresh AXME session.\n`,
-    );
   }
-  const axmeSession = createSession(projectPath);
-  try { logSessionStart(projectPath, axmeSession.id); } catch {}
-  writeClaudeSessionMapping(projectPath, claudeSessionId, axmeSession.id);
-  attachClaudeSession(projectPath, axmeSession.id, {
-    id: claudeSessionId,
-    transcriptPath,
-    role: "main",
-  });
-  return axmeSession.id;
+
+  // Slow path: need to create a new session (stale mapping OR first time).
+  // Acquire filesystem lock to prevent parallel hooks from each creating one.
+  const gotLock = waitForLock(projectPath, claudeSessionId);
+  try {
+    // Re-check inside lock — another process may have won the race.
+    // If a DIFFERENT mapping appeared (created by the lock winner), use it
+    // unconditionally — the winner just created it, so it's fresh by definition.
+    // Do NOT re-run stale checks here: the winner's process may have already
+    // exited (test workers, short-lived hooks), making the pid look dead.
+    const recheck = readClaudeSessionMapping(projectPath, claudeSessionId);
+    if (recheck && recheck !== existing) {
+      attachClaudeSession(projectPath, recheck, { id: claudeSessionId, transcriptPath, role: "main" });
+      return recheck;
+    }
+    // We won the race (or lock timed out) — create fresh session.
+    if (existing) {
+      const existingSession = loadSession(projectPath, existing);
+      process.stderr.write(
+        `AXME: stale mapping for Claude session ${claudeSessionId} → ` +
+          `AXME ${existing} (audited=${existingSession?.auditedAt ?? "no"}, ` +
+          `pid=${existingSession?.pid ?? "?"}). Creating fresh AXME session.\n`,
+      );
+    }
+    const axmeSession = createSession(projectPath);
+    try { logSessionStart(projectPath, axmeSession.id); } catch {}
+    writeClaudeSessionMapping(projectPath, claudeSessionId, axmeSession.id);
+    attachClaudeSession(projectPath, axmeSession.id, {
+      id: claudeSessionId,
+      transcriptPath,
+      role: "main",
+    });
+    return axmeSession.id;
+  } finally {
+    if (gotLock) releaseLock(projectPath, claudeSessionId);
+  }
 }
 
 // --- Legacy single-file API (DEPRECATED, kept for backward compatibility) ---