openclaw/src/agents/usage.ts

/**
 * Token usage normalization helpers.
 * Converts provider-specific usage shapes into OpenClaw's normalized input,
 * output, cache, reasoning, and total token accounting fields.
 */
import { asFiniteNumber } from "@openclaw/normalization-core/number-coercion";
import type { Usage } from "../llm/types.js";

export type ContextUsage = NonNullable<Usage["contextUsage"]>;

/** Provider/SDK usage payload variants accepted by usage normalization. */
export type UsageLike = {
  input?: number;
  output?: number;
  cacheRead?: number;
  cacheWrite?: number;
  contextUsage?: ContextUsage;
  total?: number;
  // Common alternates across providers/SDKs.
  inputTokens?: number;
  outputTokens?: number;
  promptTokens?: number;
  completionTokens?: number;
  input_tokens?: number;
  output_tokens?: number;
  prompt_tokens?: number;
  completion_tokens?: number;
  cache_read_input_tokens?: number;
  cache_creation_input_tokens?: number;
  reasoningTokens?: number;
  reasoning_tokens?: number;
  completion_tokens_details?: { reasoning_tokens?: number };
  output_tokens_details?: { reasoning_tokens?: number };
  // Moonshot/Kimi uses cached_tokens for cache read count (explicit caching API).
  cached_tokens?: number;
  // OpenAI Responses reports cached prompt reuse here.
  input_tokens_details?: { cached_tokens?: number };
  // Kimi K2 uses prompt_tokens_details.cached_tokens for automatic prefix caching.
  prompt_tokens_details?: { cached_tokens?: number };
  // Some agents/logs emit alternate naming.
  totalTokens?: number;
  total_tokens?: number;
  cache_read?: number;
  cache_write?: number;
  // llama.cpp-style streamed completion metadata.
  prompt_n?: number;
  predicted_n?: number;
  timings?: {
    prompt_n?: number;
    predicted_n?: number;
  };
  // Optional cost metadata carried through transcripts for downstream cost accounting.
  cost?: Partial<Usage["cost"]>;
};

/** Normalized token counts used by runtime accounting. */
export type NormalizedUsage = {
  input?: number;
  output?: number;
  cacheRead?: number;
  cacheWrite?: number;
  contextUsage?: ContextUsage;
  reasoningTokens?: number;
  total?: number;
};

/** OpenAI chat-completions compatible usage shape. */
export type OpenAiChatCompletionsUsage = {
  prompt_tokens: number;
  completion_tokens: number;
  total_tokens: number;
  prompt_tokens_details?: { cached_tokens: number };
  completion_tokens_details?: { reasoning_tokens: number };
};

/** Assistant usage snapshot with token counts and computed cost buckets. */
export type AssistantUsageSnapshot = Usage;

/** Build a zeroed assistant usage snapshot. */
export function makeZeroUsageSnapshot(): AssistantUsageSnapshot {
  return {
    input: 0,
    output: 0,
    cacheRead: 0,
    cacheWrite: 0,
    totalTokens: 0,
    cost: {
      input: 0,
      output: 0,
      cacheRead: 0,
      cacheWrite: 0,
      total: 0,
    },
  };
}

/** Return true when any normalized usage bucket is positive. */
export function hasNonzeroUsage(usage?: NormalizedUsage | null): usage is NormalizedUsage {
  if (!usage) {
    return false;
  }
  return (
    [
      usage.input,
      usage.output,
      usage.cacheRead,
      usage.cacheWrite,
      usage.contextUsage?.state === "available" ? usage.contextUsage.promptTokens : undefined,
      usage.contextUsage?.state === "available" ? usage.contextUsage.totalTokens : undefined,
      usage.reasoningTokens,
      usage.total,
    ].some((v) => typeof v === "number" && Number.isFinite(v) && v > 0) ||
    usage.contextUsage?.state === "unavailable"
  );
}

const normalizeTokenCount = (value: unknown): number | undefined => {
  const numeric = asFiniteNumber(value);
  if (numeric === undefined) {
    return undefined;
  }
  if (numeric <= 0) {
    return 0;
  }
  return Math.min(Math.trunc(numeric), Number.MAX_SAFE_INTEGER);
};

