fix: preserve session state across turns in claude-code hook handler

handleNotification("Done") was incorrectly emitting session.end and calling clearState at the end of each Claude turn. Since "Done" means a turn finished (not the session), clearing state caused subsequent tool calls to find no runId, storing spans without run_id and making them invisible in run-level queries. - handleNotification: remove session.end emission and clearState call; only emit run.end for the completed turn - handleSessionEnd: load state file to get runId (in-memory activeRuns is always empty in a subprocess) - handlePromptSubmit: load state file to get runId for ending previous run before starting a new one Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
2026-03-18 23:42:22 -07:00
parent f8ddea3698
commit 2e277fb138
2 changed files with 1215 additions and 0 deletions
@@ -0,0 +1,651 @@
+#!/usr/bin/env node
+import { randomUUID } from 'node:crypto';
+import { hostname, homedir } from 'node:os';
+import { readFileSync, writeFileSync, mkdirSync, unlinkSync } from 'node:fs';
+import { join } from 'node:path';
+
+const INGEST_URL = process.env.AGENTMON_INGEST_URL || 'http://localhost:8080';
+const FRAMEWORK = process.env.AGENTMON_FRAMEWORK || 'claude-code';
+const HOST = process.env.AGENTMON_HOST || hostname();
+const BATCH_SIZE = 10;
+const FLUSH_MS = 2000;
+const FETCH_TIMEOUT_MS = 500;
+
+interface Dict { [key: string]: any }
+
+// ── Persisted state (survives between hook subprocess invocations) ──────────
+interface SessionState {
+  runId?: string;
+  spans: { [key: string]: string };   // key = sessionKey:toolName, value = spanId
+  subagent?: { name: string; spanId: string };
+  compactSpanId?: string;
+}
+
+const STATE_DIR = join(homedir(), '.agentmon-state');
+
+function ensureStateDir() {
+  try { mkdirSync(STATE_DIR, { recursive: true }); } catch { /* ignore */ }
+}
+
+function loadState(sessionKey: string): SessionState {
+  try {
+    const raw = readFileSync(join(STATE_DIR, sessionKey + '.json'), 'utf8');
+    return JSON.parse(raw) as SessionState;
+  } catch {
+    return { spans: {} };
+  }
+}
+
+function saveState(sessionKey: string, state: SessionState) {
+  ensureStateDir();
+  try {
+    writeFileSync(join(STATE_DIR, sessionKey + '.json'), JSON.stringify(state), 'utf8');
+  } catch { /* ignore */ }
+}
+
+function clearState(sessionKey: string) {
+  try { unlinkSync(join(STATE_DIR, sessionKey + '.json')); } catch { /* ignore */ }
+}
+// ─────────────────────────────────────────────────────────────────────────────
+
+let buffer: Dict[] = [];
+let flushTimer: ReturnType<typeof setTimeout> | null = null;
+let isFlushing = false;
+
+const activeRuns = new Map<string, string>();
+const activeSpans = new Map<string, string>();
+const activeSubagents = new Map<string, { name: string; spanId: string }>();
+
+function isRecord(value: unknown): value is Dict {
+  return value !== null && typeof value === 'object' && !Array.isArray(value);
+}
+
+function pickString(...values: unknown[]): string | undefined {
+  for (const value of values) {
+    if (typeof value === 'string' && value.trim() !== '') {
+      return value;
+    }
+  }
+  return undefined;
+}
+
+function pickNumber(...values: unknown[]): number | undefined {
+  for (const value of values) {
+    if (typeof value === 'number' && Number.isFinite(value)) {
+      return value;
+    }
+  }
+  return undefined;
+}
+
+function truncate(value: unknown, limit: number): string | undefined {
+  if (value === undefined || value === null) {
+    return undefined;
+  }
+
+  const text = typeof value === 'string' ? value : safeJSONStringify(value);
+  if (!text) {
+    return undefined;
+  }
+
+  if (text.length <= limit) {
+    return text;
+  }
+  return text.slice(0, limit) + '...';
+}
+
+function safeJSONStringify(value: unknown): string {
+  try {
+    return JSON.stringify(value);
+  } catch {
+    return String(value);
+  }
+}
+
+function getSessionKey(input: Dict): string | undefined {
+  return pickString(
+    input.sessionId,
+    input.session_id,
+    input.conversationId,
+    input.conversation_id,
+  );
+}
+
+function getUsage(input: Dict): Dict | undefined {
+  const usage = isRecord(input.usage) ? input.usage : 
+                isRecord(input.llm) ? input.llm : 
+                undefined;
+  if (!usage) return undefined;
+  
+  const result: Dict = {};
+  if (usage.input_tokens !== undefined) result.input_tokens = usage.input_tokens;
+  if (usage.output_tokens !== undefined) result.output_tokens = usage.output_tokens;
+  if (usage.cache_creation_tokens !== undefined) result.cache_creation_tokens = usage.cache_creation_tokens;
+  if (usage.cache_read_tokens !== undefined) result.cache_read_tokens = usage.cache_read_tokens;
+  if (usage.thinking_tokens !== undefined) result.thinking_tokens = usage.thinking_tokens;
+  if (usage.total_tokens !== undefined) result.total_tokens = usage.total_tokens;
+  if (usage.total_cost !== undefined) result.total_cost = usage.total_cost;
+  
+  return Object.keys(result).length > 0 ? result : undefined;
+}
+
+function getContextWindow(input: Dict): Dict | undefined {
+  const stats = isRecord(input.context_stats) ? input.context_stats :
+                 isRecord(input.stats) ? input.stats : undefined;
+  if (!stats) return undefined;
+  
+  const result: Dict = {};
+  if (stats.input_tokens !== undefined) result.input_tokens = stats.input_tokens;
+  if (stats.output_tokens !== undefined) result.output_tokens = stats.output_tokens;
+  if (stats.used_tokens !== undefined) result.used_tokens = stats.used_tokens;
+  if (stats.max_tokens !== undefined) result.max_tokens = stats.max_tokens;
+  if (stats.tokens_remaining !== undefined) result.tokens_remaining = stats.tokens_remaining;
+  
+  return Object.keys(result).length > 0 ? result : undefined;
+}
+
+function buildEnvelope(
+  type: string,
+  sessionKey?: string,
+  opts: {
+    runId?: string;
+    spanId?: string;
+    parentSpanId?: string;
+    attributes?: Dict;
+    payload?: Dict;
+  } = {},
+): Dict {
+  const correlation: Dict = {};
+  if (sessionKey) {
+    correlation.session_id = sessionKey;
+  }
+  if (opts.runId) {
+    correlation.run_id = opts.runId;
+  }
+  if (opts.spanId) {
+    correlation.span_id = opts.spanId;
+  }
+  if (opts.parentSpanId) {
+    correlation.parent_span_id = opts.parentSpanId;
+  }
+
+  const envelope: Dict = {
+    schema: { name: 'agentmon.event', version: 1 },
+    event: {
+      id: randomUUID(),
+      type,
+      ts: new Date().toISOString(),
+      source: {
+        framework: FRAMEWORK,
+        client_id: HOST,
+        host: HOST,
+      },
+    },
+  };
+
+  if (Object.keys(correlation).length > 0) {
+    envelope.correlation = correlation;
+  }
+  if (opts.attributes && Object.keys(opts.attributes).length > 0) {
+    envelope.attributes = opts.attributes;
+  }
+  if (opts.payload && Object.keys(opts.payload).length > 0) {
+    envelope.payload = opts.payload;
+  }
+
+  return envelope;
+}
+
+function scheduleFlush() {
+  if (!flushTimer) {
+    flushTimer = setTimeout(() => {
+      void flush();
+    }, FLUSH_MS);
+  }
+}
+
+function enqueue(event: Dict) {
+  buffer.push(event);
+  if (buffer.length >= BATCH_SIZE) {
+    void flush();
+  } else {
+    scheduleFlush();
+  }
+}
+
+async function postBatch(batch: Dict[]) {
+  const controller = new AbortController();
+  const timeout = setTimeout(() => controller.abort(), FETCH_TIMEOUT_MS);
+
+  try {
+    await fetch(`${INGEST_URL}/v1/events`, {
+      method: 'POST',
+      headers: { 'Content-Type': 'application/json' },
+      body: JSON.stringify(batch),
+      signal: controller.signal,
+    });
+  } finally {
+    clearTimeout(timeout);
+  }
+}
+
+async function flush() {
+  if (flushTimer) {
+    clearTimeout(flushTimer);
+    flushTimer = null;
+  }
+  if (isFlushing || buffer.length === 0) {
+    return;
+  }
+
+  isFlushing = true;
+  const batch = buffer.splice(0, BATCH_SIZE);
+
+  try {
+    await postBatch(batch);
+  } catch {
+    console.debug(`[agentmon] failed to flush ${batch.length} events`);
+  } finally {
+    isFlushing = false;
+    if (buffer.length > 0) {
+      if (buffer.length >= BATCH_SIZE) {
+        void flush();
+      } else {
+        scheduleFlush();
+      }
+    }
+  }
+}
+
+function emitError(sessionKey: string | undefined, runId: string | undefined, spanId: string | undefined, errorValue: unknown) {
+  if (errorValue === undefined || errorValue === null || errorValue === false) {
+    return;
+  }
+
+  const errorRecord = isRecord(errorValue) ? errorValue : {};
+  const message = pickString(errorRecord.message, errorRecord.error, errorValue) || 'unknown';
+  const errType = pickString(errorRecord.type, errorRecord.code) || FRAMEWORK;
+
+  enqueue(buildEnvelope('error', sessionKey, {
+    runId,
+    spanId,
+    payload: {
+      error: {
+        type: errType,
+        message,
+      },
+    },
+  }));
+}
+
+async function handleSessionStart(input: Dict) {
+  const sessionKey = getSessionKey(input) || randomUUID();
+  const runId = randomUUID();
+  activeRuns.set(sessionKey, runId);
+  saveState(sessionKey, { runId, spans: {} });
+
+  const contextWindow = getContextWindow(input);
+
+  enqueue(buildEnvelope('session.start', sessionKey, {
+    attributes: contextWindow ? { context_window: contextWindow } : undefined,
+  }));
+
+  enqueue(buildEnvelope('run.start', sessionKey, {
+    runId,
+    attributes: {
+      trigger: pickString(input.trigger_type, input.trigger),
+    },
+    payload: {
+      prompt_preview: truncate(input.prompt, 200),
+    },
+  }));
+
+  await flush();
+}
+
+async function handleSessionEnd(input: Dict) {
+  const sessionKey = getSessionKey(input);
+  const state = sessionKey ? loadState(sessionKey) : { spans: {} };
+  const runId = sessionKey ? (activeRuns.get(sessionKey) || state.runId) : undefined;
+  const usage = getUsage(input);
+  const contextWindow = getContextWindow(input);
+  const duration = pickNumber(input.duration_ms, input.elapsed_ms);
+
+  if (runId) {
+    enqueue(buildEnvelope('run.end', sessionKey, {
+      runId,
+      payload: {
+        status: 'success',
+        duration_ms: duration,
+        model: pickString(input.model),
+        ...(usage && { usage }),
+        ...(contextWindow && { context_window: contextWindow }),
+      },
+    }));
+  }
+
+  enqueue(buildEnvelope('session.end', sessionKey, {
+    payload: {
+      ...(usage && { usage }),
+      ...(contextWindow && { context_window: contextWindow }),
+    },
+  }));
+  activeRuns.delete(sessionKey || '');
+  activeSpans.clear();
+  activeSubagents.clear();
+  if (sessionKey) clearState(sessionKey);
+
+  await flush();
+}
+
+async function handlePromptSubmit(input: Dict) {
+  const sessionKey = getSessionKey(input);
+  const state = sessionKey ? loadState(sessionKey) : { spans: {} };
+  const runId = sessionKey ? (activeRuns.get(sessionKey) || state.runId) : undefined;
+  const prompt = pickString(input.prompt, input.text, input.message);
+
+  if (runId && prompt) {
+    enqueue(buildEnvelope('run.end', sessionKey, {
+      runId,
+      payload: {
+        status: 'success',
+        duration_ms: pickNumber(input.elapsed_ms, input.duration_ms),
+      },
+    }));
+  }
+
+  const newRunId = randomUUID();
+  if (sessionKey) {
+    activeRuns.set(sessionKey, newRunId);
+    saveState(sessionKey, { runId: newRunId, spans: {} });
+  }
+
+  enqueue(buildEnvelope('run.start', sessionKey, {
+    runId: newRunId,
+    attributes: {
+      type: 'user_prompt',
+    },
+    payload: {
+      prompt_preview: truncate(prompt, 200),
+    },
+  }));
+
+  await flush();
+}
+
+async function handleToolStart(input: Dict) {
+  const sessionKey = getSessionKey(input);
+  const state = sessionKey ? loadState(sessionKey) : { spans: {} };
+  const runId = sessionKey ? (activeRuns.get(sessionKey) || state.runId) : undefined;
+  const toolName = pickString(input.tool, input.tool_name, input.name) || 'unknown';
+  const toolInput = isRecord(input.input) ? input.input : {};
+  const spanId = randomUUID();
+  const spanKey = sessionKey ? sessionKey + ':' + toolName : undefined;
+
+  if (spanKey) {
+    activeSpans.set(spanKey, spanId);
+    state.spans[spanKey] = spanId;
+    if (runId) state.runId = runId;
+    if (sessionKey) saveState(sessionKey, state);
+  }
+
+  enqueue(buildEnvelope('span.start', sessionKey, {
+    runId,
+    spanId,
+    attributes: {
+      span_kind: 'tool',
+      name: toolName,
+    },
+    payload: {
+      input: truncate(toolInput, 200),
+    },
+  }));
+
+  await flush();
+}
+
+async function handleToolEnd(input: Dict) {
+  const sessionKey = getSessionKey(input);
+  const state = sessionKey ? loadState(sessionKey) : { spans: {} };
+  const runId = sessionKey ? (activeRuns.get(sessionKey) || state.runId) : undefined;
+  const toolName = pickString(input.tool, input.tool_name, input.name) || 'unknown';
+  const spanKey = sessionKey ? sessionKey + ':' + toolName : undefined;
+  const spanId = spanKey ? (activeSpans.get(spanKey) || state.spans[spanKey]) : undefined;
+  const result = isRecord(input.result) ? input.result : isRecord(input.output) ? input.output : {};
+  const success = !result.error;
+  const duration = pickNumber(input.duration_ms, input.elapsed_ms);
+
+  enqueue(buildEnvelope('span.end', sessionKey, {
+    runId,
+    spanId,
+    attributes: {
+      span_kind: 'tool',
+      name: toolName,
+    },
+    payload: {
+      status: success ? 'success' : 'error',
+      result_preview: truncate(result.output ?? result.error ?? result, 500),
+      duration_ms: duration,
+    },
+  }));
+
+  if (!success) {
+    emitError(sessionKey, runId, spanId, result.error);
+  }
+
+  activeSpans.delete(spanKey || '');
+
+  await flush();
+}
+
+async function handleSubagentStart(input: Dict) {
+  const sessionKey = getSessionKey(input);
+  const state = sessionKey ? loadState(sessionKey) : { spans: {} };
+  const runId = sessionKey ? (activeRuns.get(sessionKey) || state.runId) : undefined;
+  const agentName = pickString(input.agent, input.agent_name, input.name) || 'unknown';
+  const spanId = randomUUID();
+
+  if (sessionKey) {
+    activeSubagents.set(sessionKey, { name: agentName, spanId });
+    state.subagent = { name: agentName, spanId };
+    if (runId) state.runId = runId;
+    saveState(sessionKey, state);
+  }
+
+  enqueue(buildEnvelope('span.start', sessionKey, {
+    runId,
+    spanId,
+    attributes: {
+      span_kind: 'agent',
+      name: agentName,
+      type: 'subagent',
+    },
+    payload: {
+      prompt_preview: truncate(input.prompt, 200),
+    },
+  }));
+
+  await flush();
+}
+
+async function handleSubagentStop(input: Dict) {
+  const sessionKey = getSessionKey(input);
+  const state = sessionKey ? loadState(sessionKey) : { spans: {} };
+  const runId = sessionKey ? (activeRuns.get(sessionKey) || state.runId) : undefined;
+  const subagent = sessionKey ? (activeSubagents.get(sessionKey) || state.subagent) : undefined;
+  const spanId = subagent?.spanId;
+  const agentName = subagent?.name || pickString(input.agent, input.agent_name) || 'unknown';
+  const duration = pickNumber(input.duration_ms, input.elapsed_ms);
+  const usage = getUsage(input);
+
+  enqueue(buildEnvelope('span.end', sessionKey, {
+    runId,
+    spanId,
+    attributes: {
+      span_kind: 'agent',
+      name: agentName,
+      type: 'subagent',
+    },
+    payload: {
+      status: 'success',
+      duration_ms: duration,
+      ...(usage && { usage }),
+    },
+  }));
+
+  activeSubagents.delete(sessionKey || '');
+
+  await flush();
+}
+
+async function handleCompactStart(input: Dict) {
+  const sessionKey = getSessionKey(input);
+  const state = sessionKey ? loadState(sessionKey) : { spans: {} };
+  const runId = sessionKey ? (activeRuns.get(sessionKey) || state.runId) : undefined;
+  const spanId = randomUUID();
+
+  if (sessionKey) {
+    activeSpans.set(sessionKey + ':compact', spanId);
+    state.compactSpanId = spanId;
+    if (runId) state.runId = runId;
+    saveState(sessionKey, state);
+  }
+
+  enqueue(buildEnvelope('span.start', sessionKey, {
+    runId,
+    spanId,
+    attributes: {
+      span_kind: 'internal',
+      name: 'context_compaction',
+    },
+  }));
+
+  await flush();
+}
+
+async function handleCompactEnd(input: Dict) {
+  const sessionKey = getSessionKey(input);
+  const state = sessionKey ? loadState(sessionKey) : { spans: {} };
+  const runId = sessionKey ? (activeRuns.get(sessionKey) || state.runId) : undefined;
+  const spanKey = sessionKey ? sessionKey + ':compact' : undefined;
+  const spanId = spanKey ? (activeSpans.get(spanKey) || state.compactSpanId) : undefined;
+  const duration = pickNumber(input.duration_ms, input.elapsed_ms);
+  const contextWindow = getContextWindow(input);
+
+  enqueue(buildEnvelope('span.end', sessionKey, {
+    runId,
+    spanId,
+    attributes: {
+      span_kind: 'internal',
+      name: 'context_compaction',
+    },
+    payload: {
+      status: 'success',
+      duration_ms: duration,
+      ...(contextWindow && { context_window: contextWindow }),
+    },
+  }));
+
+  activeSpans.delete(spanKey || '');
+
+  await flush();
+}
+
+async function handleNotification(input: Dict) {
+  const sessionKey = getSessionKey(input);
+  const notificationType = pickString(input.notification_type, input.type);
+  const usage = getUsage(input);
+  const contextWindow = getContextWindow(input);
+  const duration = pickNumber(input.duration_ms, input.elapsed_ms);
+
+  if (notificationType === 'Done' || notificationType === 'success') {
+    const state = sessionKey ? loadState(sessionKey) : { spans: {} };
+    const runId = sessionKey ? (activeRuns.get(sessionKey) || state.runId) : undefined;
+
+    if (runId) {
+      enqueue(buildEnvelope('run.end', sessionKey, {
+        runId,
+        payload: {
+          status: 'success',
+          duration_ms: duration,
+          model: pickString(input.model),
+          ...(usage && { usage }),
+          ...(contextWindow && { context_window: contextWindow }),
+        },
+      }));
+    }
+
+    // Do NOT emit session.end or clear state here — "Done" means Claude
+    // finished a turn, not the session. State must persist so tool calls
+    // between turns still have a runId. The actual session.end is emitted
+    // by handleSessionEnd when the Stop hook fires.
+  }
+
+  await flush();
+}
+
+const handler = async () => {
+  const args = process.argv.slice(2);
+  const hookType = args[0] || 'unknown';
+
+  let input: Dict = {};
+  try {
+    const stdin = await readStdin();
+    if (stdin) {
+      input = JSON.parse(stdin);
+    }
+  } catch {
+    input = {};
+  }
+
+  try {
+    switch (hookType) {
+      case 'start':
+        await handleSessionStart(input);
+        break;
+      case 'stop':
+        await handleSessionEnd(input);
+        break;
+      case 'prompt':
+        await handlePromptSubmit(input);
+        break;
+      case 'tool-start':
+        await handleToolStart(input);
+        break;
+      case 'tool-end':
+        await handleToolEnd(input);
+        break;
+      case 'subagent-start':
+        await handleSubagentStart(input);
+        break;
+      case 'subagent-stop':
+        await handleSubagentStop(input);
+        break;
+      case 'compact-start':
+        await handleCompactStart(input);
+        break;
+      case 'compact-end':
+        await handleCompactEnd(input);
+        break;
+      case 'notification':
+        await handleNotification(input);
+        break;
+      default:
+        console.debug(`[agentmon] unknown hook type: ${hookType}`);
+    }
+  } catch (err) {
+    console.debug('[agentmon] handler error:', err);
+  }
+};
+
+async function readStdin(): Promise<string> {
+  return new Promise((resolve) => {
+    let data = '';
+    process.stdin.on('data', (chunk) => data += chunk);
+    process.stdin.on('end', () => resolve(data));
+    process.stdin.on('error', () => resolve(''));
+    setTimeout(() => resolve(data), 100);
+  });
+}
+
+handler();