fix: unified session lifecycle - hooks use workspace path and active-…

…session

Three problems fixed:

1. cwd drift: hooks received cwd from Claude Code which changed
   based on shell cd commands. Now --workspace is hardcoded into
   hook commands at setup time, always pointing to workspace root.

2. Session ID mismatch: MCP server generated its own UUID, Claude Code
   passed a different session_id. Now MCP server writes active-session
   file, hooks read it. No dependency on Claude Code's session_id.

3. Storage location split: sessions were created in repo/.axme-code/
   instead of workspace/.axme-code/. Fixed by using workspace path
   consistently in all hooks.

Additional:
- Process exit handler on MCP server closes session + removes active-session
- setup regenerates hook config to include --workspace flag
- ensureSession preserved for robustness but no longer primary path

src/cli.ts

@@ -118,20 +118,28 @@ function configureHooks(projectPath: string): void {
     try { settings = JSON.parse(readFileSync(settingsPath, "utf-8")); } catch { settings = {}; }
   }
 
-  // Check if hooks already configured
-  if (settings.hooks?.PostToolUse?.some?.((h: any) => JSON.stringify(h).includes("axme-code"))) {
-    return; // already configured
+  // Remove old hooks (without --workspace) and re-create with correct path
+  if (settings.hooks?.PostToolUse) {
+    settings.hooks.PostToolUse = settings.hooks.PostToolUse.filter(
+      (h: any) => !JSON.stringify(h).includes("axme-code"),
+    );
+  }
+  if (settings.hooks?.SessionEnd) {
+    settings.hooks.SessionEnd = settings.hooks.SessionEnd.filter(
+      (h: any) => !JSON.stringify(h).includes("axme-code"),
+    );
   }
 
   if (!settings.hooks) settings.hooks = {};
 
   // PostToolUse: track filesChanged after Edit/Write
+  // --workspace is hardcoded so hooks always write to workspace root, regardless of cwd
   if (!settings.hooks.PostToolUse) settings.hooks.PostToolUse = [];
   settings.hooks.PostToolUse.push({
     matcher: "Edit|Write|NotebookEdit",
     hooks: [{
       type: "command",
-      command: "axme-code hook post-tool-use",
+      command: `axme-code hook post-tool-use --workspace ${projectPath}`,
       timeout: 10,
     }],
   });
@@ -141,7 +149,7 @@ function configureHooks(projectPath: string): void {
   settings.hooks.SessionEnd.push({
     hooks: [{
       type: "command",
-      command: "axme-code hook session-end",
+      command: `axme-code hook session-end --workspace ${projectPath}`,
       timeout: 120,
     }],
   });
@@ -292,12 +300,16 @@ async function main() {
 
     case "hook": {
       const hookName = args[1];
+      // Parse --workspace flag from CLI args
+      const wsIdx = args.indexOf("--workspace");
+      const workspacePath = wsIdx >= 0 && args[wsIdx + 1] ? args[wsIdx + 1] : undefined;
+
       if (hookName === "post-tool-use") {
         const { runPostToolUseHook } = await import("./hooks/post-tool-use.js");
-        await runPostToolUseHook();
+        await runPostToolUseHook(workspacePath);
       } else if (hookName === "session-end") {
         const { runSessionEndHook } = await import("./hooks/session-end.js");
-        await runSessionEndHook();
+        await runSessionEndHook(workspacePath);
       }
       break;
     }

src/hooks/post-tool-use.ts

@@ -2,46 +2,53 @@
  * PostToolUse hook - runs after Edit/Write tool calls.
  *
  * Tracks filesChanged in session metadata.
+ *
  * Input: JSON on stdin from Claude Code hooks system.
- * Format: { session_id, cwd, hook_event_name, tool_name, tool_input: { file_path } }
+ * Workspace path: passed via --workspace flag (hardcoded at setup time).
+ *
+ * Session ID: read from .axme-code/active-session (written by MCP server),
+ * NOT from Claude Code's session_id (which is a different ID).
  */
 
 import { trackFileChanged } from "../storage/sessions.js";
+import { readActiveSession } from "../storage/sessions.js";
 import { pathExists } from "../storage/engine.js";
 import { join } from "node:path";
 import { AXME_CODE_DIR } from "../types.js";
 
 interface HookInput {
-  session_id: string;
-  cwd: string;
   tool_name: string;
   tool_input: Record<string, any>;
 }
 
-function handlePostToolUse(event: HookInput): void {
-  const { tool_name, tool_input, session_id, cwd } = event;
+function handlePostToolUse(workspacePath: string, event: HookInput): void {
+  const { tool_name, tool_input } = event;
 
   if (!["Edit", "Write", "NotebookEdit"].includes(tool_name)) return;
-  if (!pathExists(join(cwd, AXME_CODE_DIR))) return;
+  if (!pathExists(join(workspacePath, AXME_CODE_DIR))) return;
 
   const filePath = tool_input.file_path || tool_input.path;
   if (!filePath || typeof filePath !== "string") return;
   if (filePath.includes(AXME_CODE_DIR)) return;
 
-  if (session_id) {
-    trackFileChanged(cwd, session_id, filePath);
+  const sessionId = readActiveSession(workspacePath);
+  if (sessionId) {
+    trackFileChanged(workspacePath, sessionId, filePath);
   }
 }
 
 /**
  * CLI entry point - reads JSON from stdin.
+ * @param workspacePath - from --workspace CLI flag
  */
-export async function runPostToolUseHook(): Promise<void> {
+export async function runPostToolUseHook(workspacePath?: string): Promise<void> {
+  if (!workspacePath) return; // No workspace = nothing to do
+
   try {
     const chunks: Buffer[] = [];
     for await (const chunk of process.stdin) chunks.push(chunk);
     const input = JSON.parse(Buffer.concat(chunks).toString("utf-8")) as HookInput;
-    handlePostToolUse(input);
+    handlePostToolUse(workspacePath, input);
   } catch {
     // Hook failures must be silent
   }

src/hooks/session-end.ts

@@ -9,37 +9,31 @@
  * 4. If oracle needs re-scan: runs full Oracle Scanner
  * 5. Closes session
  *
- * Input: JSON on stdin from Claude Code hooks system.
- * Format: { session_id, cwd, hook_event_name }
+ * Workspace path: passed via --workspace flag (hardcoded at setup time).
+ * Session ID: read from .axme-code/active-session.
  */
 
 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 { closeSession, loadSession } from "../storage/sessions.js";
+import { closeSession, loadSession, readActiveSession, clearActiveSession } from "../storage/sessions.js";
 import { pathExists } from "../storage/engine.js";
 import { join } from "node:path";
 import { AXME_CODE_DIR } from "../types.js";
 
-interface HookInput {
-  session_id: string;
-  cwd: string;
-  hook_event_name: string;
-  transcript_path?: string;
-}
-
-async function handleSessionEnd(input: HookInput): Promise<void> {
-  const { session_id, cwd } = input;
+async function handleSessionEnd(workspacePath: string): Promise<void> {
+  if (!pathExists(join(workspacePath, AXME_CODE_DIR))) return;
 
-  if (!pathExists(join(cwd, AXME_CODE_DIR))) return;
+  const sessionId = readActiveSession(workspacePath);
+  if (!sessionId) return;
 
-  const session = loadSession(cwd, session_id);
+  const session = loadSession(workspacePath, sessionId);
   const filesChanged = session?.filesChanged ?? [];
-  const events = readWorklog(cwd, { limit: 200 });
+  const events = readWorklog(workspacePath, { limit: 200 });
   const sessionEvents = events
-    .filter(e => e.sessionId === session_id)
+    .filter(e => e.sessionId === sessionId)
     .reverse()
     .map(e => `[${e.timestamp}] ${e.type}: ${JSON.stringify(e.data)}`)
     .join("\n");
@@ -48,56 +42,60 @@ async function handleSessionEnd(input: HookInput): Promise<void> {
     try {
       const { runSessionAudit } = await import("../agents/session-auditor.js");
 
-      const oracleSummary = oracleContext(cwd).slice(0, 500);
-      const decisionsCount = listDecisions(cwd).length;
+      const oracleSummary = oracleContext(workspacePath).slice(0, 500);
+      const decisionsCount = listDecisions(workspacePath).length;
 
       const audit = await runSessionAudit({
-        sessionId: session_id,
+        sessionId,
         sessionEvents,
         filesChanged,
-        projectPath: cwd,
+        projectPath: workspacePath,
         oracleSummary,
         decisionsCount,
       });
 
-      if (audit.memories.length > 0) saveMemories(cwd, audit.memories);
+      if (audit.memories.length > 0) saveMemories(workspacePath, audit.memories);
 
-      for (const d of audit.decisions) addDecision(cwd, d);
+      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(cwd, r.ruleType as any, r.value);
+          updateSafetyRule(workspacePath, r.ruleType as any, r.value);
         }
       }
 
       if (audit.oracleNeedsRescan && filesChanged.length > 0) {
         try {
           const { runOracleScan } = await import("../agents/scanners/oracle.js");
-          const oracleResult = await runOracleScan({ projectPath: cwd });
-          writeOracleFiles(cwd, oracleResult.files);
+          const oracleResult = await runOracleScan({ projectPath: workspacePath });
+          writeOracleFiles(workspacePath, oracleResult.files);
         } catch {}
       }
     } catch {}
   }
 
-  logSessionEnd(cwd, session_id, {
+  logSessionEnd(workspacePath, sessionId, {
     turns: session?.turns ?? 0,
     filesChanged,
   });
 
-  closeSession(cwd, session_id);
+  closeSession(workspacePath, sessionId);
+  clearActiveSession(workspacePath);
 }
 
 /**
  * CLI entry point - reads JSON from stdin.
+ * @param workspacePath - from --workspace CLI flag
  */
-export async function runSessionEndHook(): Promise<void> {
+export async function runSessionEndHook(workspacePath?: string): Promise<void> {
+  if (!workspacePath) return;
+
   try {
+    // Still consume stdin (Claude Code sends it), but we don't need its content
     const chunks: Buffer[] = [];
     for await (const chunk of process.stdin) chunks.push(chunk);
-    const input = JSON.parse(Buffer.concat(chunks).toString("utf-8")) as HookInput;
-    await handleSessionEnd(input);
+    await handleSessionEnd(workspacePath);
   } catch {
     // Hook failures must be silent
   }

src/server.ts

@@ -17,8 +17,8 @@ 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 } from "./storage/sessions.js";
-import { logSessionStart } from "./storage/worklog.js";
+import { createSession, writeActiveSession, clearActiveSession, closeSession, loadSession } from "./storage/sessions.js";
+import { logSessionStart, logSessionEnd } from "./storage/worklog.js";
 
 // --- Server state (detected at startup from cwd) ---
 
@@ -31,8 +31,28 @@ const defaultWorkspacePath = isWorkspace ? serverCwd : null;
 // --- Session (one MCP server instance = one session) ---
 
 const currentSession = createSession(defaultProjectPath);
+writeActiveSession(defaultProjectPath, currentSession.id);
 logSessionStart(defaultProjectPath, currentSession.id);
 
+// Clean up session on process exit
+function onExit() {
+  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
+  }
+}
+process.on("exit", onExit);
+process.on("SIGINT", () => { onExit(); process.exit(0); });
+process.on("SIGTERM", () => { onExit(); process.exit(0); });
+
 // --- Build instructions for Claude Code ---
 
 function buildInstructions(): string {

src/storage/sessions.ts

@@ -2,21 +2,55 @@
  * Session Manager - tracks MCP server sessions.
  *
  * Location: .axme-code/sessions/<uuid>/meta.json
+ * Active session pointer: .axme-code/active-session (contains UUID)
+ *
+ * Hooks use readActiveSession() to find the current session ID
+ * instead of relying on Claude Code's session_id.
  */
 
 import { join } from "node:path";
 import { readdirSync, readFileSync } from "node:fs";
 import { randomUUID } from "node:crypto";
-import { ensureDir, writeJson, readJson, pathExists } from "./engine.js";
+import { ensureDir, writeJson, readJson, pathExists, atomicWrite, removeFile, readSafe } from "./engine.js";
 import type { SessionMeta } from "../types.js";
 import { AXME_CODE_DIR } from "../types.js";
 
 const SESSIONS_DIR = "sessions";
+const ACTIVE_SESSION_FILE = "active-session";
 
 function sessionsRoot(projectPath: string): string {
   return join(projectPath, AXME_CODE_DIR, SESSIONS_DIR);
 }
 
+function activeSessionPath(projectPath: string): string {
+  return join(projectPath, AXME_CODE_DIR, ACTIVE_SESSION_FILE);
+}
+
+/**
+ * Write the active session ID to .axme-code/active-session.
+ * Hooks read this file to determine which session to write to.
+ */
+export function writeActiveSession(projectPath: string, sessionId: string): void {
+  ensureDir(join(projectPath, AXME_CODE_DIR));
+  atomicWrite(activeSessionPath(projectPath), sessionId);
+}
+
+/**
+ * Read the active session ID from .axme-code/active-session.
+ * Returns null if no active session.
+ */
+export function readActiveSession(projectPath: string): string | null {
+  const content = readSafe(activeSessionPath(projectPath)).trim();
+  return content || null;
+}
+
+/**
+ * Remove the active-session pointer (called on process exit).
+ */
+export function clearActiveSession(projectPath: string): void {
+  removeFile(activeSessionPath(projectPath));
+}
+
 export function initSessionStore(projectPath: string): void {
   ensureDir(sessionsRoot(projectPath));
 }