akm-cli/dist/llm/usage-telemetry.js

akm (Agent Knowledge Manager) — a portable, local-first capability library for AI agents. Discover, load, share, and improve reusable skills, scripts, workflows, and knowledge across any shell-capable coding agent, including Claude Code, OpenCode, and Cur

// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at https://mozilla.org/MPL/2.0/.
/**
 * Per-call LLM usage telemetry (#576).
 *
 * `chatCompletion` captures usage + model + finish_reason + wall-time for
 * EVERY OpenAI-compatible HTTP attempt and emits one terminal
 * {@link LlmUsageRecord} through a
 * module-level sink. The sink indirection keeps `client.ts` free of any
 * dependency on the events/db layer: the application wires the sink to
 * persistence at startup / per improve run, and tests can inspect records in
 * memory.
 *
 * The pipeline *stage* that made the call is ambient, not threaded through
 * call sites. A param-threading prototype was deliberately discarded in 0.8.5
 * (every call site would have to forward a `stage` argument it does not care
 * about). Instead callers wrap a well-delimited phase once with
 * {@link withLlmStage}; any `chatCompletion` invoked inside that async region —
 * however deeply nested — is attributed to that stage via `AsyncLocalStorage`.
 *
 * EVERYTHING here is best-effort. Telemetry must NEVER break a real LLM call:
 * a sink that throws, an unset stage, or a malformed usage block all degrade
 * silently. `emitLlmUsage` swallows sink errors; `currentLlmStage` returns
 * `undefined` outside any `withLlmStage` scope.
 */
import { AsyncLocalStorage } from "node:async_hooks";
import { asNonEmptyString } from "../core/common.js";
const attributionStorage = new AsyncLocalStorage();
let usageSink;
/**
 * Run `fn` with `stage` as the ambient LLM stage. Any `chatCompletion` call
 * made synchronously or asynchronously within `fn` (including through awaited
 * helpers and nested `withLlmStage` calls — the innermost wins) is attributed
 * to `stage`. Returns whatever `fn` returns; never alters control flow.
 */
export function withLlmStage(stage, fn, attribution = {}) {
    return attributionStorage.run({ ...attributionStorage.getStore(), ...attribution, stage }, fn);
}
/** The ambient LLM stage for the current async context, or `undefined` outside any {@link withLlmStage} scope. */
export function currentLlmStage() {
    return attributionStorage.getStore()?.stage;
}
/**
 * Install the process-wide usage sink. Replaces any previously installed sink.
 * The application wires this to persistence; tests install an in-memory
 * collector. Pair with {@link clearLlmUsageSink} in a `finally` block.
 */
export function setLlmUsageSink(sink) {
    usageSink = sink;
}
/** Remove the installed sink so subsequent calls emit nowhere. Idempotent. */
export function clearLlmUsageSink() {
    usageSink = undefined;
}
/**
 * Whether a usage sink is currently installed. Standalone entry points use
 * this to avoid clobbering a sink an enclosing run (e.g. `akm improve`) already
 * installed: they install their own only when none is active.
 */
export function hasLlmUsageSink() {
    return usageSink !== undefined;
}
/**
 * Emit one usage record to the installed sink, stamping the ambient stage.
 * Best-effort: no sink is a no-op, and a sink that throws is swallowed so
 * telemetry can never fail the LLM call that produced it.
 */
export function emitLlmUsage(record) {
    const sink = usageSink;
    if (!sink)
        return;
    try {
        const ambient = attributionStorage.getStore();
        sink({
            ...record,
            stage: record.stage ?? ambient?.stage,
            engine: record.engine ?? ambient?.engine,
            process: record.process ?? ambient?.process,
        });
    }
    catch {
        // Telemetry must never break a real run.
    }
}
function asFiniteNonNegative(value) {
    return typeof value === "number" && Number.isFinite(value) && value >= 0 ? value : undefined;
}
function asLlmUsageErrorCode(value) {
    switch (value) {
        case "rate_limited":
        case "provider_error":
        case "provider_html_error":
        case "network_error":
        case "parse_error":
        case "timeout":
        case "unknown_error":
            return value;
        default:
            return undefined;
    }
}
/** Decode durable event metadata with the same validation used by health aggregation. */
export function decodeLlmUsageRecord(value) {
    if (!value || typeof value !== "object")
        return undefined;
    const raw = value;
    const durationMs = asFiniteNonNegative(raw.durationMs);
    if (durationMs === undefined)
        return undefined;
    // Historical rows predate terminal provenance. Keep health aggregation
    // backward-compatible while conservatively treating their model as configured,
    // never response-observed.
    const outcome = raw.outcome === "success" || raw.outcome === "error" ? raw.outcome : "success";
    const modelSource = raw.modelSource === "response" || raw.modelSource === "configured" ? raw.modelSource : "configured";
    const record = { durationMs, outcome, modelSource };
    for (const key of ["stage", "engine", "process", "model", "finishReason"]) {
        const decoded = asNonEmptyString(raw[key]);
        if (decoded !== undefined)
            record[key] = decoded;
    }
    for (const key of ["promptTokens", "completionTokens", "totalTokens", "reasoningTokens"]) {
        const decoded = asFiniteNonNegative(raw[key]);
        if (decoded !== undefined)
            record[key] = decoded;
    }
    const errorCode = asLlmUsageErrorCode(raw.errorCode);
    if (errorCode !== undefined)
        record.errorCode = errorCode;
    return record;
}
/**
 * Project a provider `usage` block into the token fields of an
 * {@link LlmUsageRecord}. Missing or garbled values are omitted (not zeroed)
 * so a best-effort record still distinguishes "0 tokens" from "unknown".
 */
export function extractUsageTokens(usage) {
    if (!usage || typeof usage !== "object")
        return {};
    const out = {};
    const prompt = asFiniteNonNegative(usage.prompt_tokens);
    const completion = asFiniteNonNegative(usage.completion_tokens);
    const total = asFiniteNonNegative(usage.total_tokens);
    const reasoning = asFiniteNonNegative(usage.completion_tokens_details?.reasoning_tokens);
    if (prompt !== undefined)
        out.promptTokens = prompt;
    if (completion !== undefined)
        out.completionTokens = completion;
    if (total !== undefined)
        out.totalTokens = total;
    if (reasoning !== undefined)
        out.reasoningTokens = reasoning;
    return out;
}