@rosetears/aili-pi 0.1.0 → 0.1.3

package/extensions/zentui/index.ts

@@ -0,0 +1,580 @@
+import type {
+	ExtensionAPI,
+	ExtensionContext,
+	KeybindingsManager,
+	Theme,
+} from "@earendil-works/pi-coding-agent";
+import type { EditorTheme, TUI } from "@earendil-works/pi-tui";
+import {
+	type ColorSourcesConfig,
+	type ContextStyle,
+	type ExtensionStatusColorMode,
+	type ExtensionStatusPlacement,
+	ensureConfigExists,
+	type FixedEditorConfig,
+	type FooterSegmentsConfig,
+	type GitBranchConfig,
+	type IconMode,
+	loadConfig,
+	type PathDisplayConfig,
+	type PolishedTuiConfig,
+	type SeparatorStyle,
+	saveColorSourcesPatch,
+	saveContextStylePatch,
+	saveExtensionStatusColorMode,
+	saveExtensionStatusPlacement,
+	saveFixedEditorPatch,
+	saveFooterFormatPatch,
+	saveFooterSegmentsPatch,
+	saveGitBranchPatch,
+	saveIconsModePatch,
+	savePathDisplayPatch,
+	saveSeparatorPatch,
+	saveUiFeaturesPatch,
+	type UiFeaturesConfig,
+} from "./config.js";
+import {
+	disposeFixedEditor,
+	installFixedEditorProbe,
+	removeFixedEditorProbe,
+} from "./fixed-editor/index.js";
+import { installFooter } from "./footer.js";
+import { buildSessionDurationLabel, invalidateUsageTotalsCache } from "./format.js";
+import { emptyGitStatus, readGitStatus } from "./git.js";
+import { LiveContextController } from "./live-context.js";
+import { readPackageVersionResult } from "./package-version.js";
+import {
+	createProjectRefreshScheduler,
+	type ScheduleProjectRefreshOptions,
+	type StopProjectRefreshInterval,
+	startProjectRefreshInterval,
+} from "./project-refresh.js";
+import { applyProjectRefreshToState } from "./project-state.js";
+import { readRuntimeInfo } from "./runtime.js";
+import { installSelectorBorderStyle } from "./selector-border.js";
+import { SessionLifecycle } from "./session-lifecycle.js";
+import { registerZentuiSettingsCommand } from "./settings-command.js";
+import { createInitialState, type FooterState, syncState } from "./state.js";
+import { installThinkingMessageStyle } from "./thinking-message.js";
+import { installToolExecutionStyle } from "./tool-execution.js";
+import { PolishedEditor, WrappedPolishedEditor } from "./ui.js";
+import { installUserMessageStyle } from "./user-message.js";
+
+const ZENTUI_EDITOR_FACTORY = Symbol.for("pi-zentui.editor-factory");
+const ZENTUI_EDITOR_BASE_FACTORY = Symbol.for("pi-zentui.editor-base-factory");
+
+type EditorFactory = NonNullable<Parameters<ExtensionContext["ui"]["setEditorComponent"]>[0]>;
+
+type ZentuiEditorFactory = EditorFactory & {
+	[ZENTUI_EDITOR_FACTORY]?: true;
+	[ZENTUI_EDITOR_BASE_FACTORY]?: EditorFactory;
+};
+
+type ApplyUiResult = {
+	editorBlocked: boolean;
+};
+
+type EditorInstallMode = "none" | "standalone" | "wrapper";
+
+function isZentuiEditorFactory(factory: EditorFactory | undefined): boolean {
+	return Boolean((factory as ZentuiEditorFactory | undefined)?.[ZENTUI_EDITOR_FACTORY]);
+}
+
+function getZentuiEditorBaseFactory(factory: EditorFactory | undefined): EditorFactory | undefined {
+	return (factory as ZentuiEditorFactory | undefined)?.[ZENTUI_EDITOR_BASE_FACTORY];
+}
+
+function isTuiContext(ctx: ExtensionContext): boolean {
+	try {
+		const mode = (ctx as ExtensionContext & { mode?: string }).mode;
+		return ctx.hasUI && (mode === undefined || mode === "tui");
+	} catch {
+		return false;
+	}
+}
+
+export default function (pi: ExtensionAPI) {
+	const state: FooterState = createInitialState(emptyGitStatus());
+	const sessionLifecycle = new SessionLifecycle();
+
+	let currentConfig: PolishedTuiConfig = loadConfig();
+	let activeTheme: Theme | undefined;
+	let requestFooterRender: (() => void) | undefined;
+	let getActiveExtensionStatuses: () => ReadonlyMap<string, string> = () => new Map();
+	let stopRefreshInterval: StopProjectRefreshInterval = () => {};
+	let cleanupPrototypePatches: () => void = () => {};
+	let footerInstalled = false;
+	let editorInstalled = false;
+	let editorInstallMode: EditorInstallMode = "none";
+	let installedEditorFactory: EditorFactory | undefined;
+	let wrappedEditorFactory: EditorFactory | undefined;
+	let prototypePatchesInstalled = false;
+	let stopSessionTimer: () => void = () => {};
+	let lastDurationLabel = "";
+	let lastProjectCwd: string | undefined;
+
+	const refresh = () => {
+		if (sessionLifecycle.isCurrent()) requestFooterRender?.();
+	};
+	const liveContext = new LiveContextController(sessionLifecycle, refresh);
+	const getActiveTheme = () => activeTheme;
+	const getCurrentConfig = () => currentConfig;
+	const getThinkingLevel = () =>
+		sessionLifecycle.isCurrent() ? pi.getThinkingLevel() : ("off" as const);
+	const syncFooterState = (ctx: ExtensionContext) =>
+		syncState(state, ctx, currentConfig.icons.cacheHit);
+
+	type ProjectRefreshTarget = { cwd: string; generation: number };
+	const refreshProjectState = async ({ cwd, generation }: ProjectRefreshTarget) => {
+		if (!sessionLifecycle.isCurrent(generation)) return;
+		const gitCommitConfig = currentConfig.gitCommit;
+		const gitMetricsConfig = currentConfig.gitMetrics;
+		const segments = currentConfig.footerSegments;
+		const fmt = currentConfig.footerFormat;
+		// Enable optional probes when the segment is on OR a custom footerFormat
+		// references the relevant variable. Mirrors the session-duration timer
+		// pattern so format-only users still get data.
+		const formatNeedsTag = /\$\{?(?:git_tag|tag)\b/.test(fmt);
+		const formatNeedsCommit = /\$\{?(?:git_commit|commit)\b/.test(fmt);
+		const formatNeedsMetrics = /\$\{?(?:git_metrics|git_added|git_deleted)\b/.test(fmt);
+		const formatNeedsPackage = /\$\{?(?:package|package_version)\b/.test(fmt);
+		const wantExactTag =
+			((segments.gitCommit || formatNeedsCommit) && gitCommitConfig.showTag) || formatNeedsTag;
+		const wantMetrics = segments.gitMetrics || formatNeedsMetrics;
+		const wantPackage = segments.packageVersion || formatNeedsPackage;
+		const [git, runtime, packageVersion] = await Promise.all([
+			readGitStatus(cwd, {
+				readExactTag: wantExactTag,
+				readMetrics: wantMetrics,
+				ignoreSubmodules: gitMetricsConfig.ignoreSubmodules,
+			}),
+			readRuntimeInfo(cwd),
+			wantPackage ? readPackageVersionResult(cwd) : Promise.resolve(undefined),
+		]);
+		if (!sessionLifecycle.isCurrent(generation)) return;
+		lastProjectCwd = applyProjectRefreshToState(state, {
+			cwd,
+			previousCwd: lastProjectCwd,
+			git,
+			runtime,
+			packageVersion,
+		});
+	};
+
+	const projectRefreshScheduler = createProjectRefreshScheduler(refreshProjectState, refresh);
+	const scheduleProjectRefresh = (
+		ctx: ExtensionContext,
+		options?: ScheduleProjectRefreshOptions,
+	) => {
+		const generation = sessionLifecycle.currentGeneration();
+		if (!sessionLifecycle.isCurrent(generation)) return;
+		const cwd = ctx.cwd;
+		projectRefreshScheduler.schedule({ cwd, generation }, options);
+	};
+
+	const refreshInteractiveState = (ctx: ExtensionContext, project = false) => {
+		if (!sessionLifecycle.isCurrent() || !ctx.hasUI) return;
+		syncFooterState(ctx);
+		if (project && currentConfig.features.statusLine) scheduleProjectRefresh(ctx);
+		refresh();
+	};
+
+	const stopProjectRefresh = () => {
+		stopRefreshInterval();
+		stopRefreshInterval = () => {};
+		projectRefreshScheduler.stop();
+	};
+
+	const startSessionTimer = () => {
+		stopSessionTimer();
+		lastDurationLabel = "";
+		const segments = currentConfig.footerSegments;
+		const format = currentConfig.footerFormat ?? "";
+		// The clock belongs to the editor's upper-right metadata, not the footer.
+		// Keep a lightweight refresh tick while that editor chrome is enabled.
+		const needsWallClock =
+			currentConfig.features.editor || segments.time || /\$\{?time\b/.test(format);
+		const needsDuration =
+			segments.sessionDuration || /\$\{?(?:session_duration|duration)\b/.test(format);
+		if (!(needsWallClock || (currentConfig.features.statusLine && needsDuration))) return;
+
+		const timer = setInterval(() => {
+			if (!sessionLifecycle.isCurrent()) return;
+			if (needsWallClock) {
+				refresh();
+				return;
+			}
+			const label = state.sessionStartEpoch
+				? buildSessionDurationLabel(state.sessionStartEpoch)
+				: "";
+			if (label === lastDurationLabel) return;
+			lastDurationLabel = label;
+			refresh();
+		}, 1000);
+		timer.unref?.();
+		stopSessionTimer = () => {
+			clearInterval(timer);
+			stopSessionTimer = () => {};
+		};
+	};
+
+	const installPrototypePatches = () => {
+		if (prototypePatchesInstalled) return;
+		const cleanupSelectorBorderStyle = installSelectorBorderStyle(getActiveTheme, getCurrentConfig);
+		const cleanupUserMessageStyle = installUserMessageStyle(getActiveTheme, getCurrentConfig);
+		const cleanupToolExecutionStyle = installToolExecutionStyle(getActiveTheme);
+		const cleanupThinkingMessageStyle = installThinkingMessageStyle(getActiveTheme);
+		cleanupPrototypePatches = () => {
+			cleanupThinkingMessageStyle();
+			cleanupToolExecutionStyle();
+			cleanupSelectorBorderStyle();
+			cleanupUserMessageStyle();
+		};
+		prototypePatchesInstalled = true;
+	};
+
+	const uninstallPrototypePatches = () => {
+		cleanupPrototypePatches();
+		cleanupPrototypePatches = () => {};
+		prototypePatchesInstalled = false;
+	};
+
+	const makeEditorFactory = (ctx: ExtensionContext): ZentuiEditorFactory => {
+		const sessionTheme = ctx.ui.theme;
+		const factory = ((tui: TUI, theme: EditorTheme, keybindings: KeybindingsManager) =>
+			new PolishedEditor(
+				tui,
+				theme,
+				keybindings,
+				sessionTheme,
+				getCurrentConfig,
+				() => ({
+					modelLabel: state.modelLabel,
+					providerLabel: state.providerLabel,
+					contextLabel: state.contextLabel,
+					contextUsedLabel: state.contextUsedLabel,
+					tokenLabel: state.tokenLabel,
+				}),
+				getThinkingLevel,
+			)) as ZentuiEditorFactory;
+		factory[ZENTUI_EDITOR_FACTORY] = true;
+		return factory;
+	};
+
+	const makeWrappedEditorFactory = (
+		ctx: ExtensionContext,
+		baseFactory: EditorFactory,
+	): ZentuiEditorFactory => {
+		const sessionTheme = ctx.ui.theme;
+		const factory = ((tui: TUI, theme: EditorTheme, keybindings: KeybindingsManager) =>
+			new WrappedPolishedEditor(
+				baseFactory(tui, theme, keybindings),
+				sessionTheme,
+				getCurrentConfig,
+				() => ({
+					modelLabel: state.modelLabel,
+					providerLabel: state.providerLabel,
+					contextLabel: state.contextLabel,
+					contextUsedLabel: state.contextUsedLabel,
+					tokenLabel: state.tokenLabel,
+				}),
+				getThinkingLevel,
+			)) as ZentuiEditorFactory;
+		factory[ZENTUI_EDITOR_FACTORY] = true;
+		factory[ZENTUI_EDITOR_BASE_FACTORY] = baseFactory;
+		return factory;
+	};
+
+	const installEditor = (ctx: ExtensionContext): boolean => {
+		const currentFactory = ctx.ui.getEditorComponent();
+		if (currentFactory && currentFactory === installedEditorFactory) {
+			editorInstalled = true;
+			return true;
+		}
+
+		installPrototypePatches();
+		const currentZentuiBaseFactory = getZentuiEditorBaseFactory(currentFactory);
+		if (currentFactory && isZentuiEditorFactory(currentFactory)) {
+			wrappedEditorFactory = currentZentuiBaseFactory;
+			const nextFactory = currentZentuiBaseFactory
+				? makeWrappedEditorFactory(ctx, currentZentuiBaseFactory)
+				: makeEditorFactory(ctx);
+			ctx.ui.setEditorComponent(nextFactory);
+			installedEditorFactory = nextFactory;
+			editorInstallMode = currentZentuiBaseFactory ? "wrapper" : "standalone";
+		} else if (currentFactory) {
+			wrappedEditorFactory = currentFactory;
+			const nextFactory = makeWrappedEditorFactory(ctx, currentFactory);
+			ctx.ui.setEditorComponent(nextFactory);
+			installedEditorFactory = nextFactory;
+			editorInstallMode = "wrapper";
+		} else {
+			wrappedEditorFactory = undefined;
+			const nextFactory = makeEditorFactory(ctx);
+			ctx.ui.setEditorComponent(nextFactory);
+			installedEditorFactory = nextFactory;
+			editorInstallMode = "standalone";
+		}
+		editorInstalled = true;
+		return true;
+	};
+
+	const uninstallEditor = (ctx: ExtensionContext): boolean => {
+		const currentFactory = ctx.ui.getEditorComponent();
+		if (currentFactory && !isZentuiEditorFactory(currentFactory)) return false;
+
+		uninstallPrototypePatches();
+		ctx.ui.setEditorComponent(
+			editorInstallMode === "wrapper" && wrappedEditorFactory ? wrappedEditorFactory : undefined,
+		);
+		wrappedEditorFactory = undefined;
+		installedEditorFactory = undefined;
+		editorInstallMode = "none";
+		editorInstalled = false;
+		return true;
+	};
+
+	const installStatusLine = (ctx: ExtensionContext) => {
+		if (footerInstalled) return;
+		installFooter(ctx, state, getCurrentConfig, {
+			setRequestRender: (fn) => {
+				requestFooterRender = fn;
+			},
+			scheduleProjectRefresh,
+			setExtensionStatusesGetter(fn) {
+				getActiveExtensionStatuses = fn ?? (() => new Map());
+			},
+			getLiveContext: () => liveContext.get(),
+		});
+		footerInstalled = true;
+		stopProjectRefresh();
+		stopRefreshInterval = startProjectRefreshInterval(currentConfig.projectRefreshIntervalMs, () =>
+			scheduleProjectRefresh(ctx),
+		);
+		scheduleProjectRefresh(ctx, { force: true });
+		refresh();
+		startSessionTimer();
+	};
+
+	const uninstallStatusLine = (ctx: ExtensionContext) => {
+		stopSessionTimer();
+		stopProjectRefresh();
+		ctx.ui.setFooter(undefined);
+		footerInstalled = false;
+		requestFooterRender = undefined;
+		getActiveExtensionStatuses = () => new Map();
+	};
+
+	const applyConfiguredUi = (ctx: ExtensionContext): ApplyUiResult => {
+		const result: ApplyUiResult = { editorBlocked: false };
+		if (!isTuiContext(ctx)) return result;
+		activeTheme = ctx.ui.theme;
+		if (currentConfig.features.editor) {
+			const currentFactory = ctx.ui.getEditorComponent();
+			const editorMissingOrReplaced = !editorInstalled || !isZentuiEditorFactory(currentFactory);
+			if (editorMissingOrReplaced) result.editorBlocked = !installEditor(ctx);
+		} else if (editorInstalled || prototypePatchesInstalled) {
+			result.editorBlocked = !uninstallEditor(ctx);
+		}
+
+		if (currentConfig.features.statusLine) {
+			installStatusLine(ctx);
+		} else if (footerInstalled) {
+			uninstallStatusLine(ctx);
+		}
+		return result;
+	};
+
+	const installUi = (ctx: ExtensionContext) => {
+		if (!isTuiContext(ctx)) return;
+		activeTheme = ctx.ui.theme;
+		uninstallPrototypePatches();
+		footerInstalled = false;
+		editorInstalled = false;
+		installedEditorFactory = undefined;
+		ensureConfigExists();
+		currentConfig = loadConfig();
+		syncFooterState(ctx);
+		stopProjectRefresh();
+		applyConfiguredUi(ctx);
+		if (currentConfig.fixedEditor?.enabled) {
+			installFixedEditorProbe(ctx, getCurrentConfig, sessionLifecycle);
+		}
+		refresh();
+	};
+
+	const scheduleEditorReconciliation = (ctx: ExtensionContext) => {
+		sessionLifecycle.defer(() => {
+			if (!isTuiContext(ctx) || !currentConfig.features.editor) return;
+			const currentFactory = ctx.ui.getEditorComponent();
+			if (currentFactory && currentFactory !== installedEditorFactory) {
+				applyConfiguredUi(ctx);
+				refresh();
+			}
+		});
+	};
+
+	const cleanupUi = (ctx?: ExtensionContext) => {
+		if (!ctx || !sessionLifecycle.isCurrent()) return;
+		sessionLifecycle.shutdown();
+		try {
+			disposeFixedEditor(ctx);
+			if (isTuiContext(ctx)) removeFixedEditorProbe(ctx);
+			uninstallPrototypePatches();
+			stopSessionTimer();
+			stopProjectRefresh();
+			requestFooterRender = undefined;
+			getActiveExtensionStatuses = () => new Map();
+			if (isTuiContext(ctx)) {
+				ctx.ui.setFooter(undefined);
+				const currentFactory = ctx.ui.getEditorComponent();
+				if (!currentFactory || isZentuiEditorFactory(currentFactory)) {
+					ctx.ui.setEditorComponent(
+						getZentuiEditorBaseFactory(currentFactory) ??
+							(editorInstallMode === "wrapper" && wrappedEditorFactory
+								? wrappedEditorFactory
+								: undefined),
+					);
+				}
+			}
+			wrappedEditorFactory = undefined;
+			installedEditorFactory = undefined;
+			editorInstallMode = "none";
+			footerInstalled = false;
+			editorInstalled = false;
+			activeTheme = undefined;
+		} finally {
+			requestFooterRender = undefined;
+		}
+	};
+
+	const syncInteractiveState = (_event: unknown, ctx: ExtensionContext) => {
+		refreshInteractiveState(ctx);
+	};
+	const syncInteractiveAndProjectState = (_event: unknown, ctx: ExtensionContext) => {
+		refreshInteractiveState(ctx, true);
+	};
+
+	pi.on("session_start", async (_event, ctx) => {
+		sessionLifecycle.start();
+		liveContext.clear();
+		state.sessionStartEpoch = Date.now();
+		invalidateUsageTotalsCache();
+		lastProjectCwd = undefined;
+		installUi(ctx);
+		scheduleEditorReconciliation(ctx);
+	});
+
+	registerZentuiSettingsCommand(pi, {
+		sessionLifecycle,
+		getConfig: getCurrentConfig,
+		setColorSources(patch: Partial<ColorSourcesConfig>) {
+			currentConfig = saveColorSourcesPatch(patch);
+		},
+		setUiFeatures(patch: Partial<UiFeaturesConfig>, ctx: ExtensionContext) {
+			currentConfig = saveUiFeaturesPatch(patch);
+			const result = applyConfiguredUi(ctx);
+			return {
+				applied: !(patch.editor !== undefined && result.editorBlocked),
+				reason: result.editorBlocked
+					? "another extension is currently managing the editor; reload Pi to apply this change"
+					: undefined,
+			};
+		},
+		setFooterSegments(patch: Partial<FooterSegmentsConfig>) {
+			currentConfig = saveFooterSegmentsPatch(patch);
+			startSessionTimer();
+		},
+		setFooterFormat(value: string) {
+			currentConfig = saveFooterFormatPatch(value);
+			startSessionTimer();
+		},
+		setIconMode(mode: IconMode) {
+			currentConfig = saveIconsModePatch(mode);
+		},
+		setContextStyle(style: ContextStyle) {
+			currentConfig = saveContextStylePatch(style);
+		},
+		setSeparator(separator: SeparatorStyle) {
+			currentConfig = saveSeparatorPatch(separator);
+		},
+		setPathDisplay(patch: Partial<PathDisplayConfig>) {
+			currentConfig = savePathDisplayPatch(patch);
+		},
+		setGitBranch(patch: Partial<GitBranchConfig>) {
+			currentConfig = saveGitBranchPatch(patch);
+		},
+		getActiveExtensionStatuses() {
+			return getActiveExtensionStatuses();
+		},
+		setExtensionStatusPlacement(key: string, placement: ExtensionStatusPlacement) {
+			currentConfig = saveExtensionStatusPlacement(key, placement);
+		},
+		setExtensionStatusColorMode(key: string, colorMode: ExtensionStatusColorMode) {
+			currentConfig = saveExtensionStatusColorMode(key, colorMode);
+		},
+		setFixedEditor(patch: Partial<FixedEditorConfig>, ctx: ExtensionContext) {
+			currentConfig = saveFixedEditorPatch(patch);
+			if (patch.enabled === true) {
+				installFixedEditorProbe(ctx, getCurrentConfig, sessionLifecycle);
+			} else if (patch.enabled === false) {
+				disposeFixedEditor(ctx);
+			}
+			refresh();
+		},
+		requestRender() {
+			refresh();
+		},
+	});
+
+	pi.on("session_shutdown", async (_event, ctx) => {
+		liveContext.clear();
+		cleanupUi(ctx);
+	});
+
+	const syncInteractiveAndProjectStateWithUsage = (_event: unknown, ctx: ExtensionContext) => {
+		invalidateUsageTotalsCache();
+		refreshInteractiveState(ctx, true);
+	};
+
+	pi.on("agent_start", (event, ctx) => {
+		liveContext.clear();
+		syncInteractiveState(event, ctx);
+	});
+	pi.on("agent_end", (event, ctx) => {
+		liveContext.clear();
+		syncInteractiveAndProjectState(event, ctx);
+	});
+	pi.on("model_select", (event, ctx) => {
+		liveContext.clear();
+		syncInteractiveState(event, ctx);
+	});
+	pi.on("thinking_level_select", syncInteractiveState);
+	pi.on("message_update", (event) => {
+		liveContext.update(event.message);
+	});
+	pi.on("message_end", (event, ctx) => {
+		// Pi notifies extensions before persisting a successful message, so retain its live
+		// context until agent_end; failed messages clear immediately instead of showing stale usage.
+		if (
+			event.message.role === "assistant" &&
+			(event.message.stopReason === "error" || event.message.stopReason === "aborted")
+		) {
+			liveContext.clear();
+		}
+		syncInteractiveAndProjectStateWithUsage(event, ctx);
+	});
+	pi.on("tool_execution_start", (event, ctx) => {
+		liveContext.clear();
+		syncInteractiveState(event, ctx);
+	});
+	pi.on("tool_execution_end", syncInteractiveAndProjectState);
+	pi.on("session_compact", (event, ctx) => {
+		liveContext.clear();
+		syncInteractiveAndProjectStateWithUsage(event, ctx);
+	});
+	pi.on("session_tree", (event, ctx) => {
+		liveContext.clear();
+		syncInteractiveAndProjectStateWithUsage(event, ctx);
+	});
+}

package/extensions/zentui/live-context.ts

@@ -0,0 +1,75 @@
+import type { AssistantMessage, Usage } from "@earendil-works/pi-ai";
+import type { SessionLifecycle } from "./session-lifecycle.js";
+
+export type LiveContextOverride = {
+	tokens: number;
+};
+
+function usageComponent(value: unknown): number {
+	return typeof value === "number" && Number.isFinite(value) && value > 0 ? value : 0;
+}
+
+export function calculateLiveContextTokens(usage: Usage | undefined): number | undefined {
+	if (!usage) return undefined;
+	const totalTokens = usageComponent(usage.totalTokens);
+	if (totalTokens > 0) return totalTokens;
+	const calculated =
+		usageComponent(usage.input) +
+		usageComponent(usage.output) +
+		usageComponent(usage.cacheRead) +
+		usageComponent(usage.cacheWrite);
+	return calculated > 0 ? calculated : undefined;
+}
+
+export function liveContextFromMessage(message: unknown): LiveContextOverride | undefined {
+	if (!message || typeof message !== "object") return undefined;
+	const assistant = message as Partial<AssistantMessage>;
+	if (assistant.role !== "assistant") return undefined;
+	if (assistant.stopReason === "error" || assistant.stopReason === "aborted") return undefined;
+	const tokens = calculateLiveContextTokens(assistant.usage);
+	return tokens === undefined ? undefined : { tokens };
+}
+
+export class LiveContextController {
+	private readonly lifecycle: SessionLifecycle;
+	private readonly requestRender: () => void;
+	private override: LiveContextOverride | undefined;
+	private cancelScheduledRender: (() => void) | undefined;
+	private scheduledGeneration: number | undefined;
+
+	constructor(lifecycle: SessionLifecycle, requestRender: () => void) {
+		this.lifecycle = lifecycle;
+		this.requestRender = requestRender;
+	}
+
+	get(): LiveContextOverride | undefined {
+		return this.override;
+	}
+
+	update(message: unknown): boolean {
+		const next = liveContextFromMessage(message);
+		if (!next || !this.lifecycle.isCurrent()) return false;
+		this.override = next;
+		const generation = this.lifecycle.currentGeneration();
+		if (this.cancelScheduledRender && this.scheduledGeneration !== generation) {
+			this.cancelScheduledRender = undefined;
+			this.scheduledGeneration = undefined;
+		}
+		if (!this.cancelScheduledRender) {
+			this.scheduledGeneration = generation;
+			this.cancelScheduledRender = this.lifecycle.defer(() => {
+				this.cancelScheduledRender = undefined;
+				this.scheduledGeneration = undefined;
+				if (this.override) this.requestRender();
+			}, 250);
+		}
+		return true;
+	}
+
+	clear(): void {
+		this.override = undefined;
+		this.cancelScheduledRender?.();
+		this.cancelScheduledRender = undefined;
+		this.scheduledGeneration = undefined;
+	}
+}