harness-tools/extensions/todo-herdr/src/widget.ts

import { truncateToWidth, visibleWidth } from "@earendil-works/pi-tui";
import type { ExtensionContext, Theme } from "@earendil-works/pi-coding-agent";
import type { TaskState } from "./state.ts";
import { renderTodoPanel } from "./ui.ts";

const WIDGET_ID = "pi-todo-herdr";
const COMBINED_WIDGET_BREAKPOINT = 120;
const COMBINED_BODY_ROWS = 10;
export const TODO_ANIMATION_INTERVAL_MS = 500;

interface StatusWidgetTheme {
  fg(color: string, text: string): string;
}

interface StatusWidget {
  readonly source: object;
  readonly hasContent: () => boolean;
  readonly getBodyRowCount: () => number;
  readonly render: (width: number, maxBodyRows?: number, theme?: StatusWidgetTheme) => string[];
  readonly claim: (onChange: () => void) => () => void;
  readonly isClaimed: (onChange: () => void) => boolean;
}

interface InteractiveSubagentsGlobal {
  statusWidget?: unknown;
}

class OptionalIntegrationError extends Error {
  constructor(message: string, cause?: unknown) {
    super(message, { cause });
    this.name = "OptionalIntegrationError";
  }
}

class TuiRenderFailure extends Error {
  readonly original: unknown;

  constructor(original: unknown) {
    super("TUI render failed", { cause: original });
    this.name = "TuiRenderFailure";
    this.original = original;
  }
}

function optionalIntegrationFailure(message: string, cause?: unknown): OptionalIntegrationError {
  return new OptionalIntegrationError(message, cause);
}

function callOptionalIntegration<T>(message: string, operation: () => T): T {
  try {
    return operation();
  } catch (error) {
    if (error instanceof TuiRenderFailure) throw error.original;
    if (error instanceof OptionalIntegrationError) throw error;
    throw optionalIntegrationFailure(message, error);
  }
}

interface ActiveClaim {
  owner: object;
  source: object;
  onChange: () => void;
  release: () => void;
  generation: number;
}

// Symbols make the claim and generation survive extension reloads. A newly
// loaded module can invalidate a stale component before installing its own.
const ACTIVE_CLAIM_KEY = Symbol.for("pi-todo-herdr/active-subagents-claim");
const WIDGET_GENERATION_KEY = Symbol.for("pi-todo-herdr/widget-generation");
const ACTIVE_RUNTIME_KEY = Symbol.for("pi-todo-herdr/active-runtime");
const RUNTIME_TOKEN = Symbol("pi-todo-herdr-runtime");
const previousClaim = Reflect.get(globalThis, ACTIVE_CLAIM_KEY);
Reflect.set(globalThis, ACTIVE_RUNTIME_KEY, RUNTIME_TOKEN);
if (previousClaim !== null && typeof previousClaim === "object") {
  const release = (previousClaim as { release?: unknown }).release;
  Reflect.deleteProperty(globalThis, ACTIVE_CLAIM_KEY);
  if (typeof release === "function") {
    try {
      release.call(previousClaim);
    } catch {
      // A stale claim must not prevent the new extension runtime loading.
    }
  }
}

function isCurrentRuntime(): boolean {
  return Reflect.get(globalThis, ACTIVE_RUNTIME_KEY) === RUNTIME_TOKEN;
}

function getActiveClaim(): ActiveClaim | undefined {
  const claim = Reflect.get(globalThis, ACTIVE_CLAIM_KEY);
  return claim !== null && typeof claim === "object" ? (claim as ActiveClaim) : undefined;
}

function setActiveClaim(claim: ActiveClaim | undefined): void {
  if (claim === undefined) Reflect.deleteProperty(globalThis, ACTIVE_CLAIM_KEY);
  else Reflect.set(globalThis, ACTIVE_CLAIM_KEY, claim);
}

function claimBelongsTo(claim: ActiveClaim | undefined, owner: object, generation: number): claim is ActiveClaim {
  return claim?.owner === owner && claim.generation === generation;
}

function nextWidgetGeneration(): number {
  const current = Reflect.get(globalThis, WIDGET_GENERATION_KEY);
  const generation = typeof current === "number" && Number.isSafeInteger(current) ? current + 1 : 1;
  Reflect.set(globalThis, WIDGET_GENERATION_KEY, generation);
  return generation;
}

