@bachi/pi-coder 1.0.0

package/extensions/startup-logo/header-guard.test.ts

@@ -0,0 +1,274 @@
+import assert from "node:assert/strict";
+import { test } from "node:test";
+import {
+	HEADER_GUARD_KEY,
+	findRenderContainer,
+	freezeHeaderContainer,
+	releaseHeaderGuard,
+	type RenderableContainer,
+} from "./header-guard.ts";
+
+/**
+ * 仿 pi-tui 的 `Container`：`clear()` 会**换掉**整个 children 数组
+ * （`this.children = []`），所以守卫必须每次重读该字段而不是缓存它。
+ */
+class FakeContainer implements RenderableContainer {
+	children: unknown[] = [];
+	/** 原 render 被调用的次数（断言「内置 header 一帧都没出」用）。 */
+	originalRenders = 0;
+
+	addChild(component: unknown): void {
+		this.children.push(component);
+	}
+
+	clear(): void {
+		this.children = [];
+	}
+
+	render(width: number): string[] {
+		this.originalRenders += 1;
+		return this.children.flatMap((child) => (child as Leaf).render(width));
+	}
+}
+
+/** 仿叶子组件：行里带上自己的名字和宽度，方便看出行是谁、按哪个宽度渲染的。 */
+interface Leaf {
+	render(width: number): string[];
+}
+
+function leaf(name: string, lines = 1): Leaf {
+	return { render: (width: number) => Array.from({ length: lines }, () => `${name}#${width}`) };
+}
+
+/** 仿 `tui`：只有 children，没有 render（findRenderContainer 只认 children）。 */
+function tuiRoot(...children: unknown[]): { children: unknown[] } {
+	return { children };
+}
+
+test("findRenderContainer 按对象身份认出直接持有它的容器", () => {
+	const ours = leaf("startup-logo");
+	const inner = new FakeContainer();
+	inner.addChild(ours);
+	const outer = new FakeContainer();
+	outer.addChild(inner);
+	const root = tuiRoot(new FakeContainer(), outer);
+
+	assert.equal(findRenderContainer(root, ours), inner);
+});
+
+test("findRenderContainer 找不到就返回 undefined，不猜位置", () => {
+	const inner = new FakeContainer();
+	inner.addChild(leaf("other"));
+	const root = tuiRoot(inner);
+
+	assert.equal(findRenderContainer(root, leaf("missing")), undefined);
+	assert.equal(findRenderContainer(root, undefined), undefined);
+	assert.equal(findRenderContainer(undefined, leaf("x")), undefined);
+	assert.equal(findRenderContainer({ notAContainer: true }, leaf("x")), undefined);
+});
+
+test("findRenderContainer 容忍环状引用与只读 children", () => {
+	const a = new FakeContainer();
+	const b = new FakeContainer();
+	a.addChild(b);
+	b.addChild(a);
+
+	assert.equal(findRenderContainer(a, leaf("nope")), undefined);
+
+	const frozenRoot = tuiRoot({ children: Object.freeze([leaf("x")]) as unknown[] });
+	assert.equal(findRenderContainer(frozenRoot, leaf("y")), undefined);
+});
+
+test("我们的组件在位时原样透传，被换成内置 header 时出重放行", () => {
+	const ours = leaf("startup-logo");
+	const builtIn = leaf("pi-default-header");
+	const container = new FakeContainer();
+	container.addChild(ours);
+
+	const release = freezeHeaderContainer({
+		container,
+		ownComponent: ours,
+		frozenRender: (width) => [`frozen#${width}`],
+	});
+	assert.equal(typeof release, "function");
+
+	// 正常会话：真实渲染照旧。
+	assert.deepEqual(container.render(80), ["startup-logo#80"]);
+	assert.equal(container.originalRenders, 1);
+
+	// 换会话窗口：pi 已经 clear() 并挂上内置 header。
+	container.clear();
+	container.addChild(builtIn);
+	assert.deepEqual(container.render(80), ["frozen#80"]);
+	assert.deepEqual(container.render(40), ["frozen#40"]);
+	// 内置 header 一次都没被渲染 → 屏幕上不会出现它。
+	assert.equal(container.originalRenders, 1);
+});
+
+test("重放行按当前宽度重新截断", () => {
+	const ours = leaf("startup-logo");
+	const container = new FakeContainer();
+	container.addChild(ours);
+	const wide = " ".repeat(30);
+	freezeHeaderContainer({
+		container,
+		ownComponent: ours,
+		frozenRender: (width) => [wide.slice(0, width)],
+	});
+
+	container.clear();
+	container.addChild(leaf("pi-default-header"));
+	assert.deepEqual(container.render(10), [wide.slice(0, 10)]);
+});
+
+test("解除接管后恢复原渲染，且不留 own render 属性与接管记录", () => {
+	const ours = leaf("startup-logo");
+	const builtIn = leaf("pi-default-header");
+	const container = new FakeContainer();
+	container.addChild(ours);
+
+	const release = freezeHeaderContainer({
+		container,
+		ownComponent: ours,
+		frozenRender: () => ["frozen"],
+	})!;
+	container.clear();
+	container.addChild(builtIn);
+	assert.deepEqual(container.render(10), ["frozen"]);
+
+	release();
+	assert.equal(Object.prototype.hasOwnProperty.call(container, "render"), false);
+	assert.equal((container as { [HEADER_GUARD_KEY]?: unknown })[HEADER_GUARD_KEY], undefined);
+	assert.deepEqual(container.render(10), ["pi-default-header#10"]);
+
+	// 幂等：重复解除不再改动任何东西。
+	release();
+	assert.equal(Object.prototype.hasOwnProperty.call(container, "render"), false);
+});
+
+test("releaseHeaderGuard 只在确有接管时报告成功", () => {
+	const container = new FakeContainer();
+	assert.equal(releaseHeaderGuard(container), false);
+	assert.equal(releaseHeaderGuard(undefined), false);
+	const release = freezeHeaderContainer({ container, ownComponent: {}, frozenRender: () => ["frozen"] });
+	assert.equal(releaseHeaderGuard(container), true);
+	// 已被解除 → 再解除报 false。
+	release?.();
+	assert.equal(releaseHeaderGuard(container), false);
+});
+
+test("对同一容器重复接管会先交还上一次，不叠两层", () => {
+	const first = leaf("first");
+	const second = leaf("second");
+	const container = new FakeContainer();
+	container.addChild(first);
+	container.addChild(second);
+
+	freezeHeaderContainer({ container, ownComponent: first, frozenRender: () => ["frozen-first"] });
+	freezeHeaderContainer({ container, ownComponent: second, frozenRender: () => ["frozen-second"] });
+
+	// first 在位、second 不在位：若两层的 wrapper 还叠着，这里会出 frozen-first；
+	// 只保留最后一次接管 → 判定 second 不在位 → 出 frozen-second。
+	container.clear();
+	container.addChild(first);
+	assert.deepEqual(container.render(7), ["frozen-second"]);
+
+	// second 回到位 → 透传。
+	container.addChild(second);
+	assert.deepEqual(container.render(7), ["first#7", "second#7"]);
+});
+
+test("超时后自愈：组件在位与否都交还原渲染", () => {
+	const ours = leaf("startup-logo");
+	const builtIn = leaf("pi-default-header");
+	const container = new FakeContainer();
+	container.addChild(ours);
+	let clock = 1000;
+	const release = freezeHeaderContainer({
+		container,
+		ownComponent: ours,
+		frozenRender: () => ["frozen"],
+		maxAgeMs: 500,
+		now: () => clock,
+	})!;
+
+	// 组件仍在新位置（正常交接后的样子）+ 已超时 → 透传并交还。
+	clock = 1600;
+	assert.deepEqual(container.render(9), ["startup-logo#9"]);
+	assert.equal(Object.prototype.hasOwnProperty.call(container, "render"), false);
+
+	// 另一侧：新 header 一直没来（换会话失败）→ 超时后回到 pi 默认行为。
+	const ours2 = leaf("startup-logo2");
+	const container2 = new FakeContainer();
+	container2.addChild(ours2);
+	let clock2 = 0;
+	freezeHeaderContainer({
+		container: container2,
+		ownComponent: ours2,
+		frozenRender: () => ["frozen"],
+		maxAgeMs: 500,
+		now: () => clock2,
+	});
+	container2.clear();
+	container2.addChild(builtIn);
+	clock2 = 501;
+	assert.deepEqual(container2.render(9), ["pi-default-header#9"]);
+	release();
+});
+
+test("未超时期间无论多少帧都只出重放行", () => {
+	const ours = leaf("startup-logo");
+	const container = new FakeContainer();
+	container.addChild(ours);
+	freezeHeaderContainer({
+		container,
+		ownComponent: ours,
+		frozenRender: (width) => [`frozen#${width}`, `frozen#${width}`],
+	});
+	container.clear();
+	container.addChild(leaf("pi-default-header"));
+
+	for (let frame = 0; frame < 5; frame += 1) {
+		assert.deepEqual(container.render(80), ["frozen#80", "frozen#80"]);
+	}
+});
+
+test("容器形状不对时什么都不做（非 tui 模式 / pi 改了结构）", () => {
+	const noop = () => ["frozen"];
+	assert.equal(freezeHeaderContainer({ container: undefined, ownComponent: {}, frozenRender: noop }), undefined);
+	assert.equal(freezeHeaderContainer({ container: {}, ownComponent: {}, frozenRender: noop }), undefined);
+	assert.equal(
+		freezeHeaderContainer({ container: { children: [] }, ownComponent: {}, frozenRender: noop }),
+		undefined,
+	);
+	assert.equal(
+		freezeHeaderContainer({
+			container: { children: "not-an-array", render: noop } as unknown as RenderableContainer,
+			ownComponent: {},
+			frozenRender: noop,
+		}),
+		undefined,
+	);
+	// ownComponent 为 undefined（header 还没装上）也不接管。
+	assert.equal(
+		freezeHeaderContainer({ container: new FakeContainer(), ownComponent: undefined, frozenRender: noop }),
+		undefined,
+	);
+});
+
+test("frozenRender 抛错时回落到原渲染，不把异常漏给宿主", () => {
+	const ours = leaf("startup-logo");
+	const container = new FakeContainer();
+	container.addChild(ours);
+	freezeHeaderContainer({
+		container,
+		ownComponent: ours,
+		frozenRender: () => {
+			throw new Error("boom");
+		},
+	});
+
+	container.clear();
+	container.addChild(leaf("pi-default-header"));
+	assert.deepEqual(container.render(12), ["pi-default-header#12"]);
+});