/** Normalize provider-specific token usage fields into OpenClaw usage buckets. */
export function normalizeUsage(raw?: UsageLike | null): NormalizedUsage | undefined {
  if (!raw) {
    return undefined;
  }

  const cacheRead = normalizeTokenCount(
    raw.cacheRead ??
      raw.cache_read ??
      raw.cache_read_input_tokens ??
      raw.cached_tokens ??
      raw.input_tokens_details?.cached_tokens ??
      raw.prompt_tokens_details?.cached_tokens,
  );

  const rawInputValue =
    raw.input ??
    raw.inputTokens ??
    raw.input_tokens ??
    raw.promptTokens ??
    raw.prompt_tokens ??
    raw.prompt_n ??
    raw.timings?.prompt_n;

  const usesOpenAIStylePromptTotals =
    raw.cached_tokens !== undefined ||
    raw.input_tokens_details?.cached_tokens !== undefined ||
    raw.prompt_tokens_details?.cached_tokens !== undefined;

  // Some providers (shared model runtime OpenAI-format) pre-subtract cached_tokens from
  // prompt/input totals upstream, while OpenAI-style prompt/input aliases
  // include cached tokens in the reported prompt total. Normalize both cases
  // to uncached input tokens so downstream prompt-token math does not double-
  // count cache reads.
  const rawInput = asFiniteNumber(rawInputValue);
  const normalizedInput =
    rawInput !== undefined && usesOpenAIStylePromptTotals && cacheRead !== undefined
      ? rawInput - cacheRead
      : rawInput;
  const input = normalizeTokenCount(normalizedInput);
  const output = normalizeTokenCount(
    raw.output ??
      raw.outputTokens ??
      raw.output_tokens ??
      raw.completionTokens ??
      raw.completion_tokens ??
      raw.predicted_n ??
      raw.timings?.predicted_n,
  );
  const cacheWrite = normalizeTokenCount(
    raw.cacheWrite ?? raw.cache_write ?? raw.cache_creation_input_tokens,
  );
  const contextPromptTokens =
    raw.contextUsage?.state === "available"
      ? normalizeTokenCount(raw.contextUsage.promptTokens)
      : undefined;
  const contextTotalTokens =
    raw.contextUsage?.state === "available"
      ? normalizeTokenCount(raw.contextUsage.totalTokens)
      : undefined;
  const contextUsage =
    raw.contextUsage?.state === "unavailable"
      ? ({ state: "unavailable" } as const)
      : contextPromptTokens !== undefined &&
          contextTotalTokens !== undefined &&
          contextTotalTokens >= contextPromptTokens
        ? ({
            state: "available",
            promptTokens: contextPromptTokens,
            totalTokens: contextTotalTokens,
          } as const)
        : undefined;
  const reasoningTokens = normalizeTokenCount(
    raw.reasoningTokens ??
      raw.reasoning_tokens ??
      raw.completion_tokens_details?.reasoning_tokens ??
      raw.output_tokens_details?.reasoning_tokens,
  );
  const total = normalizeTokenCount(raw.total ?? raw.totalTokens ?? raw.total_tokens);

  if (
    input === undefined &&
    output === undefined &&
    cacheRead === undefined &&
    cacheWrite === undefined &&
    contextUsage === undefined &&
    reasoningTokens === undefined &&
    total === undefined
  ) {
    return undefined;
  }

  return {
    input,
    output,
    cacheRead,
    cacheWrite,
    ...(contextUsage ? { contextUsage } : {}),
    ...(reasoningTokens !== undefined ? { reasoningTokens } : {}),
    total,
  };
}

/**
 * Maps normalized usage to OpenAI Chat Completions `usage` fields.
 *
 * `prompt_tokens` is input + cacheRead (cache write is excluded to match the
 * OpenAI-style breakdown used by the compat endpoint).
 *
 * `total_tokens` is the greater of the component sum and aggregate `total` when
 * present, so a partial breakdown cannot discard a valid upstream total.
 *
 * `prompt_tokens_details.cached_tokens` is emitted when `cacheRead > 0` so
 * downstream chat-completions clients can compute the cache-aware blended
 * cost. Field name and shape match OpenAI's documented usage breakdown:
 * https://platform.openai.com/docs/guides/prompt-caching
 */