function currentWidgetGeneration(): number {
  const generation = Reflect.get(globalThis, WIDGET_GENERATION_KEY);
  return typeof generation === "number" && Number.isSafeInteger(generation) ? generation : 0;
}

function releaseActiveClaim(owner?: object, generation?: number): void {
  const activeClaim = getActiveClaim();
  if (
    (owner !== undefined && activeClaim?.owner !== owner) ||
    (generation !== undefined && activeClaim?.generation !== generation)
  ) return;
  setActiveClaim(undefined);
  if (activeClaim === undefined) return;
  try {
    activeClaim.release();
  } catch {
    // A broken optional integration must never prevent the Todo widget clearing.
  }
}

function getStatusWidget(): StatusWidget | undefined {
  try {
    const globals = globalThis as typeof globalThis & {
      __pi_interactive_subagents?: unknown;
    };
    const integration = globals.__pi_interactive_subagents;
    if (integration === null || typeof integration !== "object") return undefined;
    const statusWidget = (integration as InteractiveSubagentsGlobal).statusWidget;
    if (statusWidget === null || typeof statusWidget !== "object") return undefined;

    const candidate = statusWidget as Record<string, unknown>;
    if (
      typeof candidate.hasContent !== "function" ||
      typeof candidate.getBodyRowCount !== "function" ||
      typeof candidate.render !== "function" ||
      typeof candidate.claim !== "function" ||
      typeof candidate.isClaimed !== "function"
    ) {
      return undefined;
    }

    const source = statusWidget;
    const call = <T>(methodName: string, args: unknown[] = []): T =>
      callOptionalIntegration(`Subagents ${methodName} failed`, () => {
        const method = candidate[methodName];
        if (typeof method !== "function") {
          throw optionalIntegrationFailure(`invalid Subagents ${methodName} method`);
        }
        return method.call(source, ...args) as T;
      });

    return {
      source,
      hasContent: () => {
        const content = call<unknown>("hasContent");
        if (typeof content !== "boolean") throw optionalIntegrationFailure("invalid Subagents content state");
        return content;
      },
      getBodyRowCount: () => {
        const count = call<unknown>("getBodyRowCount");
        if (typeof count !== "number" || !Number.isFinite(count)) {
          throw optionalIntegrationFailure("invalid Subagents body row count");
        }
        return count;
      },
      render: (width, maxBodyRows, theme) => {
        const lines = call<unknown>("render", [width, maxBodyRows, theme]);
        const hasAllEntries =
          Array.isArray(lines) &&
          lines.length > 0 &&
          Array.from({ length: lines.length }, (_, index) => {
            return Object.prototype.hasOwnProperty.call(lines, index) && typeof lines[index] === "string";
          }).every(Boolean);
        if (!hasAllEntries) throw optionalIntegrationFailure("invalid Subagents render output");
        const validLines = lines as string[];
        if (maxBodyRows !== undefined) {
          const isBordered = validLines[0]?.includes("╭") && validLines.at(-1)?.includes("╰");
          // Body rows exclude the provider's header. Legacy bordered output has
          // top and bottom chrome; the current unboxed output has one header.
          const headerAndChromeRows = isBordered ? 2 : 1;
          const maxLines = maxBodyRows + headerAndChromeRows;
          if (validLines.length > maxLines) {
            throw optionalIntegrationFailure("Subagents render output exceeded its body limit");
          }
        }
        return validLines;
      },
      claim: (onChange) => {
        const release = call<unknown>("claim", [onChange]);
        if (typeof release !== "function") throw optionalIntegrationFailure("invalid Subagents claim release");
        return release as () => void;
      },
      isClaimed: (onChange) => {
        const claimed = call<unknown>("isClaimed", [onChange]);
        if (typeof claimed !== "boolean") throw optionalIntegrationFailure("invalid Subagents claim state");
        return claimed;
      },
    };
  } catch {
    return undefined;
  }
}

