feat: context-warn extension — alert on context threshold
This commit is contained in:
@@ -0,0 +1,152 @@
|
||||
/**
|
||||
* Context Warning Extension
|
||||
*
|
||||
* Emits a generated sound alert when the context window reaches or exceeds
|
||||
* 100 000 tokens. Also shows a UI notification.
|
||||
*
|
||||
* Usage:
|
||||
* pi -e ./context-warn-extension/index.ts
|
||||
*
|
||||
* Or place in ~/.pi/agent/extensions/context-warn/ for auto-discovery.
|
||||
*/
|
||||
|
||||
import type { ExtensionAPI } from "@earendil-works/pi-coding-agent";
|
||||
import { execFileSync } from "node:child_process";
|
||||
import { existsSync } from "node:fs";
|
||||
import { join, dirname } from "node:path";
|
||||
import { fileURLToPath } from "node:url";
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Configuration
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
const WARNING_TOKEN_THRESHOLD = 100_000; // alert at 100k tokens
|
||||
const CHECK_INTERVAL_MS = 15_000; // re-check every 15 s
|
||||
const WARNING_COOLDOWN_MS = 60_000; // min seconds between alerts
|
||||
|
||||
// Resolve the directory where this extension lives so we can find alert.wav
|
||||
const EXT_DIR = dirname(fileURLToPath(import.meta.url));
|
||||
const ALERT_WAV = join(EXT_DIR, "alert.wav");
|
||||
|
||||
/**
|
||||
* Play the bundled alert.wav sound via aplay.
|
||||
* No runtime generation needed — the WAV is committed alongside this file.
|
||||
*/
|
||||
function playAlert(): void {
|
||||
if (!existsSync(ALERT_WAV)) {
|
||||
console.error(`[context-warn] alert.wav not found at ${ALERT_WAV}`);
|
||||
return;
|
||||
}
|
||||
|
||||
try {
|
||||
execFileSync("aplay", ["-q", ALERT_WAV], { stdio: "ignore" });
|
||||
} catch (err) {
|
||||
console.error(`[context-warn] Audio play failed: ${err instanceof Error ? err.message : String(err)}`);
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Extension
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
export default function (pi: ExtensionAPI) {
|
||||
let lastAlertTimestamp: number | null = null;
|
||||
let alertedAtThreshold: boolean = false;
|
||||
|
||||
/** Check context usage and alert if threshold reached. */
|
||||
const checkContext = (_event: unknown, ctx: ReturnType<ExtensionAPI["on"]> extends (ev: string, h: (e: any, c: infer C) => any) ? C : never) => {
|
||||
const usage = ctx.getContextUsage();
|
||||
if (!usage || usage.tokens === null || usage.tokens === undefined) {
|
||||
return;
|
||||
}
|
||||
|
||||
const currentTokens = usage.tokens;
|
||||
|
||||
// Check if we crossed the threshold (went from below to >= threshold)
|
||||
const crossedThreshold = !alertedAtThreshold && currentTokens >= WARNING_TOKEN_THRESHOLD;
|
||||
alertedAtThreshold = alertedAtThreshold || crossedThreshold;
|
||||
|
||||
if (crossedThreshold) {
|
||||
// Check cooldown
|
||||
const now = Date.now();
|
||||
if (lastAlertTimestamp && (now - lastAlertTimestamp) < WARNING_COOLDOWN_MS) {
|
||||
return; // still in cooldown
|
||||
}
|
||||
lastAlertTimestamp = now;
|
||||
|
||||
// Show UI notification
|
||||
if (ctx.hasUI) {
|
||||
ctx.ui.notify(
|
||||
"⚠️ Context Warning",
|
||||
`Context reached ${currentTokens.toLocaleString()} tokens (>${WARNING_TOKEN_THRESHOLD.toLocaleString()}). Consider compacting.`,
|
||||
"warning"
|
||||
);
|
||||
ctx.ui.setStatus("context-warn", `${(currentTokens / 1000).toFixed(0)}k / ${usage.contextWindow ? `${usage.contextWindow / 1000}k` : "?"}`);
|
||||
}
|
||||
|
||||
// Play alert sound
|
||||
playAlert();
|
||||
|
||||
console.log(`[context-warn] Context: ${currentTokens.toLocaleString()} tokens (threshold: ${WARNING_TOKEN_THRESHOLD.toLocaleString()})`);
|
||||
}
|
||||
};
|
||||
|
||||
// Listen to turn_end — fired after every LLM response, which is when
|
||||
// we know the new context size.
|
||||
pi.on("turn_end", (_event, ctx) => {
|
||||
checkContext(_event, ctx);
|
||||
});
|
||||
|
||||
// Also listen to agent_end in case compaction or other events happen
|
||||
pi.on("agent_end", (_event, ctx) => {
|
||||
checkContext(_event, ctx);
|
||||
});
|
||||
|
||||
// Listen to session_start to reset state
|
||||
pi.on("session_start", (_event, ctx) => {
|
||||
alertedAtThreshold = false;
|
||||
lastAlertTimestamp = null;
|
||||
if (ctx.hasUI) {
|
||||
ctx.ui.notify("🔊 Context monitor active", `Alerts at ≥ ${WARNING_TOKEN_THRESHOLD.toLocaleString()} tokens`, "info");
|
||||
ctx.ui.setStatus("context-warn", "monitoring…");
|
||||
}
|
||||
});
|
||||
|
||||
// Listen to session_compact to reset alert state
|
||||
pi.on("session_compact", (_event, ctx) => {
|
||||
alertedAtThreshold = false;
|
||||
lastAlertTimestamp = Date.now();
|
||||
if (ctx.hasUI) {
|
||||
ctx.ui.setStatus("context-warn", "compact ✓");
|
||||
}
|
||||
});
|
||||
|
||||
// Register /context-warn-status command
|
||||
pi.registerCommand("context-warn-status", {
|
||||
description: "Show current context token usage and warning status",
|
||||
handler: async (_args, ctx) => {
|
||||
const usage = ctx.getContextUsage();
|
||||
if (!usage || usage.tokens === null || usage.tokens === undefined) {
|
||||
ctx.ui.notify("Context", "No usage data available yet.", "info");
|
||||
return;
|
||||
}
|
||||
const pct = usage.percent !== null && usage.percent !== undefined
|
||||
? `${usage.percent.toFixed(1)}%`
|
||||
: "?";
|
||||
ctx.ui.notify(
|
||||
"Context Status",
|
||||
`${usage.tokens.toLocaleString()} tokens / ${usage.contextWindow?.toLocaleString()} (${pct})\nThreshold: ${WARNING_TOKEN_THRESHOLD.toLocaleString()}\nAlerted: ${alertedAtThreshold ? "yes" : "no"}`,
|
||||
"info"
|
||||
);
|
||||
},
|
||||
});
|
||||
|
||||
// Register /context-warn-alert command — manual test
|
||||
pi.registerCommand("context-warn-alert", {
|
||||
description: "Play the alert sound immediately (test)",
|
||||
handler: async (_args, ctx) => {
|
||||
playAlert();
|
||||
ctx.ui.notify("🔔 Alert", "Sound played.", "info");
|
||||
},
|
||||
});
|
||||
}
|
||||
Reference in New Issue
Block a user