agentmon/hooks/codex/handler.ts

#!/usr/bin/env node
import { randomUUID } from 'node:crypto';
import { mkdirSync, readFileSync, unlinkSync, writeFileSync } from 'node:fs';
import { homedir, hostname } from 'node:os';
import { join } from 'node:path';
import {
  type Dict,
  isRecord,
  pickString,
  pickNumber,
  truncate,
  buildEnvelope,
  createTransport,
  readStdin,
} from '../shared/lib';

const INGEST_URL = process.env.AGENTMON_INGEST_URL || 'http://localhost:8080';
const FRAMEWORK = process.env.AGENTMON_FRAMEWORK || 'codex';
const HOST = process.env.AGENTMON_HOST || hostname();

const { enqueue, flush } = createTransport(INGEST_URL);

interface SessionState {
  sessionStarted?: boolean;
  runId?: string;
}

const STATE_DIR = join(homedir(), '.agentmon-state', 'codex');

function ensureStateDir() {
  try {
    mkdirSync(STATE_DIR, { recursive: true });
  } catch {
    // Ignore state directory creation failures so hooks stay best-effort.
  }
}

function loadState(sessionKey: string): SessionState {
  try {
    const raw = readFileSync(join(STATE_DIR, sessionKey + '.json'), 'utf8');
    return JSON.parse(raw) as SessionState;
  } catch {
    return {};
  }
}

function saveState(sessionKey: string, state: SessionState) {
  ensureStateDir();
  try {
    writeFileSync(join(STATE_DIR, sessionKey + '.json'), JSON.stringify(state), 'utf8');
  } catch {
    // Ignore state write failures so event emission can continue.
  }
}

function clearState(sessionKey: string) {
  try {
    unlinkSync(join(STATE_DIR, sessionKey + '.json'));
  } catch {
    // Ignore state cleanup failures.
  }
}

function getContext(input: Dict): Dict {
  return isRecord(input.context) ? input.context : {};
}

function getSessionRecord(input: Dict): Dict {
  return isRecord(input.session) ? input.session : {};
}

function getConversationRecord(input: Dict): Dict {
  return isRecord(input.conversation) ? input.conversation : {};
}

function getSessionKey(input: Dict): string | undefined {
  const context = getContext(input);
  const session = getSessionRecord(input);
  const conversation = getConversationRecord(input);

  return pickString(
    input.id,
    input.session,
    input.sessionId,
    input.session_id,
    input.sessionID,
    input.threadId,
    input.thread_id,
    input.threadID,
    input.chatId,
    input.chat_id,
    input.conversationId,
    input.conversation_id,
    context.id,
    context.session,
    context.sessionKey,
    context.sessionId,
    context.session_id,
    context.threadId,
    context.thread_id,
    context.conversationId,
    context.conversation_id,
    session.id,
    session.key,
    session.sessionId,
    session.session_id,
    conversation.id,
    conversation.sessionId,
    conversation.session_id,
    conversation.conversationId,
    conversation.conversation_id,
  );
}

function getUsage(input: Dict): Dict | undefined {
  const context = getContext(input);
  const usage = isRecord(input.usage) ? input.usage :
                isRecord(input.llm) ? input.llm :
                isRecord(input.tokens) ? input.tokens :
                isRecord(input.llm_usage) ? input.llm_usage :
                isRecord(context.usage) ? context.usage :
                isRecord(context.llm) ? context.llm :
                isRecord(context.tokens) ? context.tokens :
                isRecord(context.llm_usage) ? context.llm_usage : undefined;
  if (!usage) return undefined;

  const result: Dict = {};
  if (usage.input_tokens !== undefined) result.input_tokens = usage.input_tokens;
  if (usage.prompt_tokens !== undefined) result.input_tokens = usage.prompt_tokens;
  if (usage.input !== undefined) result.input_tokens = usage.input;
  if (usage.output_tokens !== undefined) result.output_tokens = usage.output_tokens;
  if (usage.completion_tokens !== undefined) result.output_tokens = usage.completion_tokens;
  if (usage.output !== undefined) result.output_tokens = usage.output;
  if (usage.total_tokens !== undefined) result.total_tokens = usage.total_tokens;
  if (usage.total !== undefined) result.total_tokens = usage.total;
  if (usage.cost !== undefined) result.total_cost = usage.cost;
  if (usage.total_cost !== undefined) result.total_cost = usage.total_cost;

  return Object.keys(result).length > 0 ? result : undefined;
}

function getModel(input: Dict): string | undefined {
  const context = getContext(input);
  return pickString(
    input.model,
    isRecord(input.llm) ? input.llm.model : undefined,
    isRecord(input.usage) ? input.usage.model : undefined,
    context.model,
    isRecord(context.llm) ? context.llm.model : undefined,
    isRecord(context.usage) ? context.usage.model : undefined,
  );
}

function getPrompt(input: Dict): string | undefined {
  const context = getContext(input);
  return pickString(
    input.prompt,
    input.text,
    input.message,
    context.prompt,
    context.text,
    context.message,
  );
}

function getDuration(input: Dict): number | undefined {
  const context = getContext(input);
  return pickNumber(
    input.duration_ms,
    input.elapsed_ms,
    input.duration,
    context.duration_ms,
    context.elapsed_ms,
    context.duration,
  );
}