function claimStatusWidget(
  statusWidget: StatusWidget,
  owner: object,
  generation: number,
  tui: { requestRender(): void },
): boolean {
  const activeClaim = getActiveClaim();
  if (activeClaim?.owner !== undefined && activeClaim.owner !== owner) return false;
  if (claimBelongsTo(activeClaim, owner, generation) && activeClaim.source === statusWidget.source) {
    try {
      if (statusWidget.isClaimed(activeClaim.onChange)) return true;
    } catch (error) {
      releaseActiveClaim(owner, generation);
      if (error instanceof OptionalIntegrationError) return false;
      throw error;
    }
    // The provider may have been reset by a reload or another lifecycle event.
    // Drop the stale local handle before asking it for a fresh claim.
  }
  releaseActiveClaim(owner, generation);

  try {
    const onChange = () => {
      if (!isCurrentRuntime()) {
        releaseActiveClaim(owner, generation);
        return;
      }
      const currentClaim = getActiveClaim();
      if (!claimBelongsTo(currentClaim, owner, generation) || currentClaim.onChange !== onChange) return;
      try {
        tui.requestRender();
      } catch (error) {
        // Let the provider restore its standalone widget when the TUI is gone.
        releaseActiveClaim(owner, generation);
        throw new TuiRenderFailure(error);
      }
    };
    const release = statusWidget.claim(onChange);
    setActiveClaim({ owner, source: statusWidget.source, onChange, release, generation });
    return true;
  } catch (error) {
    if (error instanceof OptionalIntegrationError) return false;
    throw error;
  }
}

function normalizeWidth(width: number): number {
  return Math.max(1, Math.floor(Number.isFinite(width) ? width : 1));
}

function fitLine(line: string, width: number): string {
  const fittedWidth = normalizeWidth(width);
  const fitted = truncateToWidth(line, fittedWidth);
  return `${fitted}${" ".repeat(Math.max(0, fittedWidth - visibleWidth(fitted)))}`;
}

function renderStatusPanel(
  statusWidget: StatusWidget,
  width: number,
  maxBodyRows: number,
  theme: StatusWidgetTheme,
): string[] {
  return statusWidget.render(width, maxBodyRows, theme).map((line) => fitLine(line, width));
}

function normalizePanelHeight(lines: string[], height: number, blankLine: string): string[] {
  if (lines.length >= height) return lines;
  const blanks = Array.from({ length: height - lines.length }, () => blankLine);
  const bottom = lines.at(-1);
  if (bottom?.includes("╰")) return [...lines.slice(0, -1), ...blanks, bottom];
  return [...lines, ...blanks];
}

function renderSideBySide(
  tasks: ReturnType<TaskState["list"]>,
  theme: Theme,
  width: number,
  inProgressFrame: number,
  statusWidget: StatusWidget,
): string[] {
  const leftWidth = Math.floor((width - 1) / 2);
  const rightWidth = width - 1 - leftWidth;
  const todo = renderTodoPanel(tasks, theme, leftWidth, inProgressFrame, COMBINED_BODY_ROWS);
  const subagents = renderStatusPanel(statusWidget, rightWidth, COMBINED_BODY_ROWS, theme);
  const height = Math.max(todo.length, subagents.length);
  const left = normalizePanelHeight(todo, height, " ".repeat(leftWidth)).map((line) => fitLine(line, leftWidth));
  const right = normalizePanelHeight(subagents, height, " ".repeat(rightWidth));
  return left.map((line, index) => `${line} ${right[index]}`);
}

function bodyRowsForTodo(
  tasks: ReturnType<TaskState["list"]>,
  theme: Theme,
  width: number,
  inProgressFrame: number,
): number {
  return Math.max(1, renderTodoPanel(tasks, theme, width, inProgressFrame, COMBINED_BODY_ROWS).length - 1);
}

function bodyRowsForSubagents(statusWidget: StatusWidget): number {
  return Math.max(1, Math.min(COMBINED_BODY_ROWS, Math.floor(statusWidget.getBodyRowCount())));
}

function distributeBodyRows(todoNeed: number, subagentsNeed: number): { todo: number; subagents: number } {
  let todo = Math.min(todoNeed, Math.floor(COMBINED_BODY_ROWS / 2));
  let subagents = Math.min(subagentsNeed, Math.floor(COMBINED_BODY_ROWS / 2));
  let remaining = COMBINED_BODY_ROWS - todo - subagents;

  const todoUnused = Math.max(0, todoNeed - todo);
  const todoExtra = Math.min(remaining, todoUnused);
  todo += todoExtra;
  remaining -= todoExtra;

  const subagentsUnused = Math.max(0, subagentsNeed - subagents);
  subagents += Math.min(remaining, subagentsUnused);
  return { todo: Math.max(1, todo), subagents: Math.max(1, subagents) };
}

