memind/billing-token-state.mjs

import { normalizeTokenState } from './billing.mjs';
import { fetchGooseSessionAccumulatedCostUsd } from './goose-session-cost.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);

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

  if (state.accumulatedCost != null && Number(state.accumulatedCost) >= 0) {
    return state;
  }

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

  return state;
}

async function resolveSessionCostPayload(sessionId, fetchSession, fetchSessionCostFromPg, env) {
  let sessionCost = null;
  if (typeof fetchSession === 'function' && sessionId) {
    try {
      sessionCost = await fetchSession(sessionId);
    } catch {
      sessionCost = null;
    }
  }
  if (pickSessionAccumulatedCost(sessionCost) != null) {
    return sessionCost;
  }

  const readPgCost =
    typeof fetchSessionCostFromPg === 'function'
      ? fetchSessionCostFromPg
      : (sid) => fetchGooseSessionAccumulatedCostUsd(sid, env);
  const pgUsd = env.H5_USE_BACKEND_COST === '1' ? await readPgCost(sessionId) : null;
  if (pgUsd != null) {
    return {
      ...(sessionCost && typeof sessionCost === 'object' ? sessionCost : {}),
      accumulated_cost: pgUsd,
    };
  }
  return sessionCost;
}

export async function resolveBillingTokenState(
  tokenStateRaw,
  { sessionId = null, fetchSession = null, fetchSessionCostFromPg = null } = {},
  env = process.env,
) {
  const state = normalizeTokenState(tokenStateRaw);
  const sessionCost = await resolveSessionCostPayload(
    sessionId,
    fetchSession,
    fetchSessionCostFromPg,
    env,
  );
  return enrichTokenStateForBilling(state, { sessionCost }, env);
}