openchamber/packages/ui/src/components/chat/work-status/contextUsage.ts

/**
 * Context-window usage for a specific session.
 *
 * `useSessionUIStore.getContextUsage` cannot serve this panel. It reads
 * `getSyncMessages(sessionId)` with **no directory**, which resolves to the
 * *current* directory's child store, and it keys off the store's own
 * `currentSessionId`. A session held by another directory — a worktree, or any
 * moment right after a directory switch — therefore reads as "no messages" and
 * the readout silently disappears while the header still shows a value.
 *
 * This computes the same quantity from messages the caller has already
 * subscribed to for a known session and directory, so there is no hidden
 * global read to race with.
 */

import { contextTokensFromBreakdown } from '@/stores/utils/tokenUtils';

type MessageTokens = {
  /** Server-reported window of the turn's final round-trip; absent on older servers. */
  total?: number;
  input?: number;
  output?: number;
  reasoning?: number;
  cache?: { read?: number; write?: number };
};

type MessageLike = {
  id?: string;
  role?: string;
  tokens?: MessageTokens;
};

type WorkStatusContextUsage = {
  totalTokens: number;
  /** Context limit actually used for the ratio, after the default fallback. */
  limit: number;
  /** Unrounded, so the panel and the header cannot disagree by a rounding step. */
  percent: number;
};

/** The store's own fallback when a model exposes no context limit. */
export const DEFAULT_CONTEXT_LIMIT = 200_000;

/**
 * Usage from the newest assistant message that reported a non-zero token count.
 * The latest turn describes the current fill — not a sum across turns. Within
 * a turn, the server-reported `total` is the final round-trip's window;
 * summing the breakdown fields instead overstates multi-step turns, whose
 * input/cache fields accumulate across round-trips.
 */
export const computeContextUsage = (
  messages: readonly MessageLike[],
  contextLimit: number,
): WorkStatusContextUsage | null => {
  if (messages.length === 0) return null;

  for (let index = messages.length - 1; index >= 0; index -= 1) {
    const message = messages[index];
    if (message?.role !== 'assistant' || !message.tokens) continue;

    const totalTokens = contextTokensFromBreakdown(message.tokens);
    if (totalTokens <= 0) continue;

    const limit = contextLimit > 0 ? contextLimit : DEFAULT_CONTEXT_LIMIT;
    return { totalTokens, limit, percent: (totalTokens / limit) * 100 };
  }

  return null;
};