export function toOpenAiChatCompletionsUsage(
  usage: NormalizedUsage | undefined,
): OpenAiChatCompletionsUsage {
  const input = usage?.input ?? 0;
  const output = usage?.output ?? 0;
  const cacheRead = usage?.cacheRead ?? 0;
  const promptTokens = Math.max(0, input + cacheRead);
  const completionTokens = Math.max(0, output);
  const componentTotal = promptTokens + completionTokens;
  const aggregateRaw = usage?.total;
  const aggregateTotal =
    typeof aggregateRaw === "number" && Number.isFinite(aggregateRaw)
      ? Math.max(0, aggregateRaw)
      : undefined;
  const totalTokens =
    aggregateTotal !== undefined ? Math.max(componentTotal, aggregateTotal) : componentTotal;

  const reasoningTokens = normalizeTokenCount(usage?.reasoningTokens);
  return {
    prompt_tokens: promptTokens,
    completion_tokens: completionTokens,
    total_tokens: totalTokens,
    ...(cacheRead > 0 ? { prompt_tokens_details: { cached_tokens: cacheRead } } : {}),
    ...(reasoningTokens !== undefined
      ? { completion_tokens_details: { reasoning_tokens: reasoningTokens } }
      : {}),
  };
}

/** Derive prompt/context tokens from normalized input and cache buckets. */
export function derivePromptTokens(usage?: {
  input?: number;
  cacheRead?: number;
  cacheWrite?: number;
}): number | undefined {
  if (!usage) {
    return undefined;
  }
  const input = usage.input ?? 0;
  const cacheRead = usage.cacheRead ?? 0;
  const cacheWrite = usage.cacheWrite ?? 0;
  const sum = input + cacheRead + cacheWrite;
  return sum > 0 ? sum : undefined;
}

function derivePromptTokensFromTotal(usage?: NormalizedUsage): number | undefined {
  const total = usage?.total;
  const output = usage?.output;
  if (
    typeof total !== "number" ||
    !Number.isFinite(total) ||
    total <= 0 ||
    typeof output !== "number" ||
    !Number.isFinite(output) ||
    output < 0
  ) {
    return undefined;
  }
  const promptTokens = total - output;
  return promptTokens > 0 ? promptTokens : undefined;
}

/** Resolve context prompt tokens from explicit override, last call, or aggregate usage. */
export function deriveContextPromptTokens(params: {
  lastCallUsage?: NormalizedUsage;
  promptTokens?: number;
  usage?: NormalizedUsage;
}): number | undefined {
  const promptOverride = params.promptTokens;
  if (typeof promptOverride === "number" && Number.isFinite(promptOverride) && promptOverride > 0) {
    return promptOverride;
  }

  if (params.lastCallUsage?.contextUsage?.state === "unavailable") {
    return undefined;
  }
  if (params.lastCallUsage?.contextUsage?.state === "available") {
    return params.lastCallUsage.contextUsage.promptTokens;
  }
  const lastCallPromptTokens =
    derivePromptTokens(params.lastCallUsage) ?? derivePromptTokensFromTotal(params.lastCallUsage);
  if (lastCallPromptTokens !== undefined) {
    return lastCallPromptTokens;
  }
  if (params.usage?.contextUsage?.state === "unavailable") {
    return undefined;
  }
  if (params.usage?.contextUsage?.state === "available") {
    return params.usage.contextUsage.promptTokens;
  }
  return derivePromptTokens(params.usage);
}

/** Derive the session prompt-token snapshot stored for context display. */
export function deriveSessionTotalTokens(params: {
  usage?: {
    input?: number;
    output?: number;
    total?: number;
    cacheRead?: number;
    cacheWrite?: number;
    contextUsage?: ContextUsage;
  };
  contextTokens?: number;
  promptTokens?: number;
}): number | undefined {
  const promptOverride = params.promptTokens;
  const hasPromptOverride =
    typeof promptOverride === "number" && Number.isFinite(promptOverride) && promptOverride > 0;

  const usage = params.usage;
  if (!usage && !hasPromptOverride) {
    return undefined;
  }

  // NOTE: SessionEntry.totalTokens is used as a prompt/context snapshot.
  // It intentionally excludes completion/output tokens.
  const promptTokens = deriveContextPromptTokens({
    promptTokens: hasPromptOverride ? promptOverride : undefined,
    usage,
  });

  if (!(typeof promptTokens === "number") || !Number.isFinite(promptTokens) || promptTokens <= 0) {
    return undefined;
  }

  // Keep this value unclamped; display layers are responsible for capping
  // percentages for terminal output.
  return promptTokens;
}