package/extensions/startup-logo/header-guard.ts

@@ -0,0 +1,166 @@
+/**
+ * header-guard.ts — 会话替换期间的 header「冻结」
+ *
+ * 问题和 statusline 的 footer 冻结同源：pi 每次换会话（`/new`、`/clear`、`/resume`、fork、
+ * rewind）都会在 `session_shutdown` 之后、新会话 `session_start` 之前调 `resetExtensionUI()`，
+ * 它无条件 `setExtensionHeader(undefined)` 把 header 还原成内置那只（`interactive-mode.js`）。
+ * 中间要新建会话对象、重绑扩展，几十毫秒起步，而 pi-tui 的渲染节流只有 16ms，所以这段窗口
+ * 至少会出几帧 —— 装了 logo header 的话，顶部会「闪」回 `pi vX.Y.Z` + 提示行，而且内置 header
+ * 比 logo 矮，还会带一次整块布局跳动。
+ *
+ * 修法与 footer 版一致：扩展侧没有更早的钩子（新实例要等重绑才拿到 ctx），所以把保证落在
+ * **出帧那一刻** —— 接管 header 容器的 `render`，发现里面装的已经不是我们的组件时，返回上一帧
+ * 的行，内置 header 一帧都不出现。
+ *
+ * 约束（与 statusline/footer-guard.ts 相同）：
+ *   - 只替换容器的 `render`，不改 `children`、不猜下标、不碰 pi 内部状态；
+ *   - 容器按**对象身份**在 `tui.children` 子树里认（`findRenderContainer`），认不出就什么都不做；
+ *   - 任何一步抛错都当作「没冻住」，绝不把异常漏给宿主；
+ *   - 接管记录挂在容器上的 `Symbol.for(...)` 键，`/reload` 后新实例能看到并解除旧实例的接管；
+ *     另有 `maxAgeMs` 兜底，交接万一失败也只冻一小会儿。
+ *
+ * 与 footer 版各持一份、各用各的符号键：两个扩展互不依赖，也不会互相顶掉接管。
+ */
+
+/** 挂在 header 容器上的接管记录键（全局符号注册表：跨扩展实例可见）。 */
+export const HEADER_GUARD_KEY: symbol = Symbol.for("litellm-any.pi-startup-logo.headerGuard");
+
+/** 冻结的最长时长：正常情况下新会话的 `session_start` 会远早于它解除。 */
+export const HEADER_GUARD_MAX_AGE_MS = 3000;
+
+/** pi-tui `Container` 里本模块用到的那部分结构（`children` 是公开字段）。 */
+export interface RenderableContainer {
+	children: unknown[];
+	render(width: number): string[];
+}
+
+interface GuardRecord {
+	release: () => void;
+}
+
+/**
+ * 在 `root` 的子树里按对象身份找「直接装着 `component` 的那个容器」。
+ * 只认对象身份（同形状的组件到处都是），带 seen 集与深度上限，任何异常都当作没找到。
+ */
+export function findRenderContainer(
+	root: unknown,
+	component: unknown,
+	maxDepth = 6,
+): RenderableContainer | undefined {
+	if (component === undefined || component === null) return undefined;
+	const seen = new Set<unknown>();
+
+	const walk = (node: unknown, depth: number): RenderableContainer | undefined => {
+		if (depth > maxDepth || node === null || typeof node !== "object" || seen.has(node)) {
+			return undefined;
+		}
+		seen.add(node);
+		const children = (node as { children?: unknown }).children;
+		if (!Array.isArray(children)) return undefined;
+		if (children.includes(component)) return node as RenderableContainer;
+		for (const child of children) {
+			const found = walk(child, depth + 1);
+			if (found) return found;
+		}
+		return undefined;
+	};
+
+	try {
+		return walk(root, 0);
+	} catch {
+		return undefined;
+	}
+}
+
+/** 解除 `container` 上现有的接管（可能是别的扩展实例遗留的）。返回是否解掉了什么。 */
+export function releaseHeaderGuard(container: unknown): boolean {
+	try {
+		const record = (container as { [HEADER_GUARD_KEY]?: GuardRecord } | undefined)?.[HEADER_GUARD_KEY];
+		if (!record || typeof record.release !== "function") return false;
+		record.release();
+		return true;
+	} catch {
+		return false;
+	}
+}
+
+/**
+ * 接管 header 容器的渲染，返回解除函数（幂等）。
+ *
+ * - `ownComponent` 还在容器里 → 原样透传（并在超时后顺手交还，自愈）；
+ * - 容器被 pi 换成内置 header → 返回 `frozenRender(width)` 的行；
+ * - 超过 `maxAgeMs` 仍等不到新 header → 交还，回到 pi 默认行为。
+ */
+export function freezeHeaderContainer(options: {
+	container: unknown;
+	ownComponent: unknown;
+	frozenRender: (width: number) => string[];
+	maxAgeMs?: number;
+	now?: () => number;
+}): (() => void) | undefined {
+	const container = options.container as RenderableContainer | undefined;
+	if (
+		!container ||
+		typeof container.render !== "function" ||
+		!Array.isArray(container.children) ||
+		options.ownComponent === undefined
+	) {
+		return undefined;
+	}
+	// 先解除已有接管（自己重入 / 上一个扩展实例遗留），避免叠两层 wrapper。
+	releaseHeaderGuard(container);
+
+	const ownComponent = options.ownComponent;
+	const maxAgeMs = options.maxAgeMs ?? HEADER_GUARD_MAX_AGE_MS;
+	const now = options.now ?? Date.now;
+	const armedAt = now();
+	const hadOwnRender = Object.prototype.hasOwnProperty.call(container, "render");
+	const originalRender = container.render;
+	const state: { released: boolean; record?: GuardRecord } = { released: false };
+
+	const release = (): void => {
+		if (state.released) return;
+		state.released = true;
+		try {
+			const holder = container as { [HEADER_GUARD_KEY]?: GuardRecord };
+			if (state.record && holder[HEADER_GUARD_KEY] === state.record) delete holder[HEADER_GUARD_KEY];
+			if (hadOwnRender) container.render = originalRender;
+			else delete (container as { render?: unknown }).render;
+		} catch {
+			// 容器只读 / 已被替换：wrapper 自己会因 released 标志透传，不再冻。
+		}
+	};
+
+	const wrapper = (width: number): string[] => {
+		if (state.released) return originalRender.call(container, width);
+		const expired = now() - armedAt > maxAgeMs;
+		let mounted = false;
+		try {
+			mounted = Array.isArray(container.children) && container.children.includes(ownComponent);
+		} catch {
+			mounted = true; // 读 children 都抛：别硬撑，交还给原渲染。
+		}
+		if (mounted) {
+			if (expired) release();
+			return originalRender.call(container, width);
+		}
+		if (expired) {
+			release();
+			return originalRender.call(container, width);
+		}
+		try {
+			return options.frozenRender(width);
+		} catch {
+			return originalRender.call(container, width);
+		}
+	};
+
+	try {
+		container.render = wrapper as RenderableContainer["render"];
+		state.record = { release };
+		(container as { [HEADER_GUARD_KEY]?: GuardRecord })[HEADER_GUARD_KEY] = state.record;
+	} catch {
+		return undefined;
+	}
+	return release;
+}