hermes-control-panel/lib/audit-store.cjs

"use strict"

const { withTransaction } = require("./db.cjs")

// Number of days to retain message logs (default 90)
const LOG_RETENTION_DAYS = parseInt(process.env.HERMES_LOG_RETENTION_DAYS || "90", 10)

// ─── Helpers ──────────────────────────────────────────────────────────────────

function tryParseJson(text) {
  if (!text) return null
  try {
    return JSON.parse(text)
  } catch (_e) {
    return null
  }
}

// ─── authorizeUsage ───────────────────────────────────────────────────────────

/**
 * Check rate limits and monthly token limits for an API user.
 * Uses FOR UPDATE locking on api_users to prevent race conditions.
 *
 * @param {import('pg').Pool} pool
 * @param {{ id: string, requests_per_minute: number, monthly_token_limit: number }} apiUser
 * @param {Date} now
 * @returns {Promise<{ ok: true } | { ok: false, reason: 'rate_limited', retryAfter: number } | { ok: false, reason: 'token_limit_exceeded' }>}
 */
async function authorizeUsage(pool, apiUser, now = new Date()) {
  return withTransaction(pool, async (client) => {
    // Acquire row-level lock on this api_user to prevent concurrent races
    const { rows: lockRows } = await client.query(
      "SELECT id FROM api_users WHERE id = $1 FOR UPDATE",
      [apiUser.id]
    )

    if (!lockRows.length) {
      const err = new Error("API user not found")
      err.status = 404
      throw err
    }

    // 1. Check requests per minute
    const windowStart = new Date(now.getTime() - 60 * 1000)
    const { rows: rateRows } = await client.query(
      `SELECT count(*) AS cnt FROM usage_events
       WHERE api_user_id = $1 AND request_started_at >= $2`,
      [apiUser.id, windowStart]
    )
    const recentCount = parseInt(rateRows[0].cnt, 10)

    if (recentCount >= apiUser.requests_per_minute) {
      // retryAfter: seconds until the oldest request falls out of the 60s window
      // Conservatively return 60 seconds
      return { ok: false, reason: "rate_limited", retryAfter: 60 }
    }

    // 2. Check monthly token limit
    const { rows: tokenRows } = await client.query(
      `SELECT coalesce(sum(total_tokens), 0) AS tokens FROM usage_events
       WHERE api_user_id = $1 AND request_started_at >= date_trunc('month', $2::timestamptz)`,
      [apiUser.id, now]
    )
    const monthlyTokens = parseInt(tokenRows[0].tokens, 10)

    if (monthlyTokens >= apiUser.monthly_token_limit) {
      return { ok: false, reason: "token_limit_exceeded" }
    }

    return { ok: true }
  })
}

// ─── beginUsageEvent ─────────────────────────────────────────────────────────

/**
 * Insert a new incomplete usage_events row.
 *
 * @param {import('pg').Pool} pool
 * @param {{
 *   id: string,
 *   apiUserId: string,
 *   apiUserName: string,
 *   apiKeyId: string,
 *   route: 'pre'|'post',
 *   requestStartedAt: Date,
 *   model: string|null,
 *   promptTokens: number,
 *   completionTokens: number,
 *   totalTokens: number,
 * }} input
 * @returns {Promise<object>}
 */
async function beginUsageEvent(pool, input) {
  const { rows } = await pool.query(
    `INSERT INTO usage_events
       (id, api_user_id, api_user_name, api_key_id, route, request_started_at,
        model, prompt_tokens, completion_tokens, total_tokens, audit_complete)
     VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, false)
     RETURNING *`,
    [
      input.id,
      input.apiUserId,
      input.apiUserName,
      input.apiKeyId,
      input.route,
      input.requestStartedAt,
      input.model ?? null,
      input.promptTokens,
      input.completionTokens,
      input.totalTokens,
    ]
  )
  return rows[0]
}

// ─── completeUsageEvent ───────────────────────────────────────────────────────

/**
 * Update the usage event with completion data and insert a message_logs row.
 *
 * @param {import('pg').Pool} pool
 * @param {string} id
 * @param {{
 *   requestCompletedAt: Date,
 *   httpStatus: number,
 *   model: string|null,
 *   promptTokens: number,
 *   completionTokens: number,
 *   totalTokens: number,
 *   latencyMs: number,
 *   requestJson: object|null,
 *   requestText: string|null,
 *   responseJson: object|null,
 *   responseText: string|null,
 *   responseContentType: string|null,
 *   streaming: boolean,
 *   partial: boolean,
 * }} input
 */
