266 lines
9.2 KiB
TypeScript
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");
|
|
}
|
|
},
|
|
});
|
|
}
|