function getNotificationType(input: Dict): string | undefined {
  const context = getContext(input);
  return pickString(
    input.type,
    input.notification_type,
    input.event,
    input.event_type,
    context.type,
    context.notification_type,
    context.event,
    context.event_type,
  );
}

function saveSessionState(sessionKey: string | undefined, state: SessionState) {
  if (!sessionKey) {
    return;
  }
  saveState(sessionKey, state);
}

function ensureSessionStarted(sessionKey: string | undefined, input: Dict, state: SessionState): SessionState {
  if (!sessionKey || state.sessionStarted) {
    return state;
  }

  enqueue(buildEnvelope(FRAMEWORK, HOST, 'session.start', sessionKey, {
    attributes: {
      synthetic: true,
      recovered_from: 'codex-hook',
    },
  }));

  state.sessionStarted = true;
  saveSessionState(sessionKey, state);
  return state;
}

function enqueueMetricSnapshot(sessionKey: string | undefined, runId: string | undefined, usage: Dict | undefined, input: Dict) {
  if (!usage) {
    return;
  }

  const metrics: Dict = { usage };
  const model = getModel(input);
  if (model) {
    metrics.model = model;
  }

  enqueue(buildEnvelope(FRAMEWORK, HOST, 'metric.snapshot', sessionKey, {
    runId,
    payload: { metrics },
  }));
}

function startRun(sessionKey: string | undefined, input: Dict, state: SessionState, synthetic = false): string {
  ensureSessionStarted(sessionKey, input, state);

  const runId = randomUUID();
  if (sessionKey) {
    state.runId = runId;
    saveSessionState(sessionKey, state);
  }

  enqueue(buildEnvelope(FRAMEWORK, HOST, 'run.start', sessionKey, {
    runId,
    attributes: {
      trigger: pickString(input.trigger_type, input.trigger, input.event, input.event_type),
      ...(synthetic && { synthetic: true }),
    },
    payload: {
      prompt_preview: truncate(getPrompt(input), 200),
    },
  }));

  return runId;
}

function ensureRunStarted(sessionKey: string | undefined, input: Dict, state: SessionState): string | undefined {
  if (!sessionKey) {
    return undefined;
  }
  if (state.runId) {
    return state.runId;
  }
  return startRun(sessionKey, input, state, true);
}

function endRun(sessionKey: string | undefined, state: SessionState, input: Dict, duration?: number) {
  const runId = state.runId;
  if (!runId) {
    return;
  }

  const usage = getUsage(input);
  enqueue(buildEnvelope(FRAMEWORK, HOST, 'run.end', sessionKey, {
    runId,
    payload: {
      status: 'success',
      duration_ms: duration,
      model: getModel(input),
      ...(usage && { usage }),
    },
  }));
  enqueueMetricSnapshot(sessionKey, runId, usage, input);

  state.runId = undefined;
  saveSessionState(sessionKey, state);
}

async function handleSessionStart(input: Dict) {
  const sessionKey = getSessionKey(input) || randomUUID();
  const state = loadState(sessionKey);

  if (!state.sessionStarted) {
    enqueue(buildEnvelope(FRAMEWORK, HOST, 'session.start', sessionKey));
    state.sessionStarted = true;
  }
  startRun(sessionKey, input, state);

  await flush();
}

async function handleSessionEnd(input: Dict) {
  const sessionKey = getSessionKey(input);
  const state = sessionKey ? loadState(sessionKey) : {};
  const usage = getUsage(input);
  const duration = getDuration(input);

  ensureSessionStarted(sessionKey, input, state);
  if (!state.runId && sessionKey) {
    ensureRunStarted(sessionKey, input, state);
  }
  endRun(sessionKey, state, input, duration);

  enqueue(buildEnvelope(FRAMEWORK, HOST, 'session.end', sessionKey, {
    payload: {
      model: getModel(input),
      ...(usage && { usage }),
    },
  }));

  if (sessionKey) {
    clearState(sessionKey);
  }
  await flush();
}

async function handlePromptSubmit(input: Dict) {
  const sessionKey = getSessionKey(input);
  const state = sessionKey ? loadState(sessionKey) : {};
  const duration = getDuration(input);
  const prompt = getPrompt(input);

  ensureSessionStarted(sessionKey, input, state);

  if (state.runId && prompt) {
    endRun(sessionKey, state, input, duration);
  }

  startRun(sessionKey, input, state, !state.sessionStarted);
  await flush();
}

async function handleNotification(input: Dict) {
  const sessionKey = getSessionKey(input);
  const state = sessionKey ? loadState(sessionKey) : {};
  const notificationType = getNotificationType(input);
  const usage = getUsage(input);
  const duration = getDuration(input);

  ensureSessionStarted(sessionKey, input, state);

  if (notificationType === 'agent-turn-complete' || notificationType === 'Done' || notificationType === 'turn.complete') {
    if (!state.runId && sessionKey) {
      ensureRunStarted(sessionKey, input, state);
    }
    endRun(sessionKey, state, input, duration);

    if (getPrompt(input)) {
      startRun(sessionKey, input, state);
    }
  } else if (usage) {
    enqueueMetricSnapshot(sessionKey, state.runId, usage, input);
  }

  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':
      case 'prompt-submit':
        await handlePromptSubmit(input);
        break;
      case 'notification':
        await handleNotification(input);
        break;
      default:
        console.debug(`[agentmon] unknown hook type: ${hookType}`);
    }
  } catch (err) {
    console.debug('[agentmon] handler error:', err);
  }
};

handler();