agentmon/hooks/gemini/handler.ts

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

import {
  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 || 'gemini';
const HOST = process.env.AGENTMON_HOST || hostname();

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

const activeRuns = new Map<string, string>();
const activeSpans = new Map<string, string>();

function getSessionKey(input: Dict): string | undefined {
  return pickString(
    input.sessionId,
    input.session_id,
    input.conversationId,
    input.conversation_id,
    input.sessionKey,
  );
}

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

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

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

async function handleSessionStart(input: Dict) {
  const sessionKey = getSessionKey(input) || randomUUID();
  const runId = randomUUID();
  activeRuns.set(sessionKey, runId);

  enqueue(buildEnvelope(FRAMEWORK, HOST, 'session.start', sessionKey));

  enqueue(buildEnvelope(FRAMEWORK, HOST, 'run.start', sessionKey, {
    runId,
    payload: {
      prompt_preview: truncate(input.prompt, 200),
    },
  }));

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

  if (runId) {
    enqueue(buildEnvelope(FRAMEWORK, HOST, 'run.end', sessionKey, {
      runId,
      payload: {
        status: 'success',
        duration_ms: duration,
        ...(usage && { usage }),
      },
    }));
  }

  enqueue(buildEnvelope(FRAMEWORK, HOST, 'session.end', sessionKey, {
    payload: usage ? { usage } : undefined,
  }));

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

async function handleToolCall(input: Dict) {
  const sessionKey = getSessionKey(input);
  const runId = sessionKey ? activeRuns.get(sessionKey) : undefined;
  const toolName = pickString(input.tool, input.tool_name, input.name) || 'unknown';
  const toolInput = isRecord(input.input) ? input.input : {};
  const spanId = randomUUID();

  if (sessionKey) {
    activeSpans.set(sessionKey + ':' + toolName, spanId);
  }

  enqueue(buildEnvelope(FRAMEWORK, HOST, 'span.start', sessionKey, {
    runId,
    spanId,
    attributes: {
      span_kind: 'tool',
      name: toolName,
    },
    payload: {
      input: truncate(toolInput, 200),
    },
  }));

  await flush();
}

async function handleToolResult(input: Dict) {
  const sessionKey = getSessionKey(input);
  const runId = sessionKey ? activeRuns.get(sessionKey) : undefined;
  const toolName = pickString(input.tool, input.tool_name, input.name) || 'unknown';
  const spanKey = sessionKey ? sessionKey + ':' + toolName : undefined;
  const spanId = spanKey ? activeSpans.get(spanKey) : undefined;
  const result = isRecord(input.result) ? input.result : isRecord(input.output) ? input.output : {};
  const success = !result.error && input.success !== false;
  const duration = pickNumber(input.duration_ms, input.elapsed_ms);
  const usage = getUsage(input);

  enqueue(buildEnvelope(FRAMEWORK, HOST, '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,
      ...(usage && { usage }),
    },
  }));

  activeSpans.delete(spanKey || '');
  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 'tool-call':
        await handleToolCall(input);
        break;
      case 'tool-result':
        await handleToolResult(input);
        break;
      default:
        console.debug(`[agentmon] unknown hook type: ${hookType}`);
    }
  } catch (err) {
    console.debug('[agentmon] handler error:', err);
  }
};

handler();