async function completeUsageEvent(pool, id, input) {
  // UPDATE usage_events
  await pool.query(
    `UPDATE usage_events SET
       request_completed_at = $1,
       http_status = $2,
       model = $3,
       prompt_tokens = $4,
       completion_tokens = $5,
       total_tokens = $6,
       latency_ms = $7,
       audit_complete = true
     WHERE id = $8`,
    [
      input.requestCompletedAt,
      input.httpStatus,
      input.model ?? null,
      input.promptTokens,
      input.completionTokens,
      input.totalTokens,
      input.latencyMs,
      id,
    ]
  )

  // INSERT into message_logs with delete_after
  await pool.query(
    `INSERT INTO message_logs
       (usage_event_id, request_json, request_text, response_json, response_text,
        response_content_type, streaming, partial, delete_after)
     VALUES ($1, $2, $3, $4, $5, $6, $7, $8, now() + ($9 || ' days')::interval)
     ON CONFLICT (usage_event_id) DO UPDATE SET
       request_json = EXCLUDED.request_json,
       request_text = EXCLUDED.request_text,
       response_json = EXCLUDED.response_json,
       response_text = EXCLUDED.response_text,
       response_content_type = EXCLUDED.response_content_type,
       streaming = EXCLUDED.streaming,
       partial = EXCLUDED.partial,
       delete_after = EXCLUDED.delete_after`,
    [
      id,
      input.requestJson ?? null,
      input.requestText ?? null,
      input.responseJson ?? null,
      input.responseText ?? null,
      input.responseContentType ?? null,
      input.streaming,
      input.partial,
      LOG_RETENTION_DAYS,
    ]
  )
}

// ─── failUsageEvent ───────────────────────────────────────────────────────────

/**
 * Mark a usage event as failed. Does NOT insert a message_logs row.
 *
 * @param {import('pg').Pool} pool
 * @param {string} id
 * @param {{
 *   requestCompletedAt: Date,
 *   httpStatus: number|null,
 *   latencyMs: number,
 *   errorCode: string,
 * }} input
 */
async function failUsageEvent(pool, id, input) {
  await pool.query(
    `UPDATE usage_events SET
       request_completed_at = $1,
       http_status = $2,
       latency_ms = $3,
       error_code = $4,
       audit_complete = true
     WHERE id = $5`,
    [
      input.requestCompletedAt,
      input.httpStatus ?? null,
      input.latencyMs,
      input.errorCode,
      id,
    ]
  )
}

// ─── streamJsonlLogs ─────────────────────────────────────────────────────────

/**
 * Stream usage events + message logs as JSONL to a writable stream.
 *
 * @param {import('pg').Pool} pool
 * @param {{ apiUserId: string|null, start: Date|null, end: Date|null }} filters
 * @param {import('stream').Writable} writable
 */
async function streamJsonlLogs(pool, filters, writable) {
  const conditions = ["ue.audit_complete = true"]
  const params = []
  let idx = 1

  if (filters.apiUserId != null) {
    conditions.push(`ue.api_user_id = $${idx++}`)
    params.push(filters.apiUserId)
  }
  if (filters.start != null) {
    conditions.push(`ue.request_started_at >= $${idx++}`)
    params.push(filters.start)
  }
  if (filters.end != null) {
    conditions.push(`ue.request_started_at < $${idx++}`)
    params.push(filters.end)
  }

  const whereClause = conditions.length > 0 ? `WHERE ${conditions.join(" AND ")}` : ""

  const sql = `
    SELECT
      ue.id,
      ue.api_user_id,
      ue.api_user_name,
      ue.api_key_id,
      ue.route,
      ue.request_started_at,
      ue.request_completed_at,
      ue.http_status,
      ue.model,
      ue.prompt_tokens,
      ue.completion_tokens,
      ue.total_tokens,
      ue.latency_ms,
      ue.error_code,
      ml.request_json,
      ml.request_text,
      ml.response_json,
      ml.response_text,
      ml.streaming,
      ml.partial
    FROM usage_events ue
    LEFT JOIN message_logs ml ON ml.usage_event_id = ue.id
    ${whereClause}
    ORDER BY ue.request_started_at ASC
  `

  const result = await pool.query(sql, params)

  for (const row of result.rows) {
    const entry = {
      id: row.id,
      api_user_id: row.api_user_id,
      api_user_name: row.api_user_name,
      api_key_id: row.api_key_id,
      route: row.route,
      request_started_at: row.request_started_at,
      request_completed_at: row.request_completed_at,
      http_status: row.http_status,
      model: row.model,
      prompt_tokens: row.prompt_tokens,
      completion_tokens: row.completion_tokens,
      total_tokens: row.total_tokens,
      latency_ms: row.latency_ms,
      error_code: row.error_code,
      request: row.request_json ?? (row.request_text ? tryParseJson(row.request_text) : null),
      response: row.response_json ?? (row.response_text ? tryParseJson(row.response_text) : null),
      streaming: row.streaming,
      partial: row.partial,
    }
    writable.write(JSON.stringify(entry) + "\n")
  }

  writable.end()
}

// ─── cleanupExpiredMessageLogs ────────────────────────────────────────────────

/**
 * Delete message_logs rows whose delete_after has passed.
 *
 * @param {import('pg').Pool} pool
 * @param {Date} now
 * @returns {Promise<number>} count of deleted rows
 */
async function cleanupExpiredMessageLogs(pool, now = new Date()) {
  const { rowCount } = await pool.query(
    "DELETE FROM message_logs WHERE delete_after <= $1",
    [now]
  )
  return rowCount
}

// ─── Exports ─────────────────────────────────────────────────────────────────

module.exports = {
  authorizeUsage,
  beginUsageEvent,
  completeUsageEvent,
  failUsageEvent,
  streamJsonlLogs,
  cleanupExpiredMessageLogs,
}