memind/billing-token-state.mjs

import { normalizeTokenState } from './billing.mjs';

export function loadCostEstimateConfig(env = process.env) {
  const useBackendCost = env.H5_USE_BACKEND_COST === '1';
  const enabled = useBackendCost && env.H5_COST_ESTIMATE_FROM_TOKENS !== '0';
  const inputUsdPer1M = Number(env.H5_COST_ESTIMATE_INPUT_USD_PER_1M ?? 0.27);
  const outputUsdPer1M = Number(env.H5_COST_ESTIMATE_OUTPUT_USD_PER_1M ?? 1.1);
  return {
    enabled,
    inputUsdPer1M: Number.isFinite(inputUsdPer1M) && inputUsdPer1M >= 0 ? inputUsdPer1M : 0.27,
    outputUsdPer1M: Number.isFinite(outputUsdPer1M) && outputUsdPer1M >= 0 ? outputUsdPer1M : 1.1,
  };
}

export function pickSessionAccumulatedCost(session) {
  if (!session || typeof session !== 'object') return null;
  const raw = session.accumulatedCost ?? session.accumulated_cost ?? null;
  if (raw == null) return null;
  const value = Number(raw);
  return Number.isFinite(value) && value >= 0 ? value : null;
}

export function estimateAccumulatedCostUsd(tokenStateRaw, estimateConfig = loadCostEstimateConfig()) {
  if (!estimateConfig?.enabled) return null;
  const state = normalizeTokenState(tokenStateRaw);
  if (state.accumulatedInputTokens <= 0 && state.accumulatedOutputTokens <= 0) return null;
  const inputUsd =
    (state.accumulatedInputTokens * estimateConfig.inputUsdPer1M) / 1_000_000;
  const outputUsd =
    (state.accumulatedOutputTokens * estimateConfig.outputUsdPer1M) / 1_000_000;
  const total = inputUsd + outputUsd;
  return total > 0 ? total : null;
}

export function enrichTokenStateForBilling(
  tokenStateRaw,
  { sessionCost = null } = {},
  env = process.env,
) {
  const state = normalizeTokenState(tokenStateRaw);
  if (state.accumulatedCost != null && Number(state.accumulatedCost) >= 0) {
    return state;
  }

  const sessionUsd = pickSessionAccumulatedCost(sessionCost);
  if (sessionUsd != null) {
    return { ...state, accumulatedCost: sessionUsd };
  }

  const estimatedUsd = estimateAccumulatedCostUsd(state, loadCostEstimateConfig(env));
  if (estimatedUsd != null) {
    return { ...state, accumulatedCost: estimatedUsd };
  }

  return state;
}

export async function resolveBillingTokenState(
  tokenStateRaw,
  { sessionId = null, fetchSession = null } = {},
  env = process.env,
) {
  const state = normalizeTokenState(tokenStateRaw);
  if (state.accumulatedCost != null && Number(state.accumulatedCost) >= 0) {
    return state;
  }

  if (typeof fetchSession === 'function' && sessionId) {
    try {
      const session = await fetchSession(sessionId);
      const enriched = enrichTokenStateForBilling(state, { sessionCost: session }, env);
      if (enriched.accumulatedCost != null) return enriched;
    } catch {
      // Best-effort: fall through to token estimate.
    }
  }

  return enrichTokenStateForBilling(state, {}, env);
}