agentmon/hooks/codex/handler.ts

#!/usr/bin/env node
import { randomUUID } from 'node:crypto';
import { hostname } from 'node:os';

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 BATCH_SIZE = 10;
const FLUSH_MS = 2000;
const FETCH_TIMEOUT_MS = 500;

interface Dict { [key: string]: any }

let buffer: Dict[] = [];
let flushTimer: ReturnType<typeof setTimeout> | null = null;
let isFlushing = false;

const activeRuns = new Map<string, 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.id,
    input.session,
    input.sessionId,
    input.session_id,
    input.threadId,
    input.thread_id,
    input.chatId,
    input.chat_id,
    input.conversationId,
    input.conversation_id,
  );
}

function getUsage(input: Dict): Dict | undefined {
  const usage = isRecord(input.usage) ? input.usage :
                isRecord(input.llm) ? input.llm :
                isRecord(input.tokens) ? input.tokens :
                isRecord(input.llm_usage) ? input.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 {
  return pickString(
    input.model,
    isRecord(input.llm) ? input.llm.model : undefined,
    isRecord(input.usage) ? input.usage.model : 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();
  }
}

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('metric.snapshot', sessionKey, {
    runId,
    payload: { metrics },
  }));
}

function startRun(sessionKey: string | undefined, input: Dict): string {
  const runId = randomUUID();
  if (sessionKey) {
    activeRuns.set(sessionKey, runId);
  }

  enqueue(buildEnvelope('run.start', sessionKey, {
    runId,
    attributes: {
      trigger: pickString(input.trigger_type, input.trigger, input.event, input.event_type),
    },
    payload: {
      prompt_preview: truncate(pickString(input.prompt, input.message, input.text), 200),
    },
  }));

  return runId;
}

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

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

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();
      }
    }
  }
}

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

  enqueue(buildEnvelope('session.start', sessionKey));
  startRun(sessionKey, input);

  await flush();
}

async function handleSessionEnd(input: Dict) {
  const sessionKey = getSessionKey(input);
  const runId = sessionKey ? activeRuns.get(sessionKey) : undefined;
  const usage = getUsage(input);
  const duration = pickNumber(input.duration_ms, input.elapsed_ms, input.duration);

  endRun(sessionKey, runId, input, duration);

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

  enqueueMetricSnapshot(sessionKey, runId, usage, input);

  activeRuns.delete(sessionKey || '');
  await flush();
}

async function handlePromptSubmit(input: Dict) {
  const sessionKey = getSessionKey(input);
  const runId = sessionKey ? activeRuns.get(sessionKey) : undefined;
  const duration = pickNumber(input.elapsed_ms, input.duration_ms, input.duration);
  const prompt = pickString(input.prompt, input.text, input.message);

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

  startRun(sessionKey, input);
  await flush();
}

async function handleNotification(input: Dict) {
  const sessionKey = getSessionKey(input);
  const notificationType = pickString(input.type, input.notification_type, input.event, input.event_type);
  const usage = getUsage(input);
  const duration = pickNumber(input.duration_ms, input.elapsed_ms);

  if (notificationType === 'agent-turn-complete' || notificationType === 'Done' || notificationType === 'turn.complete') {
    const runId = sessionKey ? activeRuns.get(sessionKey) : undefined;
    endRun(sessionKey, runId, input, duration);

    if (pickString(input.prompt, input.message, input.text)) {
      startRun(sessionKey, input);
    }
  } else if (usage) {
    enqueueMetricSnapshot(sessionKey, sessionKey ? activeRuns.get(sessionKey) : undefined, 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);
  }
};

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();