dev_sandbox/.pi/agent/extensions/new-reminder.ts
2026-06-18 14:52:14 -04:00

266 lines
9.2 KiB
TypeScript

/**
* new-reminder
*
* Two features for managing long-running sessions:
*
* 1. A reminder widget above the editor that nudges you to start a fresh
* session (`/new`). It shows how long the session has been running and how
* many turns it has accumulated. When the session "runs long" (exceeds a
* configurable age or turn threshold) the box turns yellow.
*
* 2. `/snew [prompt]` - summarize the current chat, then start a brand new
* session seeded with the generated summary. If a prompt is provided, the
* summary is prepended to it and the combined prompt is auto-submitted in
* the new session. If no prompt is provided, the summary lands in the
* editor as an editable draft and waits for user input.
*/
import { complete, type Message } from "@earendil-works/pi-ai";
import type { ExtensionAPI, ExtensionContext, SessionEntry, Theme } from "@earendil-works/pi-coding-agent";
import type { Component, TUI } from "@earendil-works/pi-tui";
// ---------------------------------------------------------------------------
// Tunables
// ---------------------------------------------------------------------------
/** Session is considered "long" once estimated context exceeds this many tokens. */
const TOKEN_THRESHOLD = 250_000;
/** ...or after this many user turns, whichever comes first. */
const TURN_THRESHOLD = 25;
/** Widget refresh cadence (ms). Keeps the elapsed-time readout current. */
const TICK_MS = 30_000;
// ---------------------------------------------------------------------------
// Conversation extraction helpers (for /snew)
// ---------------------------------------------------------------------------
type ContentBlock = {
type?: string;
text?: string;
name?: string;
arguments?: Record<string, unknown>;
};
function extractText(content: unknown): string[] {
if (typeof content === "string") return [content];
if (!Array.isArray(content)) return [];
const out: string[] = [];
for (const part of content) {
if (!part || typeof part !== "object") continue;
const block = part as ContentBlock;
if (block.type === "text" && typeof block.text === "string") out.push(block.text);
}
return out;
}
function buildConversationText(branch: SessionEntry[]): string {
const sections: string[] = [];
for (const entry of branch) {
if (entry.type === "compaction") {
const summary = (entry as { summary?: string }).summary;
if (summary) sections.push(`Previous summary:\n${summary}`);
continue;
}
if (entry.type !== "message") continue;
const msg = entry.message as { role?: string; content?: unknown } | undefined;
if (!msg?.role) continue;
if (msg.role !== "user" && msg.role !== "assistant") continue;
const text = extractText(msg.content).join("\n").trim();
if (text.length === 0) continue;
const label = msg.role === "user" ? "User" : "Assistant";
sections.push(`${label}: ${text}`);
}
return sections.join("\n\n");
}
const SUMMARY_SYSTEM_PROMPT = `You are a context transfer assistant. Summarize the conversation so it can be resumed in a new, self-contained thread. Capture goals, key decisions, approaches taken, important findings, files touched, and open questions. Be concise and structured with short headings. Output only the summary, with no preamble.`;
// ---------------------------------------------------------------------------
// Reminder widget
// ---------------------------------------------------------------------------
function formatTokens(tokens: number): string {
if (tokens < 1000) return `${tokens}`;
return `${(tokens / 1000).toFixed(tokens < 10_000 ? 1 : 0)}k`;
}
export default function (pi: ExtensionAPI) {
// --- shared state for the reminder widget ---
let turnCount = 0;
let activeTui: TUI | undefined;
let activeCtx: ExtensionContext | undefined;
let ticker: ReturnType<typeof setInterval> | undefined;
/** Estimated context tokens for the active model, or null if unknown. */
const contextTokens = (): number | null => activeCtx?.getContextUsage()?.tokens ?? null;
const isLong = () => {
const tokens = contextTokens();
return (tokens !== null && tokens >= TOKEN_THRESHOLD) || turnCount >= TURN_THRESHOLD;
};
const stopTicker = () => {
if (ticker) {
clearInterval(ticker);
ticker = undefined;
}
};
// Build a bordered box component. Yellow ("warning") when the session is long,
// muted otherwise.
const makeWidget = (tui: TUI, theme: Theme): Component => {
activeTui = tui;
return {
render(width: number): string[] {
const long = isLong();
const colorKey = long ? "warning" : "muted";
const color = (s: string) => theme.fg(colorKey, s);
const tokens = contextTokens();
const icon = long ? "\u26a0" : "\u23f1"; // warning sign / stopwatch
const lead = long
? `${icon} Long session \u2014 consider /new`
: `${icon} Session`;
const ctxStat = tokens === null ? "ctx ?" : `${formatTokens(tokens)} ctx`;
const stats = `${ctxStat} \u00b7 ${turnCount} turn${turnCount === 1 ? "" : "s"}`;
const label = ` ${lead} \u00b7 ${stats} `;
const inner = Math.max(0, width - 2);
const truncated = label.length > inner ? label.slice(0, inner) : label;
const pad = inner - truncated.length;
const top = color("\u256d" + "\u2500".repeat(inner) + "\u256e");
const mid = color("\u2502") + (long ? theme.bold(theme.fg("warning", truncated)) : color(truncated)) + " ".repeat(pad) + color("\u2502");
const bot = color("\u2570" + "\u2500".repeat(inner) + "\u256f");
return [top, mid, bot];
},
invalidate() {},
};
};
const startReminder = (ctx: ExtensionContext) => {
if (!ctx.hasUI) return;
activeCtx = ctx;
ctx.ui.setWidget("new-reminder", (tui, theme) => makeWidget(tui, theme), { placement: "aboveEditor" });
stopTicker();
ticker = setInterval(() => activeTui?.requestRender(), TICK_MS);
};
pi.on("session_start", (_event, ctx) => {
// Reset the turn counter for the (possibly new/resumed) session.
turnCount = 0;
startReminder(ctx);
});
pi.on("agent_start", () => {
turnCount += 1;
activeTui?.requestRender();
});
pi.on("session_shutdown", () => {
stopTicker();
activeTui = undefined;
activeCtx = undefined;
});
// -----------------------------------------------------------------------
// /snew - summarize then start a new session seeded with summary + prompt
// -----------------------------------------------------------------------
pi.registerCommand("snew", {
description: "Summarize the current chat, then start a new session. With a prompt: auto-submit summary+prompt. Without: seed the editor and wait.",
handler: async (args, ctx) => {
if (!ctx.hasUI) {
ctx.ui.notify("/snew requires interactive mode", "error");
return;
}
const userPrompt = args.trim();
if (!ctx.model) {
ctx.ui.notify("No model selected", "error");
return;
}
await ctx.waitForIdle();
const conversationText = buildConversationText(ctx.sessionManager.getBranch());
if (!conversationText.trim()) {
ctx.ui.notify("No conversation to summarize", "warning");
return;
}
const auth = await ctx.modelRegistry.getApiKeyAndHeaders(ctx.model);
if (!auth.ok || !auth.apiKey) {
ctx.ui.notify(auth.ok ? `No API key for ${ctx.model.provider}` : auth.error, "error");
return;
}
ctx.ui.notify("Summarizing current chat...", "info");
const summaryMessages: Message[] = [
{
role: "user",
content: [
{
type: "text",
text: `Summarize this conversation for handoff to a new session.\n\n<conversation>\n${conversationText}\n</conversation>`,
},
],
timestamp: Date.now(),
},
];
let summary: string;
try {
const response = await complete(
ctx.model,
{ systemPrompt: SUMMARY_SYSTEM_PROMPT, messages: summaryMessages },
{ apiKey: auth.apiKey, headers: auth.headers, reasoningEffort: "low", signal: ctx.signal },
);
summary = response.content
.filter((c): c is { type: "text"; text: string } => c.type === "text")
.map((c) => c.text)
.join("\n")
.trim();
} catch (err) {
ctx.ui.notify(`Summary failed: ${(err as Error).message}`, "error");
return;
}
if (!summary) {
ctx.ui.notify("Empty summary; aborting", "warning");
return;
}
const parentSession = ctx.sessionManager.getSessionFile();
if (userPrompt) {
// Prompt provided: combine summary + prompt and auto-submit.
const combinedPrompt = `## Context from previous session\n\n${summary}\n\n---\n\n## Task\n\n${userPrompt}`;
const result = await ctx.newSession({
parentSession,
withSession: async (replacementCtx) => {
await replacementCtx.sendUserMessage(combinedPrompt);
},
});
if (result.cancelled) {
ctx.ui.notify("New session cancelled", "info");
}
return;
}
// No prompt: seed the editor with just the summary and wait for input.
const draftPrompt = `## Context from previous session\n\n${summary}\n\n---\n\n## Task\n\n`;
const result = await ctx.newSession({
parentSession,
withSession: async (replacementCtx) => {
replacementCtx.ui.setEditorText(draftPrompt);
replacementCtx.ui.notify("New session seeded with summary. Add your prompt and submit.", "info");
},
});
if (result.cancelled) {
ctx.ui.notify("New session cancelled", "info");
}
},
});
}