wow-pi/packages/wow-statusline/src/state.ts

import type { ExtensionContext } from "@earendil-works/pi-coding-agent";
import type {
  FooterDataSnapshot,
  SessionSnapshot,
  UsageSummary,
} from "./types";

export interface StatuslineState {
  currentCtx: ExtensionContext | null;
  thinkingLevel: string;
  statusEntries: Array<[string, string]>;
  branch: string;
  requestRender: (() => void) | null;
  footerData: FooterDataSnapshot | null;
  originalSetFooter: ((...args: unknown[]) => unknown) | null;
  footerHijacked: boolean;
  lastRenderSignature: string;
  /**
   * Cached token usage summary. Refreshed by `updateStatuslineState` before
   * each signature check; the cached value is what `snapshotFromContext`
   * returns to the renderer, so display does not re-walk the session branch.
   */
  cachedUsage: UsageSummary;
  /** Leaf id at the time `cachedUsage` was computed. */
  cachedUsageLeafId: string;
  /**
   * Coalescing flag for the nextTick that re-asserts our empty footer after
   * another extension calls `setFooter`. Prevents microtask ping-pong when
   * multiple extensions compete for the slot.
   */
  nextTickResetScheduled: boolean;
  /**
   * Whether we actually installed our custom editor component. When another
   * extension's editor is already set, we defer and leave this false so
   * `clearStatuslineUI` does not clobber the other extension.
   */
  editorInstalled: boolean;
}

export function createState(): StatuslineState {
  return {
    currentCtx: null,
    thinkingLevel: "off",
    statusEntries: [],
    branch: "",
    requestRender: null,
    footerData: null,
    originalSetFooter: null,
    footerHijacked: false,
    lastRenderSignature: "",
    cachedUsage: { input: 0, output: 0, cost: 0 },
    cachedUsageLeafId: "",
    nextTickResetScheduled: false,
    editorInstalled: false,
  };
}

export function snapshotFromContext(
  state: StatuslineState,
  ctx: ExtensionContext,
): SessionSnapshot {
  return {
    cwd: ctx.cwd,
    branch: state.branch,
    modelName: ctx.model?.name || ctx.model?.id || "no-model",
    provider: ctx.model?.provider,
    thinkingLevel: state.thinkingLevel,
    usage: state.cachedUsage,
    statuses: [...state.statusEntries],
  };
}

/**
 * Refresh `state.cachedUsage` when the session leaf changes. Returns true if
 * usage was recomputed, so callers can decide whether to also recompute the
 * render signature.
 */
export function refreshUsageCache(
  state: StatuslineState,
  ctx: ExtensionContext,
): boolean {
  const leafId = readLeafId(ctx);
  if (leafId === state.cachedUsageLeafId) return false;
  state.cachedUsageLeafId = leafId;
  state.cachedUsage = collectUsage(ctx);
  return true;
}

export function syncFooterSnapshot(state: StatuslineState): void {
  if (!state.footerData) return;
  state.branch = state.footerData.getGitBranch() || "";
  const statuses = state.footerData.getExtensionStatuses();
  if (statuses.size === 0) return;
  state.statusEntries = Array.from(statuses.entries()).sort(([left], [right]) =>
    left.localeCompare(right),
  );
}

function readLeafId(ctx: ExtensionContext): string {
  try {
    return ctx.sessionManager.getLeafId() ?? "";
  } catch {
    return "";
  }
}

function collectUsage(ctx: ExtensionContext): UsageSummary {
  let input = 0;
  let output = 0;
  let cost = 0;

  try {
    for (const entry of ctx.sessionManager.getBranch()) {
      if (entry.type !== "message" || entry.message.role !== "assistant")
        continue;
      const usage = readUsage(entry.message);
      if (!usage) continue;
      input += usage.input;
      output += usage.output;
      cost += usage.cost;
    }
  } catch {
    // Session access can fail during teardown or special modes.
  }

  return { input, output, cost };
}

function readUsage(message: unknown): UsageSummary | null {
  if (!message || typeof message !== "object") return null;
  const usage = (message as { usage?: unknown }).usage;
  if (!usage || typeof usage !== "object") return null;

  const input = readNumber((usage as { input?: unknown }).input);
  const output = readNumber((usage as { output?: unknown }).output);
  const costValue =
    usage && typeof usage === "object"
      ? readNumber(
          (
            (usage as { cost?: unknown }).cost as
              | { total?: unknown }
              | undefined
          )?.total,
        )
      : 0;

  return { input, output, cost: costValue };
}

function readNumber(value: unknown): number {
  return typeof value === "number" && Number.isFinite(value) ? value : 0;
}