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;
}

export function createState(): StatuslineState {
  return {
    currentCtx: null,
    thinkingLevel: "off",
    statusEntries: [],
    branch: "",
    requestRender: null,
    footerData: null,
    originalSetFooter: null,
    footerHijacked: false,
    lastRenderSignature: "",
  };
}

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: collectUsage(ctx),
    statuses: [...state.statusEntries],
  };
}

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 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;
}