function renderStacked(
  tasks: ReturnType<TaskState["list"]>,
  theme: Theme,
  width: number,
  inProgressFrame: number,
  statusWidget: StatusWidget,
  subagentsNeed: number,
): string[] {
  const todoNeed = bodyRowsForTodo(tasks, theme, width, inProgressFrame);
  const rows = distributeBodyRows(todoNeed, subagentsNeed);
  return [
    ...renderTodoPanel(tasks, theme, width, inProgressFrame, rows.todo),
    "",
    ...renderStatusPanel(statusWidget, width, rows.subagents, theme),
  ];
}

function renderCombinedWidget(
  tasks: ReturnType<TaskState["list"]>,
  theme: Theme,
  width: number,
  inProgressFrame: number,
  tui: { requestRender(): void },
  owner: object,
  generation: number,
  composeSubagents: boolean,
): string[] {
  if (!Number.isFinite(width) || width <= 0 || tasks.length === 0) {
    releaseActiveClaim(owner, generation);
    return [];
  }

  const fittedWidth = normalizeWidth(width);
  const renderTodoFallback = (): string[] => {
    releaseActiveClaim(owner, generation);
    return renderTodoPanel(tasks, theme, fittedWidth, inProgressFrame);
  };
  if (!composeSubagents) return renderTodoFallback();

  const statusWidget = getStatusWidget();
  if (statusWidget === undefined) return renderTodoFallback();

  try {
    if (!statusWidget.hasContent()) return renderTodoFallback();
    if (!claimStatusWidget(statusWidget, owner, generation, tui)) return renderTodoFallback();

    const subagentsNeed = bodyRowsForSubagents(statusWidget);
    return fittedWidth > COMBINED_WIDGET_BREAKPOINT
      ? renderSideBySide(tasks, theme, fittedWidth, inProgressFrame, statusWidget)
      : renderStacked(tasks, theme, fittedWidth, inProgressFrame, statusWidget, subagentsNeed);
  } catch (error) {
    if (!(error instanceof OptionalIntegrationError)) {
      releaseActiveClaim(owner, generation);
      throw error;
    }
    return renderTodoFallback();
  }
}

export function clearTaskWidget(ctx: ExtensionContext): void {
  if (!isCurrentRuntime()) return;
  nextWidgetGeneration();
  releaseActiveClaim();
  ctx.ui.setWidget(WIDGET_ID, undefined);
}

export function updateTaskWidget(
  state: TaskState,
  ctx: ExtensionContext,
  visible: boolean,
  refreshMs = TODO_ANIMATION_INTERVAL_MS,
): void {
  if (!isCurrentRuntime()) return;
  const tasks = state.list();
  if (!visible || tasks.length === 0) {
    clearTaskWidget(ctx);
    return;
  }

  const generation = nextWidgetGeneration();
  // Subagents owns a native widget outside this extension. Only compose it
  // while Pi is rendering the terminal UI, never in RPC/print modes.
  const composeSubagents = ctx.mode === "tui";

  // Replacing a widget can happen before Pi calls its old dispose hook. Release
  // the optional Subagents claim eagerly so it cannot hide a standalone widget.
  releaseActiveClaim();
  const hasActiveTimer = tasks.some((task) => task.status === "in-progress" && task.startedAt !== undefined);
  ctx.ui.setWidget(WIDGET_ID, (tui, theme) => {
    let inProgressFrame = 0;
    let disposed = false;
    const owner = {};
    const isActive = (): boolean =>
      !disposed && isCurrentRuntime() && currentWidgetGeneration() === generation;
    let timer: ReturnType<typeof setInterval> | undefined;
    if (hasActiveTimer) {
      timer = setInterval(() => {
        if (!isActive()) {
          clearInterval(timer);
          return;
        }
        inProgressFrame += 1;
        try {
          tui.requestRender();
        } catch {
          disposed = true;
          clearInterval(timer);
          releaseActiveClaim(owner, generation);
        }
      }, Math.max(1, refreshMs));
    }
    timer?.unref();

    return {
      render(width: number): string[] {
        if (!isActive()) return [];
        const currentTasks = state.list();
        return renderCombinedWidget(
          currentTasks,
          theme,
          width,
          inProgressFrame,
          tui,
          owner,
          generation,
          composeSubagents,
        );
      },
      invalidate(): void {},
      dispose(): void {
        disposed = true;
        releaseActiveClaim(owner, generation);
        if (timer !== undefined) clearInterval(timer);
      },
    };
  });
}