shahondin1624/pi-extensions
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Refactor pass: shared utils, non-blocking discovery, safer monkey-patches

Browse Source 
Audit produced five concrete improvements:

1) New shared/ module (zero-dep pure utilities)
   - shared/ansi.ts: hexToRgb (throws on malformed input instead of
     silently producing NaN), fgFromHex, stripAnsi, visibleWidth,
     ANSI_RESET_FG / ANSI_RESET_ALL constants.
   - shared/format.ts: formatTokens, formatElapsed.
   - shared/ctx.ts: safely() and safelyAsync() helpers for dealing with
     pi's "stale after session replacement or reload" ExtensionRunner
     semantics.

   Removes duplicate helpers from mechanicus-footer, markdown-body-color,
   dark-mechanicus-indicator.

2) ai-server: non-blocking startup + short-race timeout
   - Factory registers STATIC_MODELS immediately so pi startup isn't
     blocked on the HTTPS round-trip.
   - Races discoverModels() against a 300ms timeout. On LAN (~40ms) the
     live list wins and pi --list-models sees the real models. Past the
     timeout, fallback remains and background discovery updates the
     provider later.
   - listModelsCached() with 5s TTL for tab completions (was firing a
     round-trip on every keystroke).
   - loadModel/unloadModel invalidate the cache.

3) dark-mechanicus-indicator: stale-ctx guard
   - Wrap the setInterval ticker body in safely() so a race between
     session_shutdown and the ticker can't crash node. Same pattern as
     the earlier footer fix.

4) Safer monkey-patches in markdown-body-color and mechanicus-thinking-label
   - Feature-detect Markdown/Editor/AssistantMessageComponent's target
     method before patching. Warn-and-skip rather than silently create
     a broken prototype if a pi-tui upgrade renames the internal method.

5) Minor
   - Replaced five `as any` casts with typed Record<string, unknown>
     access in the monkey-patch sites.
   - ai-server debug log only fires when actual discovery succeeds.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
shahondin1624
2026-04-24 00:05:18 +02:00
parent cd94a6636d
commit f1ceeb4363
6 changed files with 120 additions and 74 deletions
Download Patch File Download Diff File Expand all files Collapse all files
ai-server/admin.ts
+26 -2
View File
@@ -87,13 +87,37 @@ export async function listModels(): Promise<RouterModel[]> {
return (data?.data ?? []) as RouterModel[];
}
// Short TTL cache for listModels — tab-completion calls the completer on
// every Tab press, which would otherwise fire an HTTPS round-trip each
// time. Five seconds is long enough to dedupe back-to-back completions
// but short enough that a /ai-server-load still sees near-fresh state.
const LIST_MODELS_TTL_MS = 5_000;
let cachedList: { at: number; models: RouterModel[] } | null = null;
export async function listModelsCached(): Promise<RouterModel[]> {
if (cachedList && Date.now() - cachedList.at < LIST_MODELS_TTL_MS) {
return cachedList.models;
}
const models = await listModels();
cachedList = { at: Date.now(), models };
return models;
}
export function invalidateListModelsCache(): void {
cachedList = null;
}
export async function loadModel(id: string): Promise<unknown> {
// The router's handler reads `body["model"]`; passing `{id}` yields a 404.
return routerRequest("POST", "/models/load", { model: id });
const r = await routerRequest("POST", "/models/load", { model: id });
invalidateListModelsCache();
return r;
}
export async function unloadModel(id: string): Promise<unknown> {
return routerRequest("POST", "/models/unload", { model: id });
const r = await routerRequest("POST", "/models/unload", { model: id });
invalidateListModelsCache();
return r;
}
// A preset is "runnable" only if it has a --model path. Placeholder sections
ai-server/index.ts
+49 -24
View File
@@ -3,6 +3,7 @@ import {
discoverModels,
extractCtxSize,
listModels,
listModelsCached,
loadModel,
readPreset,
reloadOneModel,
@@ -22,7 +23,10 @@ import { streamAiServer } from "./stream.js";
async function completeModelId(prefix: string) {
try {
const models = await listModels();
// Cached for 5s. Tab-completion calls the completer on every keystroke,
// but the user typically only types one model id per command — caching
// deduplicates the network round-trip without stale-state harm.
const models = await listModelsCached();
const hits = models
.filter((m) => m.id.startsWith(prefix))
.map((m) => ({ value: m.id, label: m.id }));
@@ -53,23 +57,51 @@ function registerProviderWithModels(
});
}
export default async function (pi: ExtensionAPI) {
// Discover models from the router. Fall back to the static list if the
// server is unreachable so pi can still start up.
let models: ServerModel[] = [];
let source = "discovery";
try {
models = await discoverModels();
if (models.length === 0) {
models = STATIC_MODELS;
source = "fallback (empty discovery)";
}
} catch (err) {
models = STATIC_MODELS;
source = `fallback (${(err as Error).message})`;
}
const DISCOVERY_FAST_TIMEOUT_MS = 300;
registerProviderWithModels(pi, models);
export default async function (pi: ExtensionAPI) {
// Register the provider IMMEDIATELY with the static fallback list so pi
// startup isn't blocked on the HTTPS round-trip in the worst case.
registerProviderWithModels(pi, STATIC_MODELS);
// Then race real discovery against a short timeout. On LAN the router
// answers in ~40ms and pi --list-models sees the live list. On slow
// networks we bail at 300ms and the fallback is what the user sees; the
// background promise keeps running and re-registers later.
const discovery = discoverModels().catch((err) => {
if (process.env.PI_DEBUG) {
console.log(
`[ai-server] Discovery failed (${(err as Error).message}); fallback remains`,
);
}
return null;
});
const timeout = new Promise<null>((r) =>
setTimeout(() => r(null), DISCOVERY_FAST_TIMEOUT_MS),
);
const fastResult = await Promise.race([discovery, timeout]);
if (fastResult && fastResult.length > 0) {
registerProviderWithModels(pi, fastResult);
if (process.env.PI_DEBUG) {
console.log(
`[ai-server] Discovered ${fastResult.length} model(s) on ${AI_SERVER_URL}: ${fastResult.map((m) => m.id).join(", ")}`,
);
}
} else {
// Slow network or discovery still pending — keep waiting in the
// background and update the provider once it arrives.
discovery.then((models) => {
if (models && models.length > 0) {
registerProviderWithModels(pi, models);
if (process.env.PI_DEBUG) {
console.log(
`[ai-server] Late discovery: ${models.length} model(s)`,
);
}
}
});
}
// ─── Admin commands ──────────────────────────────────────────────────
@@ -241,11 +273,4 @@ export default async function (pi: ExtensionAPI) {
},
});
if (process.env.PI_DEBUG) {
console.log(
`[ai-server] Registered ${models.length} model(s) via ${source} on ${AI_SERVER_URL} (mTLS): ${models
.map((m) => m.id)
.join(", ")}`,
);
}
}
dark-mechanicus-indicator.ts
+14 -15
View File
@@ -20,6 +20,8 @@
*/
import type { ExtensionAPI } from "@mariozechner/pi-coding-agent";
import { safely } from "./shared/ctx.js";
import { formatElapsed } from "./shared/format.js";
// ─── Quote pools (15 per category) ───────────────────────────────────────
@@ -106,14 +108,6 @@ function shuffle<T>(source: readonly T[]): T[] {
return arr;
}
function formatElapsed(ms: number): string {
const total = Math.max(0, Math.floor(ms / 1000));
const m = Math.floor(total / 60);
const s = total % 60;
if (m > 0) return `${m}m ${s}s`;
return `${s}s`;
}
function buildFrames(quote: string, elapsedMs: number): string[] {
const elapsedStr = formatElapsed(elapsedMs);
return PULSE_SHADES.map(
@@ -153,13 +147,18 @@ export default function (pi: ExtensionAPI) {
const applyFrame = (ctx: any) => {
if (!ctx.hasUI || mode === "off" || startedAt === null) return;
const elapsed = Date.now() - startedAt;
const quoteIdx =
Math.floor(elapsed / QUOTE_DWELL_MS) % shuffled.length;
const quote = shuffled[quoteIdx]!;
ctx.ui.setWorkingIndicator({
frames: buildFrames(quote, elapsed),
intervalMs: FRAME_INTERVAL_MS,
// Guard against stale-ctx: the setInterval ticker can fire between
// session_shutdown being queued and our turn_end/session_shutdown
// handlers running. Any ctx.ui.* access on a stale runner throws.
safely(() => {
const elapsed = Date.now() - startedAt!;
const quoteIdx =
Math.floor(elapsed / QUOTE_DWELL_MS) % shuffled.length;
const quote = shuffled[quoteIdx]!;
ctx.ui.setWorkingIndicator({
frames: buildFrames(quote, elapsed),
intervalMs: FRAME_INTERVAL_MS,
});
});
};
markdown-body-color.ts
+21 -19
View File
@@ -23,21 +23,11 @@
import type { ExtensionAPI } from "@mariozechner/pi-coding-agent";
import { Editor, Markdown } from "@mariozechner/pi-tui";
import { ANSI_RESET_FG, fgFromHex } from "./shared/ansi.js";
const FALLBACK_HEX = process.env.PI_MARKDOWN_BODY_COLOR ?? "#d4c5e8";
function hexToRgb(hex: string): [number, number, number] {
const m = hex.replace(/^#/, "");
const n = parseInt(
m.length === 3 ? m.split("").map((c) => c + c).join("") : m,
16,
);
return [(n >> 16) & 0xff, (n >> 8) & 0xff, n & 0xff];
}
const [R, G, B] = hexToRgb(FALLBACK_HEX);
const OPEN = `\x1b[38;2;${R};${G};${B}m`;
const CLOSE = `\x1b[39m`;
const OPEN = fgFromHex(FALLBACK_HEX);
const CLOSE = ANSI_RESET_FG;
// Re-open the body color after any ANSI reset inside a line. Cursor inverse
// video uses \x1b[0m (full reset) and theme.fg helpers close with \x1b[39m.
@@ -56,9 +46,15 @@ export default function (_pi: ExtensionAPI) {
patched = true;
// ─ 1. Markdown component ────────────────────────────────────────────
const mdProto = Markdown.prototype as any;
const origMdRender = mdProto.render;
mdProto.render = function (width: number) {
const mdProto = Markdown.prototype as Record<string, unknown>;
const origMdRender = mdProto.render as ((width: number) => string[]) | undefined;
if (typeof origMdRender !== "function") {
console.warn(
"[markdown-body-color] pi-tui Markdown.prototype.render is missing; skipping patch. (pi-tui upgrade?)",
);
return;
}
mdProto.render = function (this: any, width: number) {
if (!this.defaultTextStyle) {
this.defaultTextStyle = { color: wrap };
if (typeof this.invalidate === "function") this.invalidate();
@@ -67,9 +63,15 @@ export default function (_pi: ExtensionAPI) {
};
// ─ 2. Editor component ──────────────────────────────────────────────
const edProto = Editor.prototype as any;
const origEdRender = edProto.render;
edProto.render = function (width: number) {
const edProto = Editor.prototype as Record<string, unknown>;
const origEdRender = edProto.render as ((width: number) => string[]) | undefined;
if (typeof origEdRender !== "function") {
console.warn(
"[markdown-body-color] pi-tui Editor.prototype.render is missing; skipping Editor patch.",
);
return;
}
edProto.render = function (this: any, width: number) {
const lines: string[] = origEdRender.call(this, width);
return lines.map(wrap);
};
mechanicus-footer.ts
+1 -11
View File
@@ -22,17 +22,7 @@ import type {
ExtensionContext,
Theme,
} from "@mariozechner/pi-coding-agent";
function formatTokens(n: number): string {
if (n < 1000) return String(n);
if (n < 10_000) return `${(n / 1000).toFixed(1)}k`;
return `${Math.round(n / 1000)}k`;
}
// Strip ANSI escape sequences to measure visible width for padding.
function visibleWidth(s: string): number {
return s.replace(/\x1b\[[0-9;]*m/g, "").length;
}
import { formatTokens } from "./shared/format.js";
export default function (pi: ExtensionAPI) {
let cachedCtx: ExtensionContext | null = null;
mechanicus-thinking-label.ts
+9 -3
View File
@@ -22,10 +22,16 @@ export default function (_pi: ExtensionAPI) {
if (patched) return;
patched = true;
const proto = AssistantMessageComponent.prototype as any;
const origUpdate = proto.updateContent;
const proto = AssistantMessageComponent.prototype as Record<string, unknown>;
const origUpdate = proto.updateContent as ((m: unknown) => void) | undefined;
if (typeof origUpdate !== "function") {
console.warn(
"[mechanicus-thinking-label] AssistantMessageComponent.updateContent is missing; skipping patch.",
);
return;
}
proto.updateContent = function (message: unknown) {
proto.updateContent = function (this: any, message: unknown) {
if (this.hiddenThinkingLabel === DEFAULT_LABEL) {
this.hiddenThinkingLabel = LABEL;
}