sync local pi changes
This commit is contained in:
11
.gitignore
vendored
11
.gitignore
vendored
@@ -1,3 +1,14 @@
|
|||||||
.worktrees/
|
.worktrees/
|
||||||
.pi/npm/
|
.pi/npm/
|
||||||
.pi/agent/sessions/
|
.pi/agent/sessions/
|
||||||
|
.pi/agent/auth.json
|
||||||
|
.pi/agent/web-search.json
|
||||||
|
.pi/subagents/
|
||||||
|
.pi/agent/extensions/.pi/
|
||||||
|
.pi/agent/extensions/tmux-subagent/events.jsonl
|
||||||
|
.pi/agent/extensions/tmux-subagent/result.json
|
||||||
|
.pi/agent/extensions/tmux-subagent/stderr.log
|
||||||
|
.pi/agent/extensions/tmux-subagent/stdout.log
|
||||||
|
.pi/agent/extensions/tmux-subagent/transcript.log
|
||||||
|
*bun.lock
|
||||||
|
*node_modules
|
||||||
|
|||||||
@@ -1,15 +0,0 @@
|
|||||||
{
|
|
||||||
"github-copilot": {
|
|
||||||
"type": "oauth",
|
|
||||||
"refresh": "ghu_j9QHUrVzPLoYOsyjarpzktAFDQWqP31gz2Ac",
|
|
||||||
"access": "tid=af454cc719f9e4daffe9b4892fa4e791;exp=1775732126;sku=plus_monthly_subscriber_quota;proxy-ep=proxy.individual.githubcopilot.com;st=dotcom;chat=1;cit=1;malfil=1;editor_preview_features=1;agent_mode=1;agent_mode_auto_approval=1;mcp=1;client_byok=0;ccr=1;8kp=1;ip=81.104.194.177;asn=AS5089:e4ff19791adbf3b64531636bad853d4de1c02e75ac46089baa3d2d799cbefadf",
|
|
||||||
"expires": 1775731826000
|
|
||||||
},
|
|
||||||
"openai-codex": {
|
|
||||||
"type": "oauth",
|
|
||||||
"access": "eyJhbGciOiJSUzI1NiIsImtpZCI6IjE5MzQ0ZTY1LWJiYzktNDRkMS1hOWQwLWY5NTdiMDc5YmQwZSIsInR5cCI6IkpXVCJ9.eyJhdWQiOlsiaHR0cHM6Ly9hcGkub3BlbmFpLmNvbS92MSJdLCJjbGllbnRfaWQiOiJhcHBfRU1vYW1FRVo3M2YwQ2tYYVhwN2hyYW5uIiwiZXhwIjoxNzc2NDE1MDUwLCJodHRwczovL2FwaS5vcGVuYWkuY29tL2F1dGgiOnsiYW1yIjpbInBvcCIsInVybjpvcGVuYWk6YW1yOnBhc3NrZXkiLCJtZmEiXSwiY2hhdGdwdF9hY2NvdW50X2lkIjoiOTY1MTZkMjYtMjljOS00Y2JjLWEwZDItNmZjODdlNzc3ZjRhIiwiY2hhdGdwdF9hY2NvdW50X3VzZXJfaWQiOiJ1c2VyLXloUkkzTjdiVHlvc0xBd1I5NmNOU25wUV9fOTY1MTZkMjYtMjljOS00Y2JjLWEwZDItNmZjODdlNzc3ZjRhIiwiY2hhdGdwdF9jb21wdXRlX3Jlc2lkZW5jeSI6Im5vX2NvbnN0cmFpbnQiLCJjaGF0Z3B0X3BsYW5fdHlwZSI6InBsdXMiLCJjaGF0Z3B0X3VzZXJfaWQiOiJ1c2VyLXloUkkzTjdiVHlvc0xBd1I5NmNOU25wUSIsImxvY2FsaG9zdCI6dHJ1ZSwidXNlcl9pZCI6InVzZXIteWhSSTNON2JUeW9zTEF3Ujk2Y05TbnBRIn0sImh0dHBzOi8vYXBpLm9wZW5haS5jb20vbWZhIjp7InJlcXVpcmVkIjoieWVzIn0sImh0dHBzOi8vYXBpLm9wZW5haS5jb20vcHJvZmlsZSI6eyJlbWFpbCI6ImNoYXRncHQuY29tLmRldGVyZ2VudDI3N0BwYXNzbWFpbC5uZXQiLCJlbWFpbF92ZXJpZmllZCI6dHJ1ZX0sImlhdCI6MTc3NTU1MTA1MCwiaXNzIjoiaHR0cHM6Ly9hdXRoLm9wZW5haS5jb20iLCJqdGkiOiI5OGZmNDRjOC02ZWFlLTRhOWQtOGQ1Yy04MTNkYWI1NjY4ZDgiLCJuYmYiOjE3NzU1NTEwNTAsInB3ZF9hdXRoX3RpbWUiOjE3NzU1NTEwNDg4NTAsInNjcCI6WyJvcGVuaWQiLCJwcm9maWxlIiwiZW1haWwiLCJvZmZsaW5lX2FjY2VzcyJdLCJzZXNzaW9uX2lkIjoiYXV0aHNlc3NfOXVQRzcwbDcyQ1dkbTVtMXkyMmp6WkpyIiwic2wiOnRydWUsInN1YiI6ImF1dGgwfG1zcmNqUDZTYkgzdTROUHFzZ1Y2SERyNyJ9.AIcfng7BnC_IUK8DYedcWI8M6AZ5r2FszzM4orhrI5Ql0nXQ-eZBtV8RVcSl6wHvkcj6XX-BcpxxJQL_w0JbRPs4utQ7ayTeEhFYut8OnsLCTcMJDF0s5Qwv4GlTJNbuG_3P6hBe8xiZ6kPkDp0ZihZOkiceghPEaBRh_npt8-zm7SQyl8R8qdfhFToYzUAgGox3aZHVQeWGWpBm39MB_WigA6jsLCK5h-SwX5iuSHppGzii8ohyiaTgHfcEKUa9kgWXHa4iOtPHxPtD3t_rWJTZuc3XfeO4V3raR8HT96m8wrAHTgKlNA5IrmVwj8pt_fUH6AbApMrJY9q5Le6ubzCbH5bmnO2PIVLKfd7Kyw-E1gtjSOH61dvgRxDFLNwjAMeKNYRnrsPRZRr1pI5Y4JV9VejsjEE-MdvN48EEIWbZn4MvKtSSd5Xr_RGZPS80wLWV0WV_5qWL62aYJjTS4Vz4B3kWFBQsPNp08ykd2NL7b5H-uuP3akY97Jasklzvhuc9BgQZBymVlGO6Fwq1GiRggCu62B6OKJlxKOqgTOHGNGhFgmgGQWxpz-cCm-qKTb81vBEbziNBmXQdhL-507cFMJwsYBYyxKI1x79Gn3odkzHWoyijTxSCColYeqOBOdba9B9y8hdNmUwhn42W27A6Hm0bojiPoerUh6ng7Nk",
|
|
||||||
"refresh": "rt_LTCkO68CsFMGg9wrJP3qgfroR-b32AXV7Uw9cmtD_nA.4ZFAy5DZCiJaIEbHiSpLyddbqWhs02ZB53NMA9PRjq8",
|
|
||||||
"expires": 1776415049237,
|
|
||||||
"accountId": "96516d26-29c9-4cbc-a0d2-6fc87e777f4a"
|
|
||||||
}
|
|
||||||
}
|
|
||||||
353
.pi/agent/extensions/context-manager/index.ts
Normal file
353
.pi/agent/extensions/context-manager/index.ts
Normal file
@@ -0,0 +1,353 @@
|
|||||||
|
import type { AgentMessage } from "@mariozechner/pi-agent-core";
|
||||||
|
import type { ExtensionAPI, ExtensionContext } from "@mariozechner/pi-coding-agent";
|
||||||
|
import { adjustPolicyForZone } from "./src/config.ts";
|
||||||
|
import { deserializeLatestSnapshot, serializeSnapshot, SNAPSHOT_ENTRY_TYPE, type RuntimeSnapshot } from "./src/persist.ts";
|
||||||
|
import { createEmptyLedger } from "./src/ledger.ts";
|
||||||
|
import { pruneContextMessages } from "./src/prune.ts";
|
||||||
|
import { createContextManagerRuntime } from "./src/runtime.ts";
|
||||||
|
import { registerContextCommands } from "./src/commands.ts";
|
||||||
|
import { buildBranchSummaryFromEntries, buildCompactionSummaryFromPreparation } from "./src/summaries.ts";
|
||||||
|
|
||||||
|
type TrackedMessage = Extract<AgentMessage, { role: "user" | "assistant" | "toolResult" }>;
|
||||||
|
type BranchEntry = ReturnType<ExtensionContext["sessionManager"]["getBranch"]>[number];
|
||||||
|
|
||||||
|
function isTextPart(part: unknown): part is { type: "text"; text?: string } {
|
||||||
|
return typeof part === "object" && part !== null && "type" in part && (part as { type?: unknown }).type === "text";
|
||||||
|
}
|
||||||
|
|
||||||
|
function toText(content: unknown): string {
|
||||||
|
if (typeof content === "string") return content;
|
||||||
|
if (!Array.isArray(content)) return "";
|
||||||
|
|
||||||
|
return content
|
||||||
|
.map((part) => {
|
||||||
|
if (!isTextPart(part)) return "";
|
||||||
|
return typeof part.text === "string" ? part.text : "";
|
||||||
|
})
|
||||||
|
.join("\n")
|
||||||
|
.trim();
|
||||||
|
}
|
||||||
|
|
||||||
|
function isMessageEntry(entry: BranchEntry): entry is Extract<BranchEntry, { type: "message" }> {
|
||||||
|
return entry.type === "message";
|
||||||
|
}
|
||||||
|
|
||||||
|
function isCompactionEntry(entry: BranchEntry): entry is Extract<BranchEntry, { type: "compaction" }> {
|
||||||
|
return entry.type === "compaction";
|
||||||
|
}
|
||||||
|
|
||||||
|
function isBranchSummaryEntry(entry: BranchEntry): entry is Extract<BranchEntry, { type: "branch_summary" }> {
|
||||||
|
return entry.type === "branch_summary";
|
||||||
|
}
|
||||||
|
|
||||||
|
function isTrackedMessage(message: AgentMessage): message is TrackedMessage {
|
||||||
|
return message.role === "user" || message.role === "assistant" || message.role === "toolResult";
|
||||||
|
}
|
||||||
|
|
||||||
|
function createDefaultSnapshot(): RuntimeSnapshot {
|
||||||
|
return {
|
||||||
|
mode: "balanced",
|
||||||
|
lastZone: "green",
|
||||||
|
ledger: createEmptyLedger(),
|
||||||
|
};
|
||||||
|
}
|
||||||
|
|
||||||
|
function getMessageContent(message: AgentMessage): string {
|
||||||
|
return "content" in message ? toText(message.content) : "";
|
||||||
|
}
|
||||||
|
|
||||||
|
function getMessageToolName(message: AgentMessage): string | undefined {
|
||||||
|
return message.role === "toolResult" ? message.toolName : undefined;
|
||||||
|
}
|
||||||
|
|
||||||
|
function rewriteContextMessage(message: { role: string; content: string; original: AgentMessage; distilled?: boolean }): AgentMessage {
|
||||||
|
if (!message.distilled || message.role !== "toolResult") {
|
||||||
|
return message.original;
|
||||||
|
}
|
||||||
|
|
||||||
|
return {
|
||||||
|
...(message.original as Extract<AgentMessage, { role: "toolResult" }>),
|
||||||
|
content: [{ type: "text", text: message.content }],
|
||||||
|
} as AgentMessage;
|
||||||
|
}
|
||||||
|
|
||||||
|
function findLatestSnapshotState(branch: BranchEntry[]): { snapshot: RuntimeSnapshot; index: number } | undefined {
|
||||||
|
for (let index = branch.length - 1; index >= 0; index -= 1) {
|
||||||
|
const entry = branch[index]!;
|
||||||
|
if (entry.type !== "custom" || entry.customType !== SNAPSHOT_ENTRY_TYPE) {
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
|
||||||
|
const snapshot = deserializeLatestSnapshot([entry]);
|
||||||
|
if (snapshot) {
|
||||||
|
return { snapshot, index };
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
return undefined;
|
||||||
|
}
|
||||||
|
|
||||||
|
function findLatestSessionSnapshot(entries: BranchEntry[]): RuntimeSnapshot | undefined {
|
||||||
|
let latest: RuntimeSnapshot | undefined;
|
||||||
|
let latestFreshness = -Infinity;
|
||||||
|
|
||||||
|
for (const entry of entries) {
|
||||||
|
if (entry.type !== "custom" || entry.customType !== SNAPSHOT_ENTRY_TYPE) {
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
|
||||||
|
const snapshot = deserializeLatestSnapshot([entry]);
|
||||||
|
if (!snapshot) {
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
|
||||||
|
const sessionItems = snapshot.ledger.items.filter((item) => item.scope === "session");
|
||||||
|
const freshness = sessionItems.length > 0 ? Math.max(...sessionItems.map((item) => item.timestamp)) : -Infinity;
|
||||||
|
if (freshness >= latestFreshness) {
|
||||||
|
latest = snapshot;
|
||||||
|
latestFreshness = freshness;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
return latest;
|
||||||
|
}
|
||||||
|
|
||||||
|
function createSessionFallbackSnapshot(source?: RuntimeSnapshot): RuntimeSnapshot {
|
||||||
|
return {
|
||||||
|
mode: source?.mode ?? "balanced",
|
||||||
|
lastZone: "green",
|
||||||
|
ledger: {
|
||||||
|
items: structuredClone((source?.ledger.items ?? []).filter((item) => item.scope === "session")),
|
||||||
|
rollingSummary: "",
|
||||||
|
},
|
||||||
|
};
|
||||||
|
}
|
||||||
|
|
||||||
|
function overlaySessionLayer(base: RuntimeSnapshot, latestSessionSnapshot?: RuntimeSnapshot): RuntimeSnapshot {
|
||||||
|
const sessionItems = latestSessionSnapshot?.ledger.items.filter((item) => item.scope === "session") ?? [];
|
||||||
|
if (sessionItems.length === 0) {
|
||||||
|
return base;
|
||||||
|
}
|
||||||
|
|
||||||
|
return {
|
||||||
|
...base,
|
||||||
|
ledger: {
|
||||||
|
...base.ledger,
|
||||||
|
items: [
|
||||||
|
...structuredClone(base.ledger.items.filter((item) => item.scope !== "session")),
|
||||||
|
...structuredClone(sessionItems),
|
||||||
|
],
|
||||||
|
},
|
||||||
|
};
|
||||||
|
}
|
||||||
|
|
||||||
|
export default function contextManager(pi: ExtensionAPI) {
|
||||||
|
const runtime = createContextManagerRuntime({
|
||||||
|
mode: "balanced",
|
||||||
|
contextWindow: 200_000,
|
||||||
|
});
|
||||||
|
let pendingResumeInjection = false;
|
||||||
|
|
||||||
|
const syncContextWindow = (ctx: Pick<ExtensionContext, "model">) => {
|
||||||
|
runtime.setContextWindow(ctx.model?.contextWindow ?? 200_000);
|
||||||
|
};
|
||||||
|
|
||||||
|
const armResumeInjection = () => {
|
||||||
|
const snapshot = runtime.getSnapshot();
|
||||||
|
pendingResumeInjection = Boolean(snapshot.lastCompactionSummary || snapshot.lastBranchSummary) && runtime.buildResumePacket().trim().length > 0;
|
||||||
|
};
|
||||||
|
|
||||||
|
const replayBranchEntry = (entry: BranchEntry) => {
|
||||||
|
if (isMessageEntry(entry) && isTrackedMessage(entry.message)) {
|
||||||
|
runtime.ingest({
|
||||||
|
entryId: entry.id,
|
||||||
|
role: entry.message.role,
|
||||||
|
text: toText(entry.message.content),
|
||||||
|
timestamp: entry.message.timestamp,
|
||||||
|
isError: entry.message.role === "toolResult" ? entry.message.isError : undefined,
|
||||||
|
});
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
|
||||||
|
if (isCompactionEntry(entry)) {
|
||||||
|
runtime.recordCompactionSummary(entry.summary, entry.id, Date.parse(entry.timestamp));
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
|
||||||
|
if (isBranchSummaryEntry(entry)) {
|
||||||
|
runtime.recordBranchSummary(entry.summary, entry.id, Date.parse(entry.timestamp));
|
||||||
|
}
|
||||||
|
};
|
||||||
|
|
||||||
|
const rebuildRuntimeFromBranch = (
|
||||||
|
ctx: Pick<ExtensionContext, "model" | "sessionManager" | "ui">,
|
||||||
|
fallbackSnapshot: RuntimeSnapshot,
|
||||||
|
options?: { preferRuntimeMode?: boolean },
|
||||||
|
) => {
|
||||||
|
syncContextWindow(ctx);
|
||||||
|
|
||||||
|
const branch = ctx.sessionManager.getBranch();
|
||||||
|
const latestSessionSnapshot = findLatestSessionSnapshot(ctx.sessionManager.getEntries() as BranchEntry[]);
|
||||||
|
const restored = findLatestSnapshotState(branch);
|
||||||
|
const baseSnapshot = restored
|
||||||
|
? overlaySessionLayer(restored.snapshot, latestSessionSnapshot)
|
||||||
|
: createSessionFallbackSnapshot(latestSessionSnapshot ?? fallbackSnapshot);
|
||||||
|
|
||||||
|
runtime.restore({
|
||||||
|
...baseSnapshot,
|
||||||
|
mode: options?.preferRuntimeMode ? fallbackSnapshot.mode : baseSnapshot.mode,
|
||||||
|
});
|
||||||
|
|
||||||
|
const replayEntries = restored ? branch.slice(restored.index + 1) : branch;
|
||||||
|
for (const entry of replayEntries) {
|
||||||
|
replayBranchEntry(entry);
|
||||||
|
}
|
||||||
|
|
||||||
|
const snapshot = runtime.getSnapshot();
|
||||||
|
ctx.ui.setStatus("context-manager", `ctx ${snapshot.lastZone}`);
|
||||||
|
};
|
||||||
|
|
||||||
|
registerContextCommands(pi, {
|
||||||
|
getSnapshot: runtime.getSnapshot,
|
||||||
|
buildPacket: runtime.buildPacket,
|
||||||
|
buildResumePacket: runtime.buildResumePacket,
|
||||||
|
setMode: runtime.setMode,
|
||||||
|
rebuildFromBranch: async (commandCtx) => {
|
||||||
|
rebuildRuntimeFromBranch(commandCtx, runtime.getSnapshot(), { preferRuntimeMode: true });
|
||||||
|
armResumeInjection();
|
||||||
|
},
|
||||||
|
isResumePending: () => pendingResumeInjection,
|
||||||
|
});
|
||||||
|
|
||||||
|
pi.on("session_start", async (_event, ctx) => {
|
||||||
|
rebuildRuntimeFromBranch(ctx, createDefaultSnapshot());
|
||||||
|
armResumeInjection();
|
||||||
|
});
|
||||||
|
|
||||||
|
pi.on("session_tree", async (event, ctx) => {
|
||||||
|
rebuildRuntimeFromBranch(ctx, createDefaultSnapshot());
|
||||||
|
|
||||||
|
if (
|
||||||
|
event.summaryEntry &&
|
||||||
|
!ctx.sessionManager.getBranch().some((entry) => isBranchSummaryEntry(entry) && entry.id === event.summaryEntry.id)
|
||||||
|
) {
|
||||||
|
runtime.recordBranchSummary(event.summaryEntry.summary, event.summaryEntry.id, Date.parse(event.summaryEntry.timestamp));
|
||||||
|
}
|
||||||
|
|
||||||
|
armResumeInjection();
|
||||||
|
|
||||||
|
if (event.summaryEntry) {
|
||||||
|
pi.appendEntry(SNAPSHOT_ENTRY_TYPE, serializeSnapshot(runtime.getSnapshot()));
|
||||||
|
}
|
||||||
|
});
|
||||||
|
|
||||||
|
pi.on("tool_result", async (event) => {
|
||||||
|
runtime.ingest({
|
||||||
|
entryId: event.toolCallId,
|
||||||
|
role: "toolResult",
|
||||||
|
text: toText(event.content),
|
||||||
|
timestamp: Date.now(),
|
||||||
|
});
|
||||||
|
});
|
||||||
|
|
||||||
|
pi.on("turn_end", async (_event, ctx) => {
|
||||||
|
rebuildRuntimeFromBranch(ctx, runtime.getSnapshot(), { preferRuntimeMode: true });
|
||||||
|
|
||||||
|
const usage = ctx.getContextUsage();
|
||||||
|
if (usage?.tokens !== null && usage?.tokens !== undefined) {
|
||||||
|
runtime.observeTokens(usage.tokens);
|
||||||
|
}
|
||||||
|
|
||||||
|
const snapshot = runtime.getSnapshot();
|
||||||
|
pi.appendEntry(SNAPSHOT_ENTRY_TYPE, serializeSnapshot(snapshot));
|
||||||
|
ctx.ui.setStatus("context-manager", `ctx ${snapshot.lastZone}`);
|
||||||
|
});
|
||||||
|
|
||||||
|
pi.on("context", async (event, ctx) => {
|
||||||
|
syncContextWindow(ctx);
|
||||||
|
const snapshot = runtime.getSnapshot();
|
||||||
|
const policy = adjustPolicyForZone(runtime.getPolicy(), snapshot.lastZone);
|
||||||
|
const normalized = event.messages.map((message) => ({
|
||||||
|
role: message.role,
|
||||||
|
content: getMessageContent(message),
|
||||||
|
toolName: getMessageToolName(message),
|
||||||
|
original: message,
|
||||||
|
}));
|
||||||
|
|
||||||
|
const pruned = pruneContextMessages(normalized, policy);
|
||||||
|
const nextMessages = pruned.map((message) =>
|
||||||
|
rewriteContextMessage(message as { role: string; content: string; original: AgentMessage; distilled?: boolean }),
|
||||||
|
);
|
||||||
|
const resumeText = pendingResumeInjection ? runtime.buildResumePacket() : "";
|
||||||
|
const packetText = pendingResumeInjection ? "" : runtime.buildPacket().text;
|
||||||
|
const injectedText = resumeText || packetText;
|
||||||
|
|
||||||
|
if (!injectedText) {
|
||||||
|
return { messages: nextMessages };
|
||||||
|
}
|
||||||
|
|
||||||
|
if (resumeText) {
|
||||||
|
pendingResumeInjection = false;
|
||||||
|
}
|
||||||
|
|
||||||
|
return {
|
||||||
|
messages: [
|
||||||
|
{
|
||||||
|
role: "custom",
|
||||||
|
customType: resumeText ? "context-manager.resume" : "context-manager.packet",
|
||||||
|
content: injectedText,
|
||||||
|
display: false,
|
||||||
|
timestamp: Date.now(),
|
||||||
|
} as any,
|
||||||
|
...nextMessages,
|
||||||
|
],
|
||||||
|
};
|
||||||
|
});
|
||||||
|
|
||||||
|
pi.on("session_before_compact", async (event, ctx) => {
|
||||||
|
syncContextWindow(ctx);
|
||||||
|
|
||||||
|
try {
|
||||||
|
return {
|
||||||
|
compaction: {
|
||||||
|
summary: buildCompactionSummaryFromPreparation({
|
||||||
|
messagesToSummarize: event.preparation.messagesToSummarize,
|
||||||
|
turnPrefixMessages: event.preparation.turnPrefixMessages,
|
||||||
|
previousSummary: event.preparation.previousSummary,
|
||||||
|
fileOps: event.preparation.fileOps,
|
||||||
|
customInstructions: event.customInstructions,
|
||||||
|
}),
|
||||||
|
firstKeptEntryId: event.preparation.firstKeptEntryId,
|
||||||
|
tokensBefore: event.preparation.tokensBefore,
|
||||||
|
details: event.preparation.fileOps,
|
||||||
|
},
|
||||||
|
};
|
||||||
|
} catch (error) {
|
||||||
|
ctx.ui.notify(`context-manager compaction fallback: ${error instanceof Error ? error.message : String(error)}`, "warning");
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
});
|
||||||
|
|
||||||
|
pi.on("session_before_tree", async (event, ctx) => {
|
||||||
|
syncContextWindow(ctx);
|
||||||
|
if (!event.preparation.userWantsSummary) return;
|
||||||
|
return {
|
||||||
|
summary: {
|
||||||
|
summary: buildBranchSummaryFromEntries({
|
||||||
|
branchLabel: "branch handoff",
|
||||||
|
entriesToSummarize: event.preparation.entriesToSummarize,
|
||||||
|
customInstructions: event.preparation.customInstructions,
|
||||||
|
replaceInstructions: event.preparation.replaceInstructions,
|
||||||
|
commonAncestorId: event.preparation.commonAncestorId,
|
||||||
|
}),
|
||||||
|
},
|
||||||
|
};
|
||||||
|
});
|
||||||
|
|
||||||
|
pi.on("session_compact", async (event, ctx) => {
|
||||||
|
runtime.recordCompactionSummary(event.compactionEntry.summary, event.compactionEntry.id, Date.parse(event.compactionEntry.timestamp));
|
||||||
|
pi.appendEntry(SNAPSHOT_ENTRY_TYPE, serializeSnapshot(runtime.getSnapshot()));
|
||||||
|
armResumeInjection();
|
||||||
|
ctx.ui.setStatus("context-manager", `ctx ${runtime.getSnapshot().lastZone}`);
|
||||||
|
});
|
||||||
|
}
|
||||||
4361
.pi/agent/extensions/context-manager/package-lock.json
generated
Normal file
4361
.pi/agent/extensions/context-manager/package-lock.json
generated
Normal file
File diff suppressed because it is too large
Load Diff
18
.pi/agent/extensions/context-manager/package.json
Normal file
18
.pi/agent/extensions/context-manager/package.json
Normal file
@@ -0,0 +1,18 @@
|
|||||||
|
{
|
||||||
|
"name": "pi-context-manager-extension",
|
||||||
|
"private": true,
|
||||||
|
"version": "0.0.0",
|
||||||
|
"type": "module",
|
||||||
|
"scripts": {
|
||||||
|
"test": "tsx --test src/*.test.ts src/**/*.test.ts"
|
||||||
|
},
|
||||||
|
"pi": {
|
||||||
|
"extensions": ["./index.ts"]
|
||||||
|
},
|
||||||
|
"devDependencies": {
|
||||||
|
"@mariozechner/pi-coding-agent": "^0.66.1",
|
||||||
|
"@types/node": "^25.5.2",
|
||||||
|
"tsx": "^4.21.0",
|
||||||
|
"typescript": "^6.0.2"
|
||||||
|
}
|
||||||
|
}
|
||||||
76
.pi/agent/extensions/context-manager/src/commands.ts
Normal file
76
.pi/agent/extensions/context-manager/src/commands.ts
Normal file
@@ -0,0 +1,76 @@
|
|||||||
|
import type { ExtensionAPI, ExtensionCommandContext } from "@mariozechner/pi-coding-agent";
|
||||||
|
import type { ContextMode } from "./config.ts";
|
||||||
|
import { serializeSnapshot, SNAPSHOT_ENTRY_TYPE, type RuntimeSnapshot } from "./persist.ts";
|
||||||
|
|
||||||
|
interface CommandRuntime {
|
||||||
|
getSnapshot(): RuntimeSnapshot;
|
||||||
|
buildPacket(): { estimatedTokens: number };
|
||||||
|
buildResumePacket(): string;
|
||||||
|
setMode(mode: ContextMode): void;
|
||||||
|
rebuildFromBranch(ctx: ExtensionCommandContext): Promise<void>;
|
||||||
|
isResumePending(): boolean;
|
||||||
|
}
|
||||||
|
|
||||||
|
export function registerContextCommands(pi: ExtensionAPI, runtime: CommandRuntime) {
|
||||||
|
pi.registerCommand("ctx-status", {
|
||||||
|
description: "Show context pressure, packet status, and persisted handoff state",
|
||||||
|
handler: async (_args, ctx) => {
|
||||||
|
const snapshot = runtime.getSnapshot();
|
||||||
|
const packet = runtime.buildPacket();
|
||||||
|
const resumePending = runtime.isResumePending();
|
||||||
|
const contextTokens = ctx.getContextUsage()?.tokens;
|
||||||
|
const nextInjectionTokens = resumePending ? Math.ceil(runtime.buildResumePacket().length / 4) : packet.estimatedTokens;
|
||||||
|
ctx.ui.notify(
|
||||||
|
[
|
||||||
|
`mode=${snapshot.mode}`,
|
||||||
|
`zone=${snapshot.lastZone}`,
|
||||||
|
`contextTokens=${contextTokens ?? "unknown"}`,
|
||||||
|
`packetTokens=${packet.estimatedTokens}`,
|
||||||
|
`nextInjectionTokens=${nextInjectionTokens}`,
|
||||||
|
`resumePending=${resumePending ? "yes" : "no"}`,
|
||||||
|
`compaction=${snapshot.lastCompactionSummary ? "yes" : "no"}`,
|
||||||
|
`branch=${snapshot.lastBranchSummary ? "yes" : "no"}`,
|
||||||
|
].join(" "),
|
||||||
|
"info",
|
||||||
|
);
|
||||||
|
},
|
||||||
|
});
|
||||||
|
|
||||||
|
pi.registerCommand("ctx-memory", {
|
||||||
|
description: "Inspect the active context ledger",
|
||||||
|
handler: async (_args, ctx) => {
|
||||||
|
const snapshot = runtime.getSnapshot();
|
||||||
|
await ctx.ui.editor("Context ledger", JSON.stringify(snapshot.ledger, null, 2));
|
||||||
|
},
|
||||||
|
});
|
||||||
|
|
||||||
|
pi.registerCommand("ctx-refresh", {
|
||||||
|
description: "Rebuild runtime state from the current branch and refresh the working packet",
|
||||||
|
handler: async (_args, ctx) => {
|
||||||
|
await runtime.rebuildFromBranch(ctx);
|
||||||
|
const packet = runtime.buildPacket();
|
||||||
|
ctx.ui.notify(`rebuilt runtime from branch (${packet.estimatedTokens} tokens)`, "info");
|
||||||
|
},
|
||||||
|
});
|
||||||
|
|
||||||
|
pi.registerCommand("ctx-compact", {
|
||||||
|
description: "Trigger compaction with optional focus instructions",
|
||||||
|
handler: async (args, ctx) => {
|
||||||
|
ctx.compact({ customInstructions: args.trim() || undefined });
|
||||||
|
},
|
||||||
|
});
|
||||||
|
|
||||||
|
pi.registerCommand("ctx-mode", {
|
||||||
|
description: "Switch context mode: conservative | balanced | aggressive",
|
||||||
|
handler: async (args, ctx) => {
|
||||||
|
const value = args.trim() as "conservative" | "balanced" | "aggressive";
|
||||||
|
if (!["conservative", "balanced", "aggressive"].includes(value)) {
|
||||||
|
ctx.ui.notify("usage: /ctx-mode conservative|balanced|aggressive", "warning");
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
runtime.setMode(value);
|
||||||
|
pi.appendEntry(SNAPSHOT_ENTRY_TYPE, serializeSnapshot(runtime.getSnapshot()));
|
||||||
|
ctx.ui.notify(`context mode set to ${value}`, "info");
|
||||||
|
},
|
||||||
|
});
|
||||||
|
}
|
||||||
86
.pi/agent/extensions/context-manager/src/config.test.ts
Normal file
86
.pi/agent/extensions/context-manager/src/config.test.ts
Normal file
@@ -0,0 +1,86 @@
|
|||||||
|
import test from "node:test";
|
||||||
|
import assert from "node:assert/strict";
|
||||||
|
import { adjustPolicyForZone, resolvePolicy, zoneForTokens } from "./config.ts";
|
||||||
|
|
||||||
|
test("resolvePolicy returns the balanced policy for a 200k context window", () => {
|
||||||
|
const policy = resolvePolicy({ mode: "balanced", contextWindow: 200_000 });
|
||||||
|
|
||||||
|
assert.deepEqual(policy, {
|
||||||
|
mode: "balanced",
|
||||||
|
recentUserTurns: 4,
|
||||||
|
packetTokenCap: 1_200,
|
||||||
|
bulkyBytes: 4_096,
|
||||||
|
bulkyLines: 150,
|
||||||
|
yellowAtTokens: 110_000,
|
||||||
|
redAtTokens: 140_000,
|
||||||
|
compactAtTokens: 164_000,
|
||||||
|
});
|
||||||
|
});
|
||||||
|
|
||||||
|
test("resolvePolicy clamps context windows below 50k before calculating thresholds", () => {
|
||||||
|
const policy = resolvePolicy({ mode: "balanced", contextWindow: 10_000 });
|
||||||
|
|
||||||
|
assert.deepEqual(policy, {
|
||||||
|
mode: "balanced",
|
||||||
|
recentUserTurns: 4,
|
||||||
|
packetTokenCap: 1_200,
|
||||||
|
bulkyBytes: 4_096,
|
||||||
|
bulkyLines: 150,
|
||||||
|
yellowAtTokens: 27_500,
|
||||||
|
redAtTokens: 35_000,
|
||||||
|
compactAtTokens: 41_000,
|
||||||
|
});
|
||||||
|
});
|
||||||
|
|
||||||
|
test("aggressive mode compacts earlier than conservative mode", () => {
|
||||||
|
const aggressive = resolvePolicy({ mode: "aggressive", contextWindow: 200_000 });
|
||||||
|
const conservative = resolvePolicy({ mode: "conservative", contextWindow: 200_000 });
|
||||||
|
|
||||||
|
assert.ok(aggressive.compactAtTokens < conservative.compactAtTokens);
|
||||||
|
});
|
||||||
|
|
||||||
|
test("aggressive mode reduces raw-window and packet budgets compared with conservative mode", () => {
|
||||||
|
const aggressive = resolvePolicy({ mode: "aggressive", contextWindow: 200_000 });
|
||||||
|
const conservative = resolvePolicy({ mode: "conservative", contextWindow: 200_000 });
|
||||||
|
|
||||||
|
assert.ok(aggressive.recentUserTurns < conservative.recentUserTurns);
|
||||||
|
assert.ok(aggressive.packetTokenCap < conservative.packetTokenCap);
|
||||||
|
assert.ok(aggressive.bulkyBytes < conservative.bulkyBytes);
|
||||||
|
assert.ok(aggressive.bulkyLines < conservative.bulkyLines);
|
||||||
|
});
|
||||||
|
|
||||||
|
test("adjustPolicyForZone tightens packet and pruning thresholds in yellow, red, and compact zones", () => {
|
||||||
|
const base = resolvePolicy({ mode: "balanced", contextWindow: 200_000 });
|
||||||
|
const yellow = adjustPolicyForZone(base, "yellow");
|
||||||
|
const red = adjustPolicyForZone(base, "red");
|
||||||
|
const compact = adjustPolicyForZone(base, "compact");
|
||||||
|
|
||||||
|
assert.ok(yellow.packetTokenCap < base.packetTokenCap);
|
||||||
|
assert.ok(yellow.bulkyBytes < base.bulkyBytes);
|
||||||
|
assert.ok(red.packetTokenCap < yellow.packetTokenCap);
|
||||||
|
assert.ok(red.recentUserTurns <= yellow.recentUserTurns);
|
||||||
|
assert.ok(red.bulkyBytes < yellow.bulkyBytes);
|
||||||
|
assert.ok(compact.packetTokenCap < red.packetTokenCap);
|
||||||
|
assert.ok(compact.recentUserTurns <= red.recentUserTurns);
|
||||||
|
assert.ok(compact.bulkyLines < red.bulkyLines);
|
||||||
|
});
|
||||||
|
|
||||||
|
test("zoneForTokens returns green, yellow, red, and compact for the balanced 200k policy", () => {
|
||||||
|
const policy = resolvePolicy({ mode: "balanced", contextWindow: 200_000 });
|
||||||
|
|
||||||
|
assert.equal(zoneForTokens(80_000, policy), "green");
|
||||||
|
assert.equal(zoneForTokens(120_000, policy), "yellow");
|
||||||
|
assert.equal(zoneForTokens(150_000, policy), "red");
|
||||||
|
assert.equal(zoneForTokens(170_000, policy), "compact");
|
||||||
|
});
|
||||||
|
|
||||||
|
test("zoneForTokens uses inclusive balanced 200k thresholds", () => {
|
||||||
|
const policy = resolvePolicy({ mode: "balanced", contextWindow: 200_000 });
|
||||||
|
|
||||||
|
assert.equal(zoneForTokens(109_999, policy), "green");
|
||||||
|
assert.equal(zoneForTokens(110_000, policy), "yellow");
|
||||||
|
assert.equal(zoneForTokens(139_999, policy), "yellow");
|
||||||
|
assert.equal(zoneForTokens(140_000, policy), "red");
|
||||||
|
assert.equal(zoneForTokens(163_999, policy), "red");
|
||||||
|
assert.equal(zoneForTokens(164_000, policy), "compact");
|
||||||
|
});
|
||||||
97
.pi/agent/extensions/context-manager/src/config.ts
Normal file
97
.pi/agent/extensions/context-manager/src/config.ts
Normal file
@@ -0,0 +1,97 @@
|
|||||||
|
export type ContextMode = "conservative" | "balanced" | "aggressive";
|
||||||
|
export type ContextZone = "green" | "yellow" | "red" | "compact";
|
||||||
|
|
||||||
|
export interface Policy {
|
||||||
|
mode: ContextMode;
|
||||||
|
recentUserTurns: number;
|
||||||
|
packetTokenCap: number;
|
||||||
|
bulkyBytes: number;
|
||||||
|
bulkyLines: number;
|
||||||
|
yellowAtTokens: number;
|
||||||
|
redAtTokens: number;
|
||||||
|
compactAtTokens: number;
|
||||||
|
}
|
||||||
|
|
||||||
|
export const MODE_PCTS: Record<ContextMode, { yellow: number; red: number; compact: number }> = {
|
||||||
|
conservative: { yellow: 0.60, red: 0.76, compact: 0.88 },
|
||||||
|
balanced: { yellow: 0.55, red: 0.70, compact: 0.82 },
|
||||||
|
aggressive: { yellow: 0.50, red: 0.64, compact: 0.76 },
|
||||||
|
};
|
||||||
|
|
||||||
|
const MODE_SETTINGS: Record<ContextMode, Pick<Policy, "recentUserTurns" | "packetTokenCap" | "bulkyBytes" | "bulkyLines">> = {
|
||||||
|
conservative: {
|
||||||
|
recentUserTurns: 5,
|
||||||
|
packetTokenCap: 1_400,
|
||||||
|
bulkyBytes: 6_144,
|
||||||
|
bulkyLines: 220,
|
||||||
|
},
|
||||||
|
balanced: {
|
||||||
|
recentUserTurns: 4,
|
||||||
|
packetTokenCap: 1_200,
|
||||||
|
bulkyBytes: 4_096,
|
||||||
|
bulkyLines: 150,
|
||||||
|
},
|
||||||
|
aggressive: {
|
||||||
|
recentUserTurns: 3,
|
||||||
|
packetTokenCap: 900,
|
||||||
|
bulkyBytes: 3_072,
|
||||||
|
bulkyLines: 100,
|
||||||
|
},
|
||||||
|
};
|
||||||
|
|
||||||
|
export function resolvePolicy(input: { mode: ContextMode; contextWindow: number }): Policy {
|
||||||
|
const contextWindow = Math.max(input.contextWindow, 50_000);
|
||||||
|
const percentages = MODE_PCTS[input.mode];
|
||||||
|
const settings = MODE_SETTINGS[input.mode];
|
||||||
|
|
||||||
|
return {
|
||||||
|
mode: input.mode,
|
||||||
|
recentUserTurns: settings.recentUserTurns,
|
||||||
|
packetTokenCap: settings.packetTokenCap,
|
||||||
|
bulkyBytes: settings.bulkyBytes,
|
||||||
|
bulkyLines: settings.bulkyLines,
|
||||||
|
yellowAtTokens: Math.floor(contextWindow * percentages.yellow),
|
||||||
|
redAtTokens: Math.floor(contextWindow * percentages.red),
|
||||||
|
compactAtTokens: Math.floor(contextWindow * percentages.compact),
|
||||||
|
};
|
||||||
|
}
|
||||||
|
|
||||||
|
export function adjustPolicyForZone(policy: Policy, zone: ContextZone): Policy {
|
||||||
|
if (zone === "green") {
|
||||||
|
return { ...policy };
|
||||||
|
}
|
||||||
|
|
||||||
|
if (zone === "yellow") {
|
||||||
|
return {
|
||||||
|
...policy,
|
||||||
|
packetTokenCap: Math.max(500, Math.floor(policy.packetTokenCap * 0.9)),
|
||||||
|
bulkyBytes: Math.max(1_536, Math.floor(policy.bulkyBytes * 0.9)),
|
||||||
|
bulkyLines: Math.max(80, Math.floor(policy.bulkyLines * 0.9)),
|
||||||
|
};
|
||||||
|
}
|
||||||
|
|
||||||
|
if (zone === "red") {
|
||||||
|
return {
|
||||||
|
...policy,
|
||||||
|
recentUserTurns: Math.max(2, policy.recentUserTurns - 1),
|
||||||
|
packetTokenCap: Math.max(400, Math.floor(policy.packetTokenCap * 0.75)),
|
||||||
|
bulkyBytes: Math.max(1_024, Math.floor(policy.bulkyBytes * 0.75)),
|
||||||
|
bulkyLines: Math.max(60, Math.floor(policy.bulkyLines * 0.75)),
|
||||||
|
};
|
||||||
|
}
|
||||||
|
|
||||||
|
return {
|
||||||
|
...policy,
|
||||||
|
recentUserTurns: Math.max(1, policy.recentUserTurns - 2),
|
||||||
|
packetTokenCap: Math.max(300, Math.floor(policy.packetTokenCap * 0.55)),
|
||||||
|
bulkyBytes: Math.max(768, Math.floor(policy.bulkyBytes * 0.5)),
|
||||||
|
bulkyLines: Math.max(40, Math.floor(policy.bulkyLines * 0.5)),
|
||||||
|
};
|
||||||
|
}
|
||||||
|
|
||||||
|
export function zoneForTokens(tokens: number, policy: Policy): ContextZone {
|
||||||
|
if (tokens >= policy.compactAtTokens) return "compact";
|
||||||
|
if (tokens >= policy.redAtTokens) return "red";
|
||||||
|
if (tokens >= policy.yellowAtTokens) return "yellow";
|
||||||
|
return "green";
|
||||||
|
}
|
||||||
30
.pi/agent/extensions/context-manager/src/distill.test.ts
Normal file
30
.pi/agent/extensions/context-manager/src/distill.test.ts
Normal file
@@ -0,0 +1,30 @@
|
|||||||
|
import test from "node:test";
|
||||||
|
import assert from "node:assert/strict";
|
||||||
|
import { distillToolResult } from "./distill.ts";
|
||||||
|
|
||||||
|
const noisy = [
|
||||||
|
"Build failed while compiling focus parser",
|
||||||
|
"Error: missing export createFocusMatcher from ./summary-focus.ts",
|
||||||
|
"at src/summaries.ts:44:12",
|
||||||
|
"line filler",
|
||||||
|
"line filler",
|
||||||
|
"line filler",
|
||||||
|
].join("\n");
|
||||||
|
|
||||||
|
test("distillToolResult prioritizes salient failure lines and truncates noise", () => {
|
||||||
|
const distilled = distillToolResult({ toolName: "bash", content: noisy });
|
||||||
|
|
||||||
|
assert.ok(distilled);
|
||||||
|
assert.match(distilled!, /Build failed while compiling focus parser/);
|
||||||
|
assert.match(distilled!, /missing export createFocusMatcher/);
|
||||||
|
assert.ok(distilled!.length < 320);
|
||||||
|
});
|
||||||
|
|
||||||
|
test("distillToolResult falls back to the first meaningful non-empty lines", () => {
|
||||||
|
const distilled = distillToolResult({
|
||||||
|
toolName: "read",
|
||||||
|
content: ["", "src/runtime.ts", "exports createContextManagerRuntime", "", "more noise"].join("\n"),
|
||||||
|
});
|
||||||
|
|
||||||
|
assert.equal(distilled, "[distilled read output] src/runtime.ts; exports createContextManagerRuntime");
|
||||||
|
});
|
||||||
47
.pi/agent/extensions/context-manager/src/distill.ts
Normal file
47
.pi/agent/extensions/context-manager/src/distill.ts
Normal file
@@ -0,0 +1,47 @@
|
|||||||
|
const ERROR_RE = /\b(?:error|failed|failure|missing|undefined|exception)\b/i;
|
||||||
|
const LOCATION_RE = /\b(?:at\s+.+:\d+(?::\d+)?)\b|(?:[A-Za-z0-9_./-]+\.(?:ts|tsx|js|mjs|json|md):\d+(?::\d+)?)/i;
|
||||||
|
const MAX_SUMMARY_LENGTH = 320;
|
||||||
|
const MAX_LINES = 2;
|
||||||
|
|
||||||
|
function unique(lines: string[]): string[] {
|
||||||
|
return lines.filter((line, index) => lines.indexOf(line) === index);
|
||||||
|
}
|
||||||
|
|
||||||
|
function pickSalientLines(content: string): string[] {
|
||||||
|
const lines = content
|
||||||
|
.split(/\n+/)
|
||||||
|
.map((line) => line.trim())
|
||||||
|
.filter(Boolean);
|
||||||
|
|
||||||
|
if (lines.length === 0) {
|
||||||
|
return [];
|
||||||
|
}
|
||||||
|
|
||||||
|
const important = unique(lines.filter((line) => ERROR_RE.test(line)));
|
||||||
|
const location = unique(lines.filter((line) => LOCATION_RE.test(line)));
|
||||||
|
const fallback = unique(lines);
|
||||||
|
|
||||||
|
const selected: string[] = [];
|
||||||
|
for (const line of [...important, ...location, ...fallback]) {
|
||||||
|
if (selected.includes(line)) {
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
|
||||||
|
selected.push(line);
|
||||||
|
if (selected.length >= MAX_LINES) {
|
||||||
|
break;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
return selected;
|
||||||
|
}
|
||||||
|
|
||||||
|
export function distillToolResult(input: { toolName?: string; content: string }): string | undefined {
|
||||||
|
const picked = pickSalientLines(input.content);
|
||||||
|
if (picked.length === 0) {
|
||||||
|
return undefined;
|
||||||
|
}
|
||||||
|
|
||||||
|
const prefix = `[distilled ${input.toolName ?? "tool"} output]`;
|
||||||
|
return `${prefix} ${picked.join("; ")}`.slice(0, MAX_SUMMARY_LENGTH);
|
||||||
|
}
|
||||||
833
.pi/agent/extensions/context-manager/src/extension.test.ts
Normal file
833
.pi/agent/extensions/context-manager/src/extension.test.ts
Normal file
@@ -0,0 +1,833 @@
|
|||||||
|
import test from "node:test";
|
||||||
|
import assert from "node:assert/strict";
|
||||||
|
import contextManagerExtension from "../index.ts";
|
||||||
|
import { deserializeLatestSnapshot, SNAPSHOT_ENTRY_TYPE, serializeSnapshot, type RuntimeSnapshot } from "./persist.ts";
|
||||||
|
|
||||||
|
type EventHandler = (event: any, ctx: any) => Promise<any> | any;
|
||||||
|
type RegisteredCommand = { description: string; handler: (args: string, ctx: any) => Promise<void> | void };
|
||||||
|
|
||||||
|
type SessionEntry =
|
||||||
|
| {
|
||||||
|
type: "message";
|
||||||
|
id: string;
|
||||||
|
parentId: string | null;
|
||||||
|
timestamp: string;
|
||||||
|
message: any;
|
||||||
|
}
|
||||||
|
| {
|
||||||
|
type: "custom";
|
||||||
|
id: string;
|
||||||
|
parentId: string | null;
|
||||||
|
timestamp: string;
|
||||||
|
customType: string;
|
||||||
|
data: unknown;
|
||||||
|
}
|
||||||
|
| {
|
||||||
|
type: "compaction";
|
||||||
|
id: string;
|
||||||
|
parentId: string | null;
|
||||||
|
timestamp: string;
|
||||||
|
summary: string;
|
||||||
|
firstKeptEntryId: string;
|
||||||
|
tokensBefore: number;
|
||||||
|
}
|
||||||
|
| {
|
||||||
|
type: "branch_summary";
|
||||||
|
id: string;
|
||||||
|
parentId: string | null;
|
||||||
|
timestamp: string;
|
||||||
|
fromId: string;
|
||||||
|
summary: string;
|
||||||
|
};
|
||||||
|
|
||||||
|
function createUsage(tokens: number) {
|
||||||
|
return {
|
||||||
|
tokens,
|
||||||
|
contextWindow: 200_000,
|
||||||
|
percent: tokens / 200_000,
|
||||||
|
};
|
||||||
|
}
|
||||||
|
|
||||||
|
function createUserMessage(content: string, timestamp: number) {
|
||||||
|
return {
|
||||||
|
role: "user",
|
||||||
|
content,
|
||||||
|
timestamp,
|
||||||
|
};
|
||||||
|
}
|
||||||
|
|
||||||
|
function createAssistantMessage(content: string, timestamp: number) {
|
||||||
|
return {
|
||||||
|
role: "assistant",
|
||||||
|
content: [{ type: "text", text: content }],
|
||||||
|
api: "openai-responses",
|
||||||
|
provider: "openai",
|
||||||
|
model: "gpt-5",
|
||||||
|
usage: {
|
||||||
|
input: 0,
|
||||||
|
output: 0,
|
||||||
|
cacheRead: 0,
|
||||||
|
cacheWrite: 0,
|
||||||
|
totalTokens: 0,
|
||||||
|
cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
|
||||||
|
},
|
||||||
|
stopReason: "stop",
|
||||||
|
timestamp,
|
||||||
|
};
|
||||||
|
}
|
||||||
|
|
||||||
|
function createToolResultMessage(content: string, timestamp: number) {
|
||||||
|
return {
|
||||||
|
role: "toolResult",
|
||||||
|
toolCallId: `tool-${timestamp}`,
|
||||||
|
toolName: "read",
|
||||||
|
content: [{ type: "text", text: content }],
|
||||||
|
isError: false,
|
||||||
|
timestamp,
|
||||||
|
};
|
||||||
|
}
|
||||||
|
|
||||||
|
function createMessageEntry(id: string, parentId: string | null, message: any): SessionEntry {
|
||||||
|
return {
|
||||||
|
type: "message",
|
||||||
|
id,
|
||||||
|
parentId,
|
||||||
|
timestamp: new Date(message.timestamp).toISOString(),
|
||||||
|
message,
|
||||||
|
};
|
||||||
|
}
|
||||||
|
|
||||||
|
function createSnapshotEntry(
|
||||||
|
id: string,
|
||||||
|
parentId: string | null,
|
||||||
|
options: {
|
||||||
|
text: string;
|
||||||
|
mode?: RuntimeSnapshot["mode"];
|
||||||
|
lastZone?: RuntimeSnapshot["lastZone"];
|
||||||
|
lastObservedTokens?: number;
|
||||||
|
lastCompactionSummary?: string;
|
||||||
|
lastBranchSummary?: string;
|
||||||
|
ledgerItems?: RuntimeSnapshot["ledger"]["items"];
|
||||||
|
rollingSummary?: string;
|
||||||
|
},
|
||||||
|
): SessionEntry {
|
||||||
|
const {
|
||||||
|
text,
|
||||||
|
mode = "aggressive",
|
||||||
|
lastZone = "red",
|
||||||
|
lastObservedTokens = 150_000,
|
||||||
|
lastCompactionSummary = "existing compaction summary",
|
||||||
|
lastBranchSummary = "existing branch summary",
|
||||||
|
ledgerItems,
|
||||||
|
rollingSummary = "stale ledger",
|
||||||
|
} = options;
|
||||||
|
|
||||||
|
return {
|
||||||
|
type: "custom",
|
||||||
|
id,
|
||||||
|
parentId,
|
||||||
|
timestamp: new Date(1).toISOString(),
|
||||||
|
customType: SNAPSHOT_ENTRY_TYPE,
|
||||||
|
data: serializeSnapshot({
|
||||||
|
mode,
|
||||||
|
lastZone,
|
||||||
|
lastObservedTokens,
|
||||||
|
lastCompactionSummary,
|
||||||
|
lastBranchSummary,
|
||||||
|
ledger: {
|
||||||
|
items: ledgerItems ?? [
|
||||||
|
{
|
||||||
|
id: `goal:session:root-goal:${id}`,
|
||||||
|
kind: "goal",
|
||||||
|
subject: "root-goal",
|
||||||
|
text,
|
||||||
|
scope: "session",
|
||||||
|
sourceEntryId: "old-user",
|
||||||
|
sourceType: "user",
|
||||||
|
timestamp: 1,
|
||||||
|
confidence: 1,
|
||||||
|
freshness: 1,
|
||||||
|
active: true,
|
||||||
|
supersedesId: undefined,
|
||||||
|
},
|
||||||
|
],
|
||||||
|
rollingSummary,
|
||||||
|
},
|
||||||
|
}),
|
||||||
|
};
|
||||||
|
}
|
||||||
|
|
||||||
|
function createHarness(initialBranch: SessionEntry[], options?: { usageTokens?: number }) {
|
||||||
|
const commands = new Map<string, RegisteredCommand>();
|
||||||
|
const handlers = new Map<string, EventHandler>();
|
||||||
|
const appendedEntries: Array<{ customType: string; data: unknown }> = [];
|
||||||
|
const statuses: Array<{ key: string; value: string }> = [];
|
||||||
|
let branch = [...initialBranch];
|
||||||
|
let entries = [...initialBranch];
|
||||||
|
|
||||||
|
const ctx = {
|
||||||
|
model: { contextWindow: 200_000 },
|
||||||
|
sessionManager: {
|
||||||
|
getBranch() {
|
||||||
|
return branch;
|
||||||
|
},
|
||||||
|
getEntries() {
|
||||||
|
return entries;
|
||||||
|
},
|
||||||
|
},
|
||||||
|
ui: {
|
||||||
|
setStatus(key: string, value: string) {
|
||||||
|
statuses.push({ key, value });
|
||||||
|
},
|
||||||
|
notify() {},
|
||||||
|
editor: async () => {},
|
||||||
|
},
|
||||||
|
getContextUsage() {
|
||||||
|
return options?.usageTokens === undefined ? undefined : createUsage(options.usageTokens);
|
||||||
|
},
|
||||||
|
compact() {},
|
||||||
|
};
|
||||||
|
|
||||||
|
contextManagerExtension({
|
||||||
|
registerCommand(name: string, command: RegisteredCommand) {
|
||||||
|
commands.set(name, command);
|
||||||
|
},
|
||||||
|
on(name: string, handler: EventHandler) {
|
||||||
|
handlers.set(name, handler);
|
||||||
|
},
|
||||||
|
appendEntry(customType: string, data: unknown) {
|
||||||
|
appendedEntries.push({ customType, data });
|
||||||
|
const entry = {
|
||||||
|
type: "custom" as const,
|
||||||
|
id: `custom-${appendedEntries.length}`,
|
||||||
|
parentId: branch.at(-1)?.id ?? null,
|
||||||
|
timestamp: new Date(10_000 + appendedEntries.length).toISOString(),
|
||||||
|
customType,
|
||||||
|
data,
|
||||||
|
};
|
||||||
|
branch.push(entry);
|
||||||
|
entries.push(entry);
|
||||||
|
},
|
||||||
|
} as any);
|
||||||
|
|
||||||
|
return {
|
||||||
|
commands,
|
||||||
|
handlers,
|
||||||
|
appendedEntries,
|
||||||
|
statuses,
|
||||||
|
ctx,
|
||||||
|
setBranch(nextBranch: SessionEntry[]) {
|
||||||
|
branch = [...nextBranch];
|
||||||
|
const byId = new Map(entries.map((entry) => [entry.id, entry]));
|
||||||
|
for (const entry of nextBranch) {
|
||||||
|
byId.set(entry.id, entry);
|
||||||
|
}
|
||||||
|
entries = [...byId.values()];
|
||||||
|
},
|
||||||
|
};
|
||||||
|
}
|
||||||
|
|
||||||
|
test("the extension registers the expected hooks and commands", () => {
|
||||||
|
const harness = createHarness([]);
|
||||||
|
|
||||||
|
assert.deepEqual([...harness.commands.keys()].sort(), ["ctx-compact", "ctx-memory", "ctx-mode", "ctx-refresh", "ctx-status"]);
|
||||||
|
assert.deepEqual(
|
||||||
|
[...harness.handlers.keys()].sort(),
|
||||||
|
["context", "session_before_compact", "session_before_tree", "session_compact", "session_start", "session_tree", "tool_result", "turn_end"],
|
||||||
|
);
|
||||||
|
});
|
||||||
|
|
||||||
|
test("turn_end persists a rebuilt snapshot that includes branch user and assistant facts", async () => {
|
||||||
|
const branch: SessionEntry[] = [
|
||||||
|
createSnapshotEntry("snapshot-1", null, { text: "Stale snapshot fact" }),
|
||||||
|
createMessageEntry("user-1", "snapshot-1", createUserMessage("Goal: Fix Task 6\nPrefer keeping the public API stable", 2)),
|
||||||
|
createMessageEntry(
|
||||||
|
"assistant-1",
|
||||||
|
"user-1",
|
||||||
|
createAssistantMessage("Decision: rebuild from ctx.sessionManager.getBranch()\nNext: add integration tests", 3),
|
||||||
|
),
|
||||||
|
createMessageEntry(
|
||||||
|
"tool-1",
|
||||||
|
"assistant-1",
|
||||||
|
createToolResultMessage("Opened .pi/agent/extensions/context-manager/index.ts", 4),
|
||||||
|
),
|
||||||
|
];
|
||||||
|
|
||||||
|
const harness = createHarness(branch, { usageTokens: 120_000 });
|
||||||
|
await harness.handlers.get("session_start")?.({ type: "session_start" }, harness.ctx);
|
||||||
|
|
||||||
|
await harness.handlers.get("turn_end")?.(
|
||||||
|
{
|
||||||
|
type: "turn_end",
|
||||||
|
turnIndex: 1,
|
||||||
|
message: createAssistantMessage("done", 5),
|
||||||
|
toolResults: [],
|
||||||
|
},
|
||||||
|
harness.ctx,
|
||||||
|
);
|
||||||
|
|
||||||
|
assert.equal(harness.appendedEntries.length, 1);
|
||||||
|
assert.equal(harness.appendedEntries[0]?.customType, SNAPSHOT_ENTRY_TYPE);
|
||||||
|
|
||||||
|
const snapshot = harness.appendedEntries[0]!.data as any;
|
||||||
|
const activeTexts = snapshot.ledger.items.filter((item: any) => item.active).map((item: any) => item.text);
|
||||||
|
|
||||||
|
assert.equal(snapshot.mode, "aggressive");
|
||||||
|
assert.equal(snapshot.lastCompactionSummary, "existing compaction summary");
|
||||||
|
assert.equal(snapshot.lastBranchSummary, "existing branch summary");
|
||||||
|
assert.equal(snapshot.lastObservedTokens, 120_000);
|
||||||
|
assert.equal(snapshot.lastZone, "yellow");
|
||||||
|
assert.deepEqual(activeTexts, ["Stale snapshot fact", "Fix Task 6", "Prefer keeping the public API stable", "rebuild from ctx.sessionManager.getBranch()", "add integration tests", ".pi/agent/extensions/context-manager/index.ts"]);
|
||||||
|
assert.deepEqual(harness.statuses.at(-1), { key: "context-manager", value: "ctx yellow" });
|
||||||
|
});
|
||||||
|
|
||||||
|
test("session_tree rebuilds runtime from snapshot-only branches before injecting the next packet", async () => {
|
||||||
|
const oldBranch: SessionEntry[] = [createSnapshotEntry("snapshot-old", null, { text: "Old branch goal" })];
|
||||||
|
const newBranch: SessionEntry[] = [
|
||||||
|
createSnapshotEntry("snapshot-new", null, {
|
||||||
|
text: "Snapshot-only branch goal",
|
||||||
|
ledgerItems: [
|
||||||
|
{
|
||||||
|
id: "goal:session:root-goal:snapshot-new",
|
||||||
|
kind: "goal",
|
||||||
|
subject: "root-goal",
|
||||||
|
text: "Snapshot-only branch goal",
|
||||||
|
scope: "session",
|
||||||
|
sourceEntryId: "snapshot-new",
|
||||||
|
sourceType: "user",
|
||||||
|
timestamp: 11,
|
||||||
|
confidence: 1,
|
||||||
|
freshness: 11,
|
||||||
|
active: true,
|
||||||
|
supersedesId: undefined,
|
||||||
|
},
|
||||||
|
{
|
||||||
|
id: "decision:branch:branch-decision:snapshot-new",
|
||||||
|
kind: "decision",
|
||||||
|
subject: "branch-decision",
|
||||||
|
text: "Use the snapshot-backed branch state immediately",
|
||||||
|
scope: "branch",
|
||||||
|
sourceEntryId: "snapshot-new",
|
||||||
|
sourceType: "assistant",
|
||||||
|
timestamp: 12,
|
||||||
|
confidence: 0.9,
|
||||||
|
freshness: 12,
|
||||||
|
active: true,
|
||||||
|
supersedesId: undefined,
|
||||||
|
},
|
||||||
|
],
|
||||||
|
rollingSummary: "snapshot-only branch state",
|
||||||
|
}),
|
||||||
|
];
|
||||||
|
|
||||||
|
const harness = createHarness(oldBranch);
|
||||||
|
await harness.handlers.get("session_start")?.({ type: "session_start" }, harness.ctx);
|
||||||
|
|
||||||
|
harness.setBranch(newBranch);
|
||||||
|
await harness.handlers.get("session_tree")?.(
|
||||||
|
{
|
||||||
|
type: "session_tree",
|
||||||
|
oldLeafId: "snapshot-old",
|
||||||
|
newLeafId: "snapshot-new",
|
||||||
|
},
|
||||||
|
harness.ctx,
|
||||||
|
);
|
||||||
|
|
||||||
|
const result = await harness.handlers.get("context")?.(
|
||||||
|
{
|
||||||
|
type: "context",
|
||||||
|
messages: [createUserMessage("What should happen next?", 13)],
|
||||||
|
},
|
||||||
|
harness.ctx,
|
||||||
|
);
|
||||||
|
|
||||||
|
assert.ok(result);
|
||||||
|
assert.equal(result.messages[0]?.role, "custom");
|
||||||
|
assert.equal(result.messages[0]?.customType, "context-manager.resume");
|
||||||
|
assert.match(result.messages[0]?.content, /Snapshot-only branch goal/);
|
||||||
|
assert.match(result.messages[0]?.content, /Use the snapshot-backed branch state immediately/);
|
||||||
|
assert.doesNotMatch(result.messages[0]?.content, /Old branch goal/);
|
||||||
|
assert.deepEqual(harness.statuses.at(-1), { key: "context-manager", value: "ctx red" });
|
||||||
|
});
|
||||||
|
|
||||||
|
test("context keeps a distilled stale tool result visible after pruning bulky output", async () => {
|
||||||
|
const bulkyFailure = [
|
||||||
|
"Build failed while compiling focus parser",
|
||||||
|
"Error: missing export createFocusMatcher from ./summary-focus.ts",
|
||||||
|
...Array.from({ length: 220 }, () => "stack frame"),
|
||||||
|
].join("\n");
|
||||||
|
|
||||||
|
const harness = createHarness([]);
|
||||||
|
await harness.handlers.get("session_start")?.({ type: "session_start" }, harness.ctx);
|
||||||
|
|
||||||
|
const result = await harness.handlers.get("context")?.(
|
||||||
|
{
|
||||||
|
type: "context",
|
||||||
|
messages: [
|
||||||
|
createUserMessage("turn 1", 1),
|
||||||
|
createToolResultMessage(bulkyFailure, 2),
|
||||||
|
createAssistantMessage("observed turn 1", 3),
|
||||||
|
createUserMessage("turn 2", 4),
|
||||||
|
createAssistantMessage("observed turn 2", 5),
|
||||||
|
createUserMessage("turn 3", 6),
|
||||||
|
createAssistantMessage("observed turn 3", 7),
|
||||||
|
createUserMessage("turn 4", 8),
|
||||||
|
createAssistantMessage("observed turn 4", 9),
|
||||||
|
createUserMessage("turn 5", 10),
|
||||||
|
],
|
||||||
|
},
|
||||||
|
harness.ctx,
|
||||||
|
);
|
||||||
|
|
||||||
|
const toolResult = result.messages.find((message: any) => message.role === "toolResult");
|
||||||
|
assert.ok(toolResult);
|
||||||
|
assert.match(toolResult.content[0].text, /missing export createFocusMatcher/);
|
||||||
|
assert.ok(toolResult.content[0].text.length < 320);
|
||||||
|
});
|
||||||
|
|
||||||
|
test("session_tree preserves session-scoped facts but drops stale branch handoff metadata on an empty destination branch", async () => {
|
||||||
|
const sourceBranch: SessionEntry[] = [
|
||||||
|
createSnapshotEntry("snapshot-session", null, {
|
||||||
|
text: "Ship the context manager extension",
|
||||||
|
mode: "balanced",
|
||||||
|
lastZone: "yellow",
|
||||||
|
lastObservedTokens: 120_000,
|
||||||
|
lastCompactionSummary: "## Key Decisions\n- Keep summaries deterministic.",
|
||||||
|
lastBranchSummary: "# Handoff for branch\n\n## Key Decisions\n- Do not leak branch-local goals.",
|
||||||
|
}),
|
||||||
|
];
|
||||||
|
|
||||||
|
const harness = createHarness(sourceBranch);
|
||||||
|
await harness.handlers.get("session_start")?.({ type: "session_start" }, harness.ctx);
|
||||||
|
|
||||||
|
harness.setBranch([]);
|
||||||
|
await harness.handlers.get("session_tree")?.(
|
||||||
|
{
|
||||||
|
type: "session_tree",
|
||||||
|
oldLeafId: "snapshot-session",
|
||||||
|
newLeafId: null,
|
||||||
|
},
|
||||||
|
harness.ctx,
|
||||||
|
);
|
||||||
|
|
||||||
|
const result = await harness.handlers.get("context")?.(
|
||||||
|
{ type: "context", messages: [createUserMessage("continue", 30)] },
|
||||||
|
harness.ctx,
|
||||||
|
);
|
||||||
|
|
||||||
|
assert.equal(result.messages[0]?.customType, "context-manager.packet");
|
||||||
|
assert.match(result.messages[0]?.content, /Ship the context manager extension/);
|
||||||
|
assert.doesNotMatch(result.messages[0]?.content, /Do not leak branch-local goals/);
|
||||||
|
assert.doesNotMatch(result.messages[0]?.content, /Keep summaries deterministic/);
|
||||||
|
});
|
||||||
|
|
||||||
|
test("session_tree overlays newer session-scoped facts onto a destination branch with an older snapshot", async () => {
|
||||||
|
const newerSessionSnapshot = createSnapshotEntry("snapshot-newer", null, {
|
||||||
|
text: "Ship the context manager extension",
|
||||||
|
ledgerItems: [
|
||||||
|
{
|
||||||
|
id: "goal:session:root-goal:snapshot-newer",
|
||||||
|
kind: "goal",
|
||||||
|
subject: "root-goal",
|
||||||
|
text: "Ship the context manager extension",
|
||||||
|
scope: "session",
|
||||||
|
sourceEntryId: "snapshot-newer",
|
||||||
|
sourceType: "user",
|
||||||
|
timestamp: 1,
|
||||||
|
confidence: 1,
|
||||||
|
freshness: 1,
|
||||||
|
active: true,
|
||||||
|
supersedesId: undefined,
|
||||||
|
},
|
||||||
|
{
|
||||||
|
id: "constraint:session:must-session-newer:2",
|
||||||
|
kind: "constraint",
|
||||||
|
subject: "must-session-newer",
|
||||||
|
text: "Prefer concise reports across the whole session.",
|
||||||
|
scope: "session",
|
||||||
|
sourceEntryId: "snapshot-newer",
|
||||||
|
sourceType: "user",
|
||||||
|
timestamp: 2,
|
||||||
|
confidence: 0.9,
|
||||||
|
freshness: 2,
|
||||||
|
active: true,
|
||||||
|
supersedesId: undefined,
|
||||||
|
},
|
||||||
|
],
|
||||||
|
lastCompactionSummary: "",
|
||||||
|
lastBranchSummary: "",
|
||||||
|
});
|
||||||
|
const olderBranchSnapshot = createSnapshotEntry("snapshot-older", null, {
|
||||||
|
text: "Ship the context manager extension",
|
||||||
|
ledgerItems: [
|
||||||
|
{
|
||||||
|
id: "goal:session:root-goal:snapshot-older",
|
||||||
|
kind: "goal",
|
||||||
|
subject: "root-goal",
|
||||||
|
text: "Ship the context manager extension",
|
||||||
|
scope: "session",
|
||||||
|
sourceEntryId: "snapshot-older",
|
||||||
|
sourceType: "user",
|
||||||
|
timestamp: 1,
|
||||||
|
confidence: 1,
|
||||||
|
freshness: 1,
|
||||||
|
active: true,
|
||||||
|
supersedesId: undefined,
|
||||||
|
},
|
||||||
|
],
|
||||||
|
lastCompactionSummary: "",
|
||||||
|
lastBranchSummary: "",
|
||||||
|
});
|
||||||
|
|
||||||
|
const harness = createHarness([newerSessionSnapshot]);
|
||||||
|
await harness.handlers.get("session_start")?.({ type: "session_start" }, harness.ctx);
|
||||||
|
|
||||||
|
harness.setBranch([olderBranchSnapshot]);
|
||||||
|
await harness.handlers.get("session_tree")?.(
|
||||||
|
{
|
||||||
|
type: "session_tree",
|
||||||
|
oldLeafId: "snapshot-newer",
|
||||||
|
newLeafId: "snapshot-older",
|
||||||
|
},
|
||||||
|
harness.ctx,
|
||||||
|
);
|
||||||
|
|
||||||
|
const result = await harness.handlers.get("context")?.(
|
||||||
|
{ type: "context", messages: [createUserMessage("continue", 32)] },
|
||||||
|
harness.ctx,
|
||||||
|
);
|
||||||
|
|
||||||
|
assert.match(result.messages[0]?.content, /Prefer concise reports across the whole session/);
|
||||||
|
});
|
||||||
|
|
||||||
|
test("ctx-refresh preserves session memory without leaking old handoff summaries on a snapshot-less branch", async () => {
|
||||||
|
const sourceBranch: SessionEntry[] = [
|
||||||
|
createSnapshotEntry("snapshot-refresh", null, {
|
||||||
|
text: "Ship the context manager extension",
|
||||||
|
mode: "balanced",
|
||||||
|
lastZone: "yellow",
|
||||||
|
lastObservedTokens: 120_000,
|
||||||
|
lastCompactionSummary: "## Key Decisions\n- Keep summaries deterministic.",
|
||||||
|
lastBranchSummary: "# Handoff for branch\n\n## Key Decisions\n- Do not leak branch-local goals.",
|
||||||
|
}),
|
||||||
|
];
|
||||||
|
|
||||||
|
const harness = createHarness(sourceBranch);
|
||||||
|
await harness.handlers.get("session_start")?.({ type: "session_start" }, harness.ctx);
|
||||||
|
|
||||||
|
harness.setBranch([]);
|
||||||
|
await harness.commands.get("ctx-refresh")?.handler("", harness.ctx);
|
||||||
|
|
||||||
|
const result = await harness.handlers.get("context")?.(
|
||||||
|
{ type: "context", messages: [createUserMessage("continue", 31)] },
|
||||||
|
harness.ctx,
|
||||||
|
);
|
||||||
|
|
||||||
|
assert.equal(result.messages[0]?.customType, "context-manager.packet");
|
||||||
|
assert.match(result.messages[0]?.content, /Ship the context manager extension/);
|
||||||
|
assert.doesNotMatch(result.messages[0]?.content, /Do not leak branch-local goals/);
|
||||||
|
assert.doesNotMatch(result.messages[0]?.content, /Keep summaries deterministic/);
|
||||||
|
});
|
||||||
|
|
||||||
|
test("session_start replays default pi compaction blockers into resume state", async () => {
|
||||||
|
const branch: SessionEntry[] = [
|
||||||
|
createSnapshotEntry("snapshot-default", null, {
|
||||||
|
text: "Ship the context manager extension",
|
||||||
|
lastCompactionSummary: undefined,
|
||||||
|
lastBranchSummary: undefined,
|
||||||
|
}),
|
||||||
|
{
|
||||||
|
type: "compaction",
|
||||||
|
id: "compaction-default-1",
|
||||||
|
parentId: "snapshot-default",
|
||||||
|
timestamp: new Date(40).toISOString(),
|
||||||
|
summary: [
|
||||||
|
"## Progress",
|
||||||
|
"### Blocked",
|
||||||
|
"- confirm whether /tree replaceInstructions should override defaults",
|
||||||
|
].join("\n"),
|
||||||
|
firstKeptEntryId: "snapshot-default",
|
||||||
|
tokensBefore: 123_000,
|
||||||
|
},
|
||||||
|
];
|
||||||
|
|
||||||
|
const harness = createHarness(branch);
|
||||||
|
await harness.handlers.get("session_start")?.({ type: "session_start" }, harness.ctx);
|
||||||
|
|
||||||
|
const result = await harness.handlers.get("context")?.(
|
||||||
|
{ type: "context", messages: [createUserMessage("continue", 41)] },
|
||||||
|
harness.ctx,
|
||||||
|
);
|
||||||
|
|
||||||
|
assert.match(result.messages[0]?.content, /confirm whether \/tree replaceInstructions should override defaults/);
|
||||||
|
});
|
||||||
|
|
||||||
|
test("session_before_compact honors preparation inputs and custom focus", async () => {
|
||||||
|
const harness = createHarness([]);
|
||||||
|
await harness.handlers.get("session_start")?.({ type: "session_start" }, harness.ctx);
|
||||||
|
|
||||||
|
const result = await harness.handlers.get("session_before_compact")?.(
|
||||||
|
{
|
||||||
|
type: "session_before_compact",
|
||||||
|
customInstructions: "Focus on decisions and relevant files.",
|
||||||
|
preparation: {
|
||||||
|
messagesToSummarize: [createUserMessage("Decision: keep compaction summaries deterministic", 1)],
|
||||||
|
turnPrefixMessages: [createToolResultMessage("Opened .pi/agent/extensions/context-manager/index.ts", 2)],
|
||||||
|
previousSummary: "## Goal\n- Ship the context manager extension",
|
||||||
|
fileOps: {
|
||||||
|
readFiles: [".pi/agent/extensions/context-manager/index.ts"],
|
||||||
|
modifiedFiles: [".pi/agent/extensions/context-manager/src/summaries.ts"],
|
||||||
|
},
|
||||||
|
tokensBefore: 120_000,
|
||||||
|
firstKeptEntryId: "keep-1",
|
||||||
|
},
|
||||||
|
branchEntries: [],
|
||||||
|
signal: AbortSignal.timeout(1_000),
|
||||||
|
},
|
||||||
|
harness.ctx,
|
||||||
|
);
|
||||||
|
|
||||||
|
assert.equal(result.compaction.firstKeptEntryId, "keep-1");
|
||||||
|
assert.equal(result.compaction.tokensBefore, 120_000);
|
||||||
|
assert.match(result.compaction.summary, /keep compaction summaries deterministic/);
|
||||||
|
assert.match(result.compaction.summary, /index.ts/);
|
||||||
|
});
|
||||||
|
|
||||||
|
test("session_before_tree honors abandoned-branch entries and focus text", async () => {
|
||||||
|
const harness = createHarness([]);
|
||||||
|
await harness.handlers.get("session_start")?.({ type: "session_start" }, harness.ctx);
|
||||||
|
|
||||||
|
const result = await harness.handlers.get("session_before_tree")?.(
|
||||||
|
{
|
||||||
|
type: "session_before_tree",
|
||||||
|
preparation: {
|
||||||
|
targetId: "target-1",
|
||||||
|
oldLeafId: "old-1",
|
||||||
|
commonAncestorId: "root",
|
||||||
|
userWantsSummary: true,
|
||||||
|
customInstructions: "Focus on goals and decisions.",
|
||||||
|
replaceInstructions: false,
|
||||||
|
entriesToSummarize: [
|
||||||
|
createMessageEntry("user-1", null, createUserMessage("Goal: explore tree handoff", 1)),
|
||||||
|
createMessageEntry("assistant-1", "user-1", createAssistantMessage("Decision: do not leak branch-local goals", 2)),
|
||||||
|
],
|
||||||
|
},
|
||||||
|
signal: AbortSignal.timeout(1_000),
|
||||||
|
},
|
||||||
|
harness.ctx,
|
||||||
|
);
|
||||||
|
|
||||||
|
assert.ok(result?.summary?.summary);
|
||||||
|
assert.match(result.summary.summary, /explore tree handoff/);
|
||||||
|
assert.match(result.summary.summary, /do not leak branch-local goals/);
|
||||||
|
});
|
||||||
|
|
||||||
|
test("session_compact persists the latest compaction summary into a fresh snapshot and injects a resume packet once", async () => {
|
||||||
|
const harness = createHarness([]);
|
||||||
|
await harness.handlers.get("session_start")?.({ type: "session_start" }, harness.ctx);
|
||||||
|
|
||||||
|
await harness.handlers.get("session_compact")?.(
|
||||||
|
{
|
||||||
|
type: "session_compact",
|
||||||
|
fromExtension: true,
|
||||||
|
compactionEntry: {
|
||||||
|
type: "compaction",
|
||||||
|
id: "cmp-1",
|
||||||
|
parentId: "prev",
|
||||||
|
timestamp: new Date(10).toISOString(),
|
||||||
|
summary: "## Key Decisions\n- Keep summaries deterministic.\n\n## Open questions and blockers\n- Verify /tree replaceInstructions behavior.",
|
||||||
|
firstKeptEntryId: "keep-1",
|
||||||
|
tokensBefore: 140_000,
|
||||||
|
},
|
||||||
|
},
|
||||||
|
harness.ctx,
|
||||||
|
);
|
||||||
|
|
||||||
|
assert.equal(harness.appendedEntries.at(-1)?.customType, SNAPSHOT_ENTRY_TYPE);
|
||||||
|
|
||||||
|
const context = await harness.handlers.get("context")?.(
|
||||||
|
{ type: "context", messages: [createUserMessage("continue", 11)] },
|
||||||
|
harness.ctx,
|
||||||
|
);
|
||||||
|
assert.match(context.messages[0]?.content, /Keep summaries deterministic/);
|
||||||
|
assert.match(context.messages[0]?.content, /Verify \/tree replaceInstructions behavior/);
|
||||||
|
|
||||||
|
const nextContext = await harness.handlers.get("context")?.(
|
||||||
|
{ type: "context", messages: [createUserMessage("continue again", 12)] },
|
||||||
|
harness.ctx,
|
||||||
|
);
|
||||||
|
assert.equal(nextContext.messages[0]?.customType, "context-manager.packet");
|
||||||
|
assert.doesNotMatch(nextContext.messages[0]?.content ?? "", /## Latest compaction handoff/);
|
||||||
|
});
|
||||||
|
|
||||||
|
test("session_tree replays branch summaries newer than the latest snapshot before the next packet is injected", async () => {
|
||||||
|
const branch: SessionEntry[] = [
|
||||||
|
createSnapshotEntry("snapshot-1", null, { text: "Ship the context manager extension" }),
|
||||||
|
{
|
||||||
|
type: "branch_summary",
|
||||||
|
id: "branch-summary-1",
|
||||||
|
parentId: "snapshot-1",
|
||||||
|
timestamp: new Date(20).toISOString(),
|
||||||
|
fromId: "old-leaf",
|
||||||
|
summary: "# Handoff for branch\n\n## Key Decisions\n- Do not leak branch-local goals.",
|
||||||
|
},
|
||||||
|
];
|
||||||
|
|
||||||
|
const harness = createHarness(branch);
|
||||||
|
await harness.handlers.get("session_start")?.({ type: "session_start" }, harness.ctx);
|
||||||
|
|
||||||
|
const context = await harness.handlers.get("context")?.(
|
||||||
|
{ type: "context", messages: [createUserMessage("what next", 21)] },
|
||||||
|
harness.ctx,
|
||||||
|
);
|
||||||
|
assert.match(context.messages[0]?.content, /Do not leak branch-local goals/);
|
||||||
|
});
|
||||||
|
|
||||||
|
|
||||||
|
test("session_tree records event summaryEntry before persisting the next snapshot", async () => {
|
||||||
|
const branch: SessionEntry[] = [
|
||||||
|
createSnapshotEntry("snapshot-1", null, {
|
||||||
|
text: "Ship the context manager extension",
|
||||||
|
lastCompactionSummary: "",
|
||||||
|
lastBranchSummary: "",
|
||||||
|
}),
|
||||||
|
];
|
||||||
|
|
||||||
|
const harness = createHarness(branch);
|
||||||
|
await harness.handlers.get("session_start")?.({ type: "session_start" }, harness.ctx);
|
||||||
|
|
||||||
|
await harness.handlers.get("session_tree")?.(
|
||||||
|
{
|
||||||
|
type: "session_tree",
|
||||||
|
fromExtension: true,
|
||||||
|
summaryEntry: {
|
||||||
|
type: "branch_summary",
|
||||||
|
id: "branch-summary-event",
|
||||||
|
parentId: "snapshot-1",
|
||||||
|
timestamp: new Date(20).toISOString(),
|
||||||
|
fromId: "old-leaf",
|
||||||
|
summary: "# Handoff for branch\n\n## Key Decisions\n- Preserve the latest branch summary from the event payload.",
|
||||||
|
},
|
||||||
|
},
|
||||||
|
harness.ctx,
|
||||||
|
);
|
||||||
|
|
||||||
|
const snapshot = harness.appendedEntries.at(-1)?.data as RuntimeSnapshot | undefined;
|
||||||
|
assert.match(snapshot?.lastBranchSummary ?? "", /Preserve the latest branch summary/);
|
||||||
|
|
||||||
|
const context = await harness.handlers.get("context")?.(
|
||||||
|
{ type: "context", messages: [createUserMessage("what changed", 21)] },
|
||||||
|
harness.ctx,
|
||||||
|
);
|
||||||
|
assert.match(context.messages[0]?.content, /Preserve the latest branch summary/);
|
||||||
|
});
|
||||||
|
|
||||||
|
test("ctx-status reports mode, zone, packet size, and summary-artifact presence", async () => {
|
||||||
|
const branch = [
|
||||||
|
createSnapshotEntry("snapshot-1", null, {
|
||||||
|
text: "Ship the context manager extension",
|
||||||
|
mode: "balanced",
|
||||||
|
lastZone: "yellow",
|
||||||
|
lastObservedTokens: 120_000,
|
||||||
|
lastCompactionSummary: "## Key Decisions\n- Keep summaries deterministic.",
|
||||||
|
lastBranchSummary: "# Handoff for branch\n\n## Key Decisions\n- Do not leak branch-local goals.",
|
||||||
|
}),
|
||||||
|
];
|
||||||
|
|
||||||
|
const notifications: string[] = [];
|
||||||
|
const harness = createHarness(branch);
|
||||||
|
harness.ctx.ui.notify = (message: string) => notifications.push(message);
|
||||||
|
|
||||||
|
await harness.handlers.get("session_start")?.({ type: "session_start" }, harness.ctx);
|
||||||
|
await harness.commands.get("ctx-status")?.handler("", harness.ctx);
|
||||||
|
|
||||||
|
assert.match(notifications.at(-1) ?? "", /mode=balanced/);
|
||||||
|
assert.match(notifications.at(-1) ?? "", /zone=yellow/);
|
||||||
|
assert.match(notifications.at(-1) ?? "", /compaction=yes/);
|
||||||
|
assert.match(notifications.at(-1) ?? "", /branch=yes/);
|
||||||
|
});
|
||||||
|
|
||||||
|
test("ctx-refresh rebuilds runtime from the current branch instead of only re-rendering the packet", async () => {
|
||||||
|
const harness = createHarness([createSnapshotEntry("snapshot-1", null, { text: "Old goal" })]);
|
||||||
|
await harness.handlers.get("session_start")?.({ type: "session_start" }, harness.ctx);
|
||||||
|
|
||||||
|
harness.setBranch([
|
||||||
|
createSnapshotEntry("snapshot-2", null, {
|
||||||
|
text: "New branch goal",
|
||||||
|
lastBranchSummary: "# Handoff for branch\n\n## Key Decisions\n- Use the new branch immediately.",
|
||||||
|
}),
|
||||||
|
]);
|
||||||
|
|
||||||
|
await harness.commands.get("ctx-refresh")?.handler("", harness.ctx);
|
||||||
|
const result = await harness.handlers.get("context")?.(
|
||||||
|
{ type: "context", messages: [createUserMessage("continue", 3)] },
|
||||||
|
harness.ctx,
|
||||||
|
);
|
||||||
|
|
||||||
|
assert.match(result.messages[0]?.content, /New branch goal/);
|
||||||
|
assert.doesNotMatch(result.messages[0]?.content, /Old goal/);
|
||||||
|
});
|
||||||
|
|
||||||
|
test("ctx-mode persists the updated mode immediately without waiting for turn_end", async () => {
|
||||||
|
const branch: SessionEntry[] = [
|
||||||
|
createSnapshotEntry("snapshot-1", null, {
|
||||||
|
text: "Persist the updated mode",
|
||||||
|
mode: "balanced",
|
||||||
|
lastZone: "green",
|
||||||
|
lastObservedTokens: 90_000,
|
||||||
|
}),
|
||||||
|
];
|
||||||
|
|
||||||
|
const harness = createHarness(branch);
|
||||||
|
await harness.handlers.get("session_start")?.({ type: "session_start" }, harness.ctx);
|
||||||
|
|
||||||
|
assert.equal(deserializeLatestSnapshot(harness.ctx.sessionManager.getBranch())?.mode, "balanced");
|
||||||
|
|
||||||
|
const modeCommand = harness.commands.get("ctx-mode");
|
||||||
|
assert.ok(modeCommand);
|
||||||
|
await modeCommand.handler("aggressive", harness.ctx);
|
||||||
|
|
||||||
|
assert.equal(harness.appendedEntries.length, 1);
|
||||||
|
assert.equal(harness.appendedEntries[0]?.customType, SNAPSHOT_ENTRY_TYPE);
|
||||||
|
assert.equal(deserializeLatestSnapshot(harness.ctx.sessionManager.getBranch())?.mode, "aggressive");
|
||||||
|
});
|
||||||
|
|
||||||
|
test("ctx-mode changes survive turn_end and persist into the next snapshot", async () => {
|
||||||
|
const branch: SessionEntry[] = [
|
||||||
|
createSnapshotEntry("snapshot-1", null, {
|
||||||
|
text: "Persist the updated mode",
|
||||||
|
mode: "balanced",
|
||||||
|
lastZone: "green",
|
||||||
|
lastObservedTokens: 90_000,
|
||||||
|
}),
|
||||||
|
];
|
||||||
|
|
||||||
|
const harness = createHarness(branch, { usageTokens: 105_000 });
|
||||||
|
await harness.handlers.get("session_start")?.({ type: "session_start" }, harness.ctx);
|
||||||
|
|
||||||
|
const modeCommand = harness.commands.get("ctx-mode");
|
||||||
|
assert.ok(modeCommand);
|
||||||
|
await modeCommand.handler("aggressive", harness.ctx);
|
||||||
|
|
||||||
|
await harness.handlers.get("turn_end")?.(
|
||||||
|
{
|
||||||
|
type: "turn_end",
|
||||||
|
turnIndex: 1,
|
||||||
|
message: createAssistantMessage("done", 5),
|
||||||
|
toolResults: [],
|
||||||
|
},
|
||||||
|
harness.ctx,
|
||||||
|
);
|
||||||
|
|
||||||
|
assert.equal(harness.appendedEntries.length, 2);
|
||||||
|
|
||||||
|
const immediateSnapshot = harness.appendedEntries[0]!.data as any;
|
||||||
|
assert.equal(immediateSnapshot.mode, "aggressive");
|
||||||
|
assert.equal(immediateSnapshot.lastObservedTokens, 90_000);
|
||||||
|
assert.equal(immediateSnapshot.lastZone, "green");
|
||||||
|
|
||||||
|
const snapshot = harness.appendedEntries[1]!.data as any;
|
||||||
|
assert.equal(snapshot.mode, "aggressive");
|
||||||
|
assert.equal(snapshot.lastObservedTokens, 105_000);
|
||||||
|
assert.equal(snapshot.lastZone, "yellow");
|
||||||
|
});
|
||||||
280
.pi/agent/extensions/context-manager/src/extract.test.ts
Normal file
280
.pi/agent/extensions/context-manager/src/extract.test.ts
Normal file
@@ -0,0 +1,280 @@
|
|||||||
|
import test from "node:test";
|
||||||
|
import assert from "node:assert/strict";
|
||||||
|
import { extractCandidates } from "./extract.ts";
|
||||||
|
import { createEmptyLedger, getActiveItems, mergeCandidates } from "./ledger.ts";
|
||||||
|
|
||||||
|
test("extractCandidates pulls goals, constraints, decisions, next steps, and file references", () => {
|
||||||
|
const candidates = extractCandidates({
|
||||||
|
entryId: "u1",
|
||||||
|
role: "user",
|
||||||
|
text: [
|
||||||
|
"Goal: Build a context manager extension for pi.",
|
||||||
|
"We must adapt to the active model context window.",
|
||||||
|
"Decision: keep the MVP quiet and avoid new LLM-facing tools.",
|
||||||
|
"Next: inspect .pi/agent/extensions/web-search/index.ts and docs/extensions.md.",
|
||||||
|
].join("\n"),
|
||||||
|
timestamp: 1,
|
||||||
|
});
|
||||||
|
|
||||||
|
assert.deepEqual(
|
||||||
|
candidates.map((candidate) => [
|
||||||
|
candidate.kind,
|
||||||
|
candidate.subject,
|
||||||
|
candidate.scope,
|
||||||
|
candidate.sourceEntryId,
|
||||||
|
candidate.sourceType,
|
||||||
|
candidate.timestamp,
|
||||||
|
]),
|
||||||
|
[
|
||||||
|
["goal", "root-goal", "session", "u1", "user", 1],
|
||||||
|
["constraint", "must-u1-0", "branch", "u1", "user", 1],
|
||||||
|
["decision", "decision-u1-0", "branch", "u1", "user", 1],
|
||||||
|
["activeTask", "next-step-u1-0", "branch", "u1", "user", 1],
|
||||||
|
["relevantFile", ".pi/agent/extensions/web-search/index.ts", "branch", "u1", "user", 1],
|
||||||
|
["relevantFile", "docs/extensions.md", "branch", "u1", "user", 1],
|
||||||
|
],
|
||||||
|
);
|
||||||
|
});
|
||||||
|
|
||||||
|
test("extractCandidates promotes only the first durable goal to session scope", () => {
|
||||||
|
const firstGoal = extractCandidates(
|
||||||
|
{
|
||||||
|
entryId: "u-goal-1",
|
||||||
|
role: "user",
|
||||||
|
text: "Goal: Ship the context manager extension.",
|
||||||
|
timestamp: 10,
|
||||||
|
},
|
||||||
|
{ hasSessionGoal: false },
|
||||||
|
);
|
||||||
|
|
||||||
|
const branchGoal = extractCandidates(
|
||||||
|
{
|
||||||
|
entryId: "u-goal-2",
|
||||||
|
role: "user",
|
||||||
|
text: "Goal: prototype a branch-local tree handoff.",
|
||||||
|
timestamp: 11,
|
||||||
|
},
|
||||||
|
{ hasSessionGoal: true },
|
||||||
|
);
|
||||||
|
|
||||||
|
assert.deepEqual(
|
||||||
|
firstGoal.map((candidate) => [candidate.kind, candidate.subject, candidate.scope, candidate.text]),
|
||||||
|
[["goal", "root-goal", "session", "Ship the context manager extension."]],
|
||||||
|
);
|
||||||
|
assert.deepEqual(
|
||||||
|
branchGoal.map((candidate) => [candidate.kind, candidate.subject, candidate.scope, candidate.text]),
|
||||||
|
[["goal", "goal-u-goal-2-0", "branch", "prototype a branch-local tree handoff."]],
|
||||||
|
);
|
||||||
|
});
|
||||||
|
|
||||||
|
test("mergeCandidates keeps independently extracted decisions, constraints, and next steps active", () => {
|
||||||
|
const ledger = mergeCandidates(createEmptyLedger(), [
|
||||||
|
...extractCandidates({
|
||||||
|
entryId: "u1",
|
||||||
|
role: "user",
|
||||||
|
text: [
|
||||||
|
"We must adapt to the active model context window.",
|
||||||
|
"Decision: keep snapshots tiny.",
|
||||||
|
"Next: inspect src/extract.ts.",
|
||||||
|
].join("\n"),
|
||||||
|
timestamp: 1,
|
||||||
|
}),
|
||||||
|
...extractCandidates({
|
||||||
|
entryId: "u2",
|
||||||
|
role: "user",
|
||||||
|
text: [
|
||||||
|
"We prefer concise reports across the whole session.",
|
||||||
|
"Decision: persist snapshots after each turn_end.",
|
||||||
|
"Task: add regression coverage.",
|
||||||
|
].join("\n"),
|
||||||
|
timestamp: 2,
|
||||||
|
}),
|
||||||
|
]);
|
||||||
|
|
||||||
|
assert.deepEqual(
|
||||||
|
getActiveItems(ledger, "constraint").map((item) => [item.subject, item.sourceEntryId, item.text]),
|
||||||
|
[
|
||||||
|
["must-u1-0", "u1", "We must adapt to the active model context window."],
|
||||||
|
["must-u2-0", "u2", "We prefer concise reports across the whole session."],
|
||||||
|
],
|
||||||
|
);
|
||||||
|
assert.deepEqual(
|
||||||
|
getActiveItems(ledger, "decision").map((item) => [item.subject, item.sourceEntryId, item.text]),
|
||||||
|
[
|
||||||
|
["decision-u1-0", "u1", "keep snapshots tiny."],
|
||||||
|
["decision-u2-0", "u2", "persist snapshots after each turn_end."],
|
||||||
|
],
|
||||||
|
);
|
||||||
|
assert.deepEqual(
|
||||||
|
getActiveItems(ledger, "activeTask").map((item) => [item.subject, item.sourceEntryId, item.text]),
|
||||||
|
[
|
||||||
|
["next-step-u1-0", "u1", "inspect src/extract.ts."],
|
||||||
|
["next-step-u2-0", "u2", "add regression coverage."],
|
||||||
|
],
|
||||||
|
);
|
||||||
|
});
|
||||||
|
|
||||||
|
test("user constraints default to branch scope unless they explicitly signal durable session scope", () => {
|
||||||
|
const candidates = extractCandidates({
|
||||||
|
entryId: "u3",
|
||||||
|
role: "user",
|
||||||
|
text: [
|
||||||
|
"We should keep this branch experimental for now.",
|
||||||
|
"We should keep the MVP branch experimental.",
|
||||||
|
"We should rename the context window helper in this module.",
|
||||||
|
"Avoid touching docs/extensions.md.",
|
||||||
|
"Avoid touching docs/extensions.md across the whole session.",
|
||||||
|
"Prefer concise reports across the whole session.",
|
||||||
|
].join("\n"),
|
||||||
|
timestamp: 3,
|
||||||
|
});
|
||||||
|
|
||||||
|
assert.deepEqual(
|
||||||
|
candidates
|
||||||
|
.filter((candidate) => candidate.kind === "constraint")
|
||||||
|
.map((candidate) => [candidate.text, candidate.scope, candidate.subject]),
|
||||||
|
[
|
||||||
|
["We should keep this branch experimental for now.", "branch", "must-u3-0"],
|
||||||
|
["We should keep the MVP branch experimental.", "branch", "must-u3-1"],
|
||||||
|
["We should rename the context window helper in this module.", "branch", "must-u3-2"],
|
||||||
|
["Avoid touching docs/extensions.md.", "branch", "must-u3-3"],
|
||||||
|
["Avoid touching docs/extensions.md across the whole session.", "session", "must-u3-4"],
|
||||||
|
["Prefer concise reports across the whole session.", "session", "must-u3-5"],
|
||||||
|
],
|
||||||
|
);
|
||||||
|
});
|
||||||
|
|
||||||
|
test("extractCandidates treats spelled-out do not as a constraint trigger", () => {
|
||||||
|
const candidates = extractCandidates({
|
||||||
|
entryId: "u4",
|
||||||
|
role: "user",
|
||||||
|
text: "Do not add new LLM-facing tools across the whole session.",
|
||||||
|
timestamp: 4,
|
||||||
|
});
|
||||||
|
|
||||||
|
assert.deepEqual(
|
||||||
|
candidates.map((candidate) => [candidate.kind, candidate.text, candidate.scope, candidate.subject]),
|
||||||
|
[["constraint", "Do not add new LLM-facing tools across the whole session.", "session", "must-u4-0"]],
|
||||||
|
);
|
||||||
|
});
|
||||||
|
|
||||||
|
test("extractCandidates keeps compaction goals branch-scoped unless they are explicitly session-wide", () => {
|
||||||
|
const candidates = extractCandidates(
|
||||||
|
{
|
||||||
|
entryId: "cmp-goal-1",
|
||||||
|
role: "compaction",
|
||||||
|
text: "## Goal\n- prototype a branch-local tree handoff.",
|
||||||
|
timestamp: 19,
|
||||||
|
},
|
||||||
|
{ hasSessionGoal: false },
|
||||||
|
);
|
||||||
|
|
||||||
|
assert.deepEqual(
|
||||||
|
candidates.map((candidate) => [candidate.kind, candidate.subject, candidate.scope, candidate.text]),
|
||||||
|
[["goal", "goal-cmp-goal-1-0", "branch", "prototype a branch-local tree handoff."]],
|
||||||
|
);
|
||||||
|
});
|
||||||
|
|
||||||
|
test("extractCandidates captures blockers from direct lines, tool errors, and structured summaries", () => {
|
||||||
|
const direct = extractCandidates(
|
||||||
|
{
|
||||||
|
entryId: "a-blocked-1",
|
||||||
|
role: "assistant",
|
||||||
|
text: "Blocked: confirm whether /tree summaries should replace instructions.",
|
||||||
|
timestamp: 20,
|
||||||
|
},
|
||||||
|
{ hasSessionGoal: true },
|
||||||
|
);
|
||||||
|
|
||||||
|
const tool = extractCandidates(
|
||||||
|
{
|
||||||
|
entryId: "t-blocked-1",
|
||||||
|
role: "toolResult",
|
||||||
|
text: "Error: missing export createFocusMatcher\nstack...",
|
||||||
|
timestamp: 21,
|
||||||
|
isError: true,
|
||||||
|
},
|
||||||
|
{ hasSessionGoal: true },
|
||||||
|
);
|
||||||
|
|
||||||
|
const summary = extractCandidates(
|
||||||
|
{
|
||||||
|
entryId: "cmp-1",
|
||||||
|
role: "compaction",
|
||||||
|
text: [
|
||||||
|
"## Open questions and blockers",
|
||||||
|
"- Need to confirm whether /tree summaries should replace instructions.",
|
||||||
|
"",
|
||||||
|
"## Relevant files",
|
||||||
|
"- .pi/agent/extensions/context-manager/index.ts",
|
||||||
|
].join("\n"),
|
||||||
|
timestamp: 22,
|
||||||
|
},
|
||||||
|
{ hasSessionGoal: true },
|
||||||
|
);
|
||||||
|
|
||||||
|
assert.deepEqual(
|
||||||
|
direct.map((candidate) => [candidate.kind, candidate.subject, candidate.text]),
|
||||||
|
[["openQuestion", "open-question-a-blocked-1-0", "confirm whether /tree summaries should replace instructions."]],
|
||||||
|
);
|
||||||
|
assert.equal(tool[0]?.kind, "openQuestion");
|
||||||
|
assert.match(tool[0]?.text ?? "", /missing export createFocusMatcher/);
|
||||||
|
assert.deepEqual(
|
||||||
|
summary.map((candidate) => [candidate.kind, candidate.text]),
|
||||||
|
[
|
||||||
|
["openQuestion", "Need to confirm whether /tree summaries should replace instructions."],
|
||||||
|
["relevantFile", ".pi/agent/extensions/context-manager/index.ts"],
|
||||||
|
],
|
||||||
|
);
|
||||||
|
});
|
||||||
|
|
||||||
|
test("extractCandidates parses pi fallback progress and blocked summary sections", () => {
|
||||||
|
const candidates = extractCandidates(
|
||||||
|
{
|
||||||
|
entryId: "cmp-default-1",
|
||||||
|
role: "compaction",
|
||||||
|
text: [
|
||||||
|
"## Constraints and preferences",
|
||||||
|
"- Keep the public API stable.",
|
||||||
|
"",
|
||||||
|
"## Progress",
|
||||||
|
"### In Progress",
|
||||||
|
"- Wire runtime hydration.",
|
||||||
|
"",
|
||||||
|
"### Blocked",
|
||||||
|
"- confirm whether /tree replaceInstructions should override defaults",
|
||||||
|
].join("\n"),
|
||||||
|
timestamp: 23,
|
||||||
|
},
|
||||||
|
{ hasSessionGoal: true },
|
||||||
|
);
|
||||||
|
|
||||||
|
assert.ok(candidates.some((candidate) => candidate.kind === "constraint" && candidate.text === "Keep the public API stable."));
|
||||||
|
assert.ok(candidates.some((candidate) => candidate.kind === "activeTask" && candidate.text === "Wire runtime hydration."));
|
||||||
|
assert.ok(
|
||||||
|
candidates.some(
|
||||||
|
(candidate) => candidate.kind === "openQuestion" && candidate.text === "confirm whether /tree replaceInstructions should override defaults",
|
||||||
|
),
|
||||||
|
);
|
||||||
|
});
|
||||||
|
|
||||||
|
test("assistant decisions and tool-result file references are extracted as branch facts", () => {
|
||||||
|
const assistant = extractCandidates({
|
||||||
|
entryId: "a1",
|
||||||
|
role: "assistant",
|
||||||
|
text: "Decision: persist snapshots after each turn_end.",
|
||||||
|
timestamp: 2,
|
||||||
|
});
|
||||||
|
|
||||||
|
const tool = extractCandidates({
|
||||||
|
entryId: "t1",
|
||||||
|
role: "toolResult",
|
||||||
|
text: "Updated file: .pi/agent/extensions/context-manager/src/runtime.ts",
|
||||||
|
timestamp: 3,
|
||||||
|
});
|
||||||
|
|
||||||
|
assert.equal(assistant[0]?.kind, "decision");
|
||||||
|
assert.equal(assistant[0]?.subject, "decision-a1-0");
|
||||||
|
assert.equal(assistant[0]?.scope, "branch");
|
||||||
|
assert.equal(tool[0]?.kind, "relevantFile");
|
||||||
|
});
|
||||||
314
.pi/agent/extensions/context-manager/src/extract.ts
Normal file
314
.pi/agent/extensions/context-manager/src/extract.ts
Normal file
@@ -0,0 +1,314 @@
|
|||||||
|
import type { MemoryCandidate, MemoryScope, MemorySourceType } from "./ledger.ts";
|
||||||
|
|
||||||
|
export interface ExtractOptions {
|
||||||
|
hasSessionGoal?: boolean;
|
||||||
|
}
|
||||||
|
|
||||||
|
export interface TranscriptSlice {
|
||||||
|
entryId: string;
|
||||||
|
role: "user" | "assistant" | "toolResult" | "compaction" | "branchSummary";
|
||||||
|
text: string;
|
||||||
|
timestamp: number;
|
||||||
|
isError?: boolean;
|
||||||
|
}
|
||||||
|
|
||||||
|
const FILE_RE = /(?:\.?\/?[A-Za-z0-9_./-]+\.(?:ts|tsx|js|mjs|json|md))/g;
|
||||||
|
const BRANCH_LOCAL_CONSTRAINT_RE =
|
||||||
|
/\b(?:this|current)\s+(?:branch|task|change|step|file|module|test|command|implementation|worktree)\b|\b(?:for now|right now|in this branch|on this branch|in this file|in this module|here)\b/i;
|
||||||
|
const DURABLE_SESSION_CONSTRAINT_RE =
|
||||||
|
/\b(?:whole|entire|rest of (?:the )?|remaining)\s+(?:session|project|codebase)\b|\bacross (?:the )?(?:whole )?(?:session|project|codebase)\b|\bacross (?:all |every )?branches\b|\b(?:session|project|codebase)[-\s]?wide\b|\bthroughout (?:the )?(?:session|project|codebase)\b/i;
|
||||||
|
const CONSTRAINT_RE = /\b(?:must|should|don't|do not|avoid|prefer)\b/i;
|
||||||
|
const GOAL_RE = /^(goal|session goal|overall goal):/i;
|
||||||
|
const OPEN_QUESTION_RE = /^(?:blocked|blocker|open question|question):/i;
|
||||||
|
const ERROR_LINE_RE = /\b(?:error|failed|failure|missing|undefined|exception)\b/i;
|
||||||
|
|
||||||
|
type SummarySectionKind = "goal" | "constraint" | "decision" | "activeTask" | "openQuestion" | "relevantFile";
|
||||||
|
|
||||||
|
function sourceTypeForRole(role: TranscriptSlice["role"]): MemorySourceType {
|
||||||
|
if (role === "compaction") return "compaction";
|
||||||
|
if (role === "branchSummary") return "branchSummary";
|
||||||
|
return role;
|
||||||
|
}
|
||||||
|
|
||||||
|
function pushCandidate(
|
||||||
|
list: MemoryCandidate[],
|
||||||
|
candidate: Omit<MemoryCandidate, "sourceEntryId" | "sourceType" | "timestamp">,
|
||||||
|
slice: TranscriptSlice,
|
||||||
|
) {
|
||||||
|
list.push({
|
||||||
|
...candidate,
|
||||||
|
sourceEntryId: slice.entryId,
|
||||||
|
sourceType: sourceTypeForRole(slice.role),
|
||||||
|
timestamp: slice.timestamp,
|
||||||
|
});
|
||||||
|
}
|
||||||
|
|
||||||
|
function createIndexedSubject(prefix: string, slice: TranscriptSlice, index: number): string {
|
||||||
|
return `${prefix}-${slice.entryId}-${index}`;
|
||||||
|
}
|
||||||
|
|
||||||
|
function inferConstraintScope(slice: TranscriptSlice, line: string): MemoryScope {
|
||||||
|
if (slice.role !== "user") {
|
||||||
|
return "branch";
|
||||||
|
}
|
||||||
|
|
||||||
|
if (BRANCH_LOCAL_CONSTRAINT_RE.test(line)) {
|
||||||
|
return "branch";
|
||||||
|
}
|
||||||
|
|
||||||
|
if (DURABLE_SESSION_CONSTRAINT_RE.test(line)) {
|
||||||
|
return "session";
|
||||||
|
}
|
||||||
|
|
||||||
|
if ((line.match(FILE_RE) ?? []).length > 0) {
|
||||||
|
return "branch";
|
||||||
|
}
|
||||||
|
|
||||||
|
return "branch";
|
||||||
|
}
|
||||||
|
|
||||||
|
function nextGoalCandidate(
|
||||||
|
line: string,
|
||||||
|
slice: TranscriptSlice,
|
||||||
|
options: ExtractOptions,
|
||||||
|
index: number,
|
||||||
|
): Omit<MemoryCandidate, "sourceEntryId" | "sourceType" | "timestamp"> {
|
||||||
|
const text = line.replace(GOAL_RE, "").trim();
|
||||||
|
const explicitSessionGoal = /^(session goal|overall goal):/i.test(line);
|
||||||
|
const canSeedSessionGoal = slice.role === "user";
|
||||||
|
const shouldPromoteRootGoal = explicitSessionGoal || (!options.hasSessionGoal && canSeedSessionGoal);
|
||||||
|
|
||||||
|
if (shouldPromoteRootGoal) {
|
||||||
|
return {
|
||||||
|
kind: "goal",
|
||||||
|
subject: "root-goal",
|
||||||
|
text,
|
||||||
|
scope: "session",
|
||||||
|
confidence: 1,
|
||||||
|
};
|
||||||
|
}
|
||||||
|
|
||||||
|
return {
|
||||||
|
kind: "goal",
|
||||||
|
subject: createIndexedSubject("goal", slice, index),
|
||||||
|
text,
|
||||||
|
scope: "branch",
|
||||||
|
confidence: 0.9,
|
||||||
|
};
|
||||||
|
}
|
||||||
|
|
||||||
|
function nextOpenQuestionCandidate(
|
||||||
|
text: string,
|
||||||
|
slice: TranscriptSlice,
|
||||||
|
index: number,
|
||||||
|
): Omit<MemoryCandidate, "sourceEntryId" | "sourceType" | "timestamp"> {
|
||||||
|
return {
|
||||||
|
kind: "openQuestion",
|
||||||
|
subject: createIndexedSubject("open-question", slice, index),
|
||||||
|
text,
|
||||||
|
scope: "branch",
|
||||||
|
confidence: slice.role === "toolResult" ? 0.85 : 0.8,
|
||||||
|
};
|
||||||
|
}
|
||||||
|
|
||||||
|
function summarySectionToKind(line: string): SummarySectionKind | undefined {
|
||||||
|
const heading = line.replace(/^##\s+/i, "").trim().toLowerCase();
|
||||||
|
|
||||||
|
if (heading === "goal") return "goal";
|
||||||
|
if (heading === "constraints" || heading === "constraints & preferences" || heading === "constraints and preferences") {
|
||||||
|
return "constraint";
|
||||||
|
}
|
||||||
|
if (heading === "decisions" || heading === "key decisions") return "decision";
|
||||||
|
if (heading === "active work" || heading === "next steps" || heading === "current task" || heading === "progress") {
|
||||||
|
return "activeTask";
|
||||||
|
}
|
||||||
|
if (heading === "open questions and blockers" || heading === "open questions / blockers") return "openQuestion";
|
||||||
|
if (heading === "relevant files" || heading === "critical context") return "relevantFile";
|
||||||
|
|
||||||
|
return undefined;
|
||||||
|
}
|
||||||
|
|
||||||
|
function pushRelevantFiles(list: MemoryCandidate[], slice: TranscriptSlice, line: string) {
|
||||||
|
const fileMatches = line.match(FILE_RE) ?? [];
|
||||||
|
for (const match of fileMatches) {
|
||||||
|
pushCandidate(
|
||||||
|
list,
|
||||||
|
{
|
||||||
|
kind: "relevantFile",
|
||||||
|
subject: match,
|
||||||
|
text: match,
|
||||||
|
scope: "branch",
|
||||||
|
confidence: 0.7,
|
||||||
|
},
|
||||||
|
slice,
|
||||||
|
);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
export function extractCandidates(slice: TranscriptSlice, options: ExtractOptions = {}): MemoryCandidate[] {
|
||||||
|
const out: MemoryCandidate[] = [];
|
||||||
|
const lines = slice.text
|
||||||
|
.split(/\n+/)
|
||||||
|
.map((line) => line.trim())
|
||||||
|
.filter(Boolean);
|
||||||
|
|
||||||
|
let currentSection: SummarySectionKind | undefined;
|
||||||
|
let goalIndex = 0;
|
||||||
|
let decisionIndex = 0;
|
||||||
|
let nextStepIndex = 0;
|
||||||
|
let mustIndex = 0;
|
||||||
|
let openQuestionIndex = 0;
|
||||||
|
|
||||||
|
for (const line of lines) {
|
||||||
|
if (/^##\s+/i.test(line)) {
|
||||||
|
currentSection = summarySectionToKind(line);
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
|
||||||
|
if (/^###\s+/i.test(line)) {
|
||||||
|
const subheading = line.replace(/^###\s+/i, "").trim().toLowerCase();
|
||||||
|
if (subheading === "blocked") {
|
||||||
|
currentSection = "openQuestion";
|
||||||
|
} else if (subheading === "in progress" || subheading === "done") {
|
||||||
|
currentSection = "activeTask";
|
||||||
|
}
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
|
||||||
|
const bullet = line.match(/^-\s+(.*)$/)?.[1]?.trim();
|
||||||
|
const isGoal = GOAL_RE.test(line);
|
||||||
|
const isDecision = /^decision:/i.test(line);
|
||||||
|
const isNextStep = /^(next|task):/i.test(line);
|
||||||
|
const isOpenQuestion = OPEN_QUESTION_RE.test(line);
|
||||||
|
|
||||||
|
if (isGoal) {
|
||||||
|
pushCandidate(out, nextGoalCandidate(line, slice, options, goalIndex++), slice);
|
||||||
|
pushRelevantFiles(out, slice, line);
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
|
||||||
|
if (isOpenQuestion) {
|
||||||
|
pushCandidate(
|
||||||
|
out,
|
||||||
|
nextOpenQuestionCandidate(line.replace(OPEN_QUESTION_RE, "").trim(), slice, openQuestionIndex++),
|
||||||
|
slice,
|
||||||
|
);
|
||||||
|
pushRelevantFiles(out, slice, line);
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
|
||||||
|
if (isDecision) {
|
||||||
|
pushCandidate(
|
||||||
|
out,
|
||||||
|
{
|
||||||
|
kind: "decision",
|
||||||
|
subject: createIndexedSubject("decision", slice, decisionIndex++),
|
||||||
|
text: line.replace(/^decision:\s*/i, "").trim(),
|
||||||
|
scope: "branch",
|
||||||
|
confidence: 0.9,
|
||||||
|
},
|
||||||
|
slice,
|
||||||
|
);
|
||||||
|
pushRelevantFiles(out, slice, line);
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
|
||||||
|
if (isNextStep) {
|
||||||
|
pushCandidate(
|
||||||
|
out,
|
||||||
|
{
|
||||||
|
kind: "activeTask",
|
||||||
|
subject: createIndexedSubject("next-step", slice, nextStepIndex++),
|
||||||
|
text: line.replace(/^(next|task):\s*/i, "").trim(),
|
||||||
|
scope: "branch",
|
||||||
|
confidence: 0.8,
|
||||||
|
},
|
||||||
|
slice,
|
||||||
|
);
|
||||||
|
pushRelevantFiles(out, slice, line);
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
|
||||||
|
if (bullet && currentSection === "goal") {
|
||||||
|
pushCandidate(out, nextGoalCandidate(`Goal: ${bullet}`, slice, options, goalIndex++), slice);
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
|
||||||
|
if (bullet && currentSection === "constraint") {
|
||||||
|
pushCandidate(
|
||||||
|
out,
|
||||||
|
{
|
||||||
|
kind: "constraint",
|
||||||
|
subject: createIndexedSubject("must", slice, mustIndex++),
|
||||||
|
text: bullet,
|
||||||
|
scope: inferConstraintScope(slice, bullet),
|
||||||
|
confidence: 0.8,
|
||||||
|
},
|
||||||
|
slice,
|
||||||
|
);
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
|
||||||
|
if (bullet && currentSection === "decision") {
|
||||||
|
pushCandidate(
|
||||||
|
out,
|
||||||
|
{
|
||||||
|
kind: "decision",
|
||||||
|
subject: createIndexedSubject("decision", slice, decisionIndex++),
|
||||||
|
text: bullet,
|
||||||
|
scope: "branch",
|
||||||
|
confidence: 0.9,
|
||||||
|
},
|
||||||
|
slice,
|
||||||
|
);
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
|
||||||
|
if (bullet && currentSection === "activeTask") {
|
||||||
|
pushCandidate(
|
||||||
|
out,
|
||||||
|
{
|
||||||
|
kind: "activeTask",
|
||||||
|
subject: createIndexedSubject("next-step", slice, nextStepIndex++),
|
||||||
|
text: bullet,
|
||||||
|
scope: "branch",
|
||||||
|
confidence: 0.8,
|
||||||
|
},
|
||||||
|
slice,
|
||||||
|
);
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
|
||||||
|
if (bullet && currentSection === "openQuestion") {
|
||||||
|
pushCandidate(out, nextOpenQuestionCandidate(bullet, slice, openQuestionIndex++), slice);
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
|
||||||
|
if (bullet && currentSection === "relevantFile") {
|
||||||
|
pushRelevantFiles(out, slice, bullet);
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
|
||||||
|
if (slice.role === "toolResult" && (slice.isError || ERROR_LINE_RE.test(line))) {
|
||||||
|
pushCandidate(out, nextOpenQuestionCandidate(line, slice, openQuestionIndex++), slice);
|
||||||
|
}
|
||||||
|
|
||||||
|
if (CONSTRAINT_RE.test(line)) {
|
||||||
|
pushCandidate(
|
||||||
|
out,
|
||||||
|
{
|
||||||
|
kind: "constraint",
|
||||||
|
subject: createIndexedSubject("must", slice, mustIndex++),
|
||||||
|
text: line,
|
||||||
|
scope: inferConstraintScope(slice, line),
|
||||||
|
confidence: 0.8,
|
||||||
|
},
|
||||||
|
slice,
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
pushRelevantFiles(out, slice, line);
|
||||||
|
}
|
||||||
|
|
||||||
|
return out;
|
||||||
|
}
|
||||||
132
.pi/agent/extensions/context-manager/src/ledger.test.ts
Normal file
132
.pi/agent/extensions/context-manager/src/ledger.test.ts
Normal file
@@ -0,0 +1,132 @@
|
|||||||
|
import test from "node:test";
|
||||||
|
import assert from "node:assert/strict";
|
||||||
|
import { createEmptyLedger, getActiveItems, mergeCandidates, type MemoryCandidate } from "./ledger.ts";
|
||||||
|
|
||||||
|
const base: Omit<MemoryCandidate, "kind" | "subject" | "text"> = {
|
||||||
|
scope: "branch",
|
||||||
|
sourceEntryId: "u1",
|
||||||
|
sourceType: "user",
|
||||||
|
timestamp: 1,
|
||||||
|
confidence: 0.9,
|
||||||
|
};
|
||||||
|
|
||||||
|
test("mergeCandidates adds new active items to an empty ledger", () => {
|
||||||
|
const ledger = mergeCandidates(createEmptyLedger(), [
|
||||||
|
{ ...base, kind: "goal", subject: "root-goal", text: "Build a pi context manager extension" },
|
||||||
|
]);
|
||||||
|
|
||||||
|
assert.equal(getActiveItems(ledger).length, 1);
|
||||||
|
assert.equal(getActiveItems(ledger)[0]?.text, "Build a pi context manager extension");
|
||||||
|
});
|
||||||
|
|
||||||
|
test("mergeCandidates archives older items when a new item supersedes the same subject", () => {
|
||||||
|
const first = mergeCandidates(createEmptyLedger(), [
|
||||||
|
{ ...base, kind: "decision", subject: "persistence", text: "Persist snapshots with appendEntry()", timestamp: 1 },
|
||||||
|
]);
|
||||||
|
|
||||||
|
const second = mergeCandidates(first, [
|
||||||
|
{ ...base, kind: "decision", subject: "persistence", text: "Persist snapshots after each turn_end", timestamp: 2 },
|
||||||
|
]);
|
||||||
|
|
||||||
|
const active = getActiveItems(second, "decision");
|
||||||
|
assert.equal(active.length, 1);
|
||||||
|
assert.equal(active[0]?.text, "Persist snapshots after each turn_end");
|
||||||
|
assert.equal(active[0]?.supersedesId, "decision:branch:persistence:1");
|
||||||
|
assert.equal(second.items.find((item) => item.id === "decision:branch:persistence:1")?.active, false);
|
||||||
|
});
|
||||||
|
|
||||||
|
test("mergeCandidates keeps the newest item active when same-slot candidates arrive out of order", () => {
|
||||||
|
const ledger = mergeCandidates(createEmptyLedger(), [
|
||||||
|
{ ...base, kind: "decision", subject: "persistence", text: "Persist snapshots after each turn_end", timestamp: 2 },
|
||||||
|
{ ...base, kind: "decision", subject: "persistence", text: "Persist snapshots with appendEntry()", timestamp: 1 },
|
||||||
|
]);
|
||||||
|
|
||||||
|
const active = getActiveItems(ledger, "decision");
|
||||||
|
const stale = ledger.items.find((item) => item.text === "Persist snapshots with appendEntry()");
|
||||||
|
|
||||||
|
assert.equal(active.length, 1);
|
||||||
|
assert.equal(active[0]?.text, "Persist snapshots after each turn_end");
|
||||||
|
assert.equal(stale?.active, false);
|
||||||
|
assert.equal(stale?.supersedesId, undefined);
|
||||||
|
});
|
||||||
|
|
||||||
|
test("mergeCandidates gives same-slot same-timestamp candidates distinct ids", () => {
|
||||||
|
const ledger = mergeCandidates(createEmptyLedger(), [
|
||||||
|
{ ...base, kind: "decision", subject: "persistence", text: "Persist snapshots with appendEntry()", timestamp: 1 },
|
||||||
|
{ ...base, kind: "decision", subject: "persistence", text: "Persist snapshots after each turn_end", timestamp: 1 },
|
||||||
|
]);
|
||||||
|
|
||||||
|
const ids = ledger.items.map((item) => item.id);
|
||||||
|
const active = getActiveItems(ledger, "decision")[0];
|
||||||
|
const archived = ledger.items.find((item) => item.text === "Persist snapshots with appendEntry()");
|
||||||
|
|
||||||
|
assert.equal(new Set(ids).size, ledger.items.length);
|
||||||
|
assert.equal(active?.text, "Persist snapshots after each turn_end");
|
||||||
|
assert.equal(active?.supersedesId, archived?.id);
|
||||||
|
assert.notEqual(active?.id, active?.supersedesId);
|
||||||
|
});
|
||||||
|
|
||||||
|
test("mergeCandidates keeps same-slot same-timestamp snapshots deterministic regardless of input order", () => {
|
||||||
|
const appendEntryCandidate = {
|
||||||
|
...base,
|
||||||
|
kind: "decision" as const,
|
||||||
|
subject: "persistence",
|
||||||
|
text: "Persist snapshots with appendEntry()",
|
||||||
|
timestamp: 1,
|
||||||
|
};
|
||||||
|
const turnEndCandidate = {
|
||||||
|
...base,
|
||||||
|
kind: "decision" as const,
|
||||||
|
subject: "persistence",
|
||||||
|
text: "Persist snapshots after each turn_end",
|
||||||
|
timestamp: 1,
|
||||||
|
};
|
||||||
|
|
||||||
|
const forward = mergeCandidates(createEmptyLedger(), [appendEntryCandidate, turnEndCandidate]);
|
||||||
|
const reversed = mergeCandidates(createEmptyLedger(), [turnEndCandidate, appendEntryCandidate]);
|
||||||
|
|
||||||
|
assert.deepEqual(forward, reversed);
|
||||||
|
assert.deepEqual(forward.items, [
|
||||||
|
{
|
||||||
|
...turnEndCandidate,
|
||||||
|
id: "decision:branch:persistence:1",
|
||||||
|
freshness: 1,
|
||||||
|
active: true,
|
||||||
|
supersedesId: "decision:branch:persistence:1:2",
|
||||||
|
},
|
||||||
|
{
|
||||||
|
...appendEntryCandidate,
|
||||||
|
id: "decision:branch:persistence:1:2",
|
||||||
|
freshness: 1,
|
||||||
|
active: false,
|
||||||
|
supersedesId: undefined,
|
||||||
|
},
|
||||||
|
]);
|
||||||
|
});
|
||||||
|
|
||||||
|
test("session-scoped memory can coexist with branch-scoped memory for the same kind", () => {
|
||||||
|
const ledger = mergeCandidates(createEmptyLedger(), [
|
||||||
|
{
|
||||||
|
kind: "constraint",
|
||||||
|
subject: "llm-tools",
|
||||||
|
text: "Do not add new LLM-facing tools in the MVP",
|
||||||
|
scope: "session",
|
||||||
|
sourceEntryId: "u1",
|
||||||
|
sourceType: "user",
|
||||||
|
timestamp: 1,
|
||||||
|
confidence: 1,
|
||||||
|
},
|
||||||
|
{
|
||||||
|
kind: "constraint",
|
||||||
|
subject: "branch-policy",
|
||||||
|
text: "Keep branch A experimental",
|
||||||
|
scope: "branch",
|
||||||
|
sourceEntryId: "u2",
|
||||||
|
sourceType: "user",
|
||||||
|
timestamp: 2,
|
||||||
|
confidence: 0.8,
|
||||||
|
},
|
||||||
|
]);
|
||||||
|
|
||||||
|
assert.equal(getActiveItems(ledger, "constraint").length, 2);
|
||||||
|
});
|
||||||
196
.pi/agent/extensions/context-manager/src/ledger.ts
Normal file
196
.pi/agent/extensions/context-manager/src/ledger.ts
Normal file
@@ -0,0 +1,196 @@
|
|||||||
|
export type MemoryKind = "goal" | "constraint" | "decision" | "activeTask" | "openQuestion" | "relevantFile";
|
||||||
|
export type MemoryScope = "branch" | "session";
|
||||||
|
export type MemorySourceType = "user" | "assistant" | "toolResult" | "compaction" | "branchSummary";
|
||||||
|
|
||||||
|
export interface MemoryCandidate {
|
||||||
|
kind: MemoryKind;
|
||||||
|
subject: string;
|
||||||
|
text: string;
|
||||||
|
scope: MemoryScope;
|
||||||
|
sourceEntryId: string;
|
||||||
|
sourceType: MemorySourceType;
|
||||||
|
timestamp: number;
|
||||||
|
confidence: number;
|
||||||
|
}
|
||||||
|
|
||||||
|
export interface MemoryItem extends MemoryCandidate {
|
||||||
|
id: string;
|
||||||
|
freshness: number;
|
||||||
|
active: boolean;
|
||||||
|
supersedesId?: string;
|
||||||
|
}
|
||||||
|
|
||||||
|
export interface LedgerState {
|
||||||
|
items: MemoryItem[];
|
||||||
|
rollingSummary: string;
|
||||||
|
}
|
||||||
|
|
||||||
|
type MemorySlot = Pick<MemoryCandidate, "kind" | "scope" | "subject">;
|
||||||
|
|
||||||
|
export function createEmptyLedger(): LedgerState {
|
||||||
|
return { items: [], rollingSummary: "" };
|
||||||
|
}
|
||||||
|
|
||||||
|
function createId(candidate: MemoryCandidate): string {
|
||||||
|
return `${candidate.kind}:${candidate.scope}:${candidate.subject}:${candidate.timestamp}`;
|
||||||
|
}
|
||||||
|
|
||||||
|
function ensureUniqueId(items: Pick<MemoryItem, "id">[], baseId: string): string {
|
||||||
|
let id = baseId;
|
||||||
|
let suffix = 2;
|
||||||
|
|
||||||
|
while (items.some((item) => item.id === id)) {
|
||||||
|
id = `${baseId}:${suffix}`;
|
||||||
|
suffix += 1;
|
||||||
|
}
|
||||||
|
|
||||||
|
return id;
|
||||||
|
}
|
||||||
|
|
||||||
|
function sameSlot(left: MemorySlot, right: MemorySlot) {
|
||||||
|
return left.kind === right.kind && left.scope === right.scope && left.subject === right.subject;
|
||||||
|
}
|
||||||
|
|
||||||
|
function createSlotKey(slot: MemorySlot): string {
|
||||||
|
return `${slot.kind}\u0000${slot.scope}\u0000${slot.subject}`;
|
||||||
|
}
|
||||||
|
|
||||||
|
function compareStrings(left: string, right: string): number {
|
||||||
|
if (left === right) {
|
||||||
|
return 0;
|
||||||
|
}
|
||||||
|
|
||||||
|
return left < right ? -1 : 1;
|
||||||
|
}
|
||||||
|
|
||||||
|
function compareSameTimestampCandidates(
|
||||||
|
left: Pick<MemoryCandidate, "text" | "sourceType" | "sourceEntryId" | "confidence">,
|
||||||
|
right: Pick<MemoryCandidate, "text" | "sourceType" | "sourceEntryId" | "confidence">
|
||||||
|
): number {
|
||||||
|
// Exact-timestamp ties should resolve the same way no matter which candidate is processed first.
|
||||||
|
const textComparison = compareStrings(left.text, right.text);
|
||||||
|
if (textComparison !== 0) {
|
||||||
|
return textComparison;
|
||||||
|
}
|
||||||
|
|
||||||
|
const sourceTypeComparison = compareStrings(left.sourceType, right.sourceType);
|
||||||
|
if (sourceTypeComparison !== 0) {
|
||||||
|
return sourceTypeComparison;
|
||||||
|
}
|
||||||
|
|
||||||
|
const sourceEntryIdComparison = compareStrings(left.sourceEntryId, right.sourceEntryId);
|
||||||
|
if (sourceEntryIdComparison !== 0) {
|
||||||
|
return sourceEntryIdComparison;
|
||||||
|
}
|
||||||
|
|
||||||
|
if (left.confidence !== right.confidence) {
|
||||||
|
return left.confidence > right.confidence ? -1 : 1;
|
||||||
|
}
|
||||||
|
|
||||||
|
return 0;
|
||||||
|
}
|
||||||
|
|
||||||
|
function candidateSupersedesPrevious(candidate: MemoryCandidate, previous?: MemoryItem): boolean {
|
||||||
|
if (!previous) {
|
||||||
|
return true;
|
||||||
|
}
|
||||||
|
|
||||||
|
if (candidate.timestamp !== previous.timestamp) {
|
||||||
|
return candidate.timestamp > previous.timestamp;
|
||||||
|
}
|
||||||
|
|
||||||
|
return compareSameTimestampCandidates(candidate, previous) < 0;
|
||||||
|
}
|
||||||
|
|
||||||
|
function compareSlotItems(left: MemoryItem, right: MemoryItem): number {
|
||||||
|
if (left.timestamp !== right.timestamp) {
|
||||||
|
return right.timestamp - left.timestamp;
|
||||||
|
}
|
||||||
|
|
||||||
|
return compareSameTimestampCandidates(left, right);
|
||||||
|
}
|
||||||
|
|
||||||
|
function normalizeSlotItems(items: MemoryItem[], slot: MemorySlot): MemoryItem[] {
|
||||||
|
const slotIndices: number[] = [];
|
||||||
|
const slotItems: MemoryItem[] = [];
|
||||||
|
|
||||||
|
items.forEach((item, index) => {
|
||||||
|
if (!sameSlot(item, slot)) {
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
|
||||||
|
slotIndices.push(index);
|
||||||
|
slotItems.push(item);
|
||||||
|
});
|
||||||
|
|
||||||
|
if (slotItems.length <= 1) {
|
||||||
|
return items;
|
||||||
|
}
|
||||||
|
|
||||||
|
const sortedSlotItems = [...slotItems].sort(compareSlotItems);
|
||||||
|
const slotIds = new Map<string, number>();
|
||||||
|
const sortedSlotItemsWithIds = sortedSlotItems.map((item) => {
|
||||||
|
const baseId = createId(item);
|
||||||
|
const nextSlotIdCount = (slotIds.get(baseId) ?? 0) + 1;
|
||||||
|
slotIds.set(baseId, nextSlotIdCount);
|
||||||
|
|
||||||
|
return {
|
||||||
|
item,
|
||||||
|
id: nextSlotIdCount === 1 ? baseId : `${baseId}:${nextSlotIdCount}`,
|
||||||
|
};
|
||||||
|
});
|
||||||
|
|
||||||
|
const normalizedSlotItems = sortedSlotItemsWithIds.map(({ item, id }, index) => ({
|
||||||
|
...item,
|
||||||
|
id,
|
||||||
|
freshness: index === 0 ? item.timestamp : sortedSlotItemsWithIds[index - 1]!.item.timestamp,
|
||||||
|
active: index === 0,
|
||||||
|
supersedesId: sortedSlotItemsWithIds[index + 1]?.id,
|
||||||
|
}));
|
||||||
|
|
||||||
|
const normalizedItems = [...items];
|
||||||
|
slotIndices.forEach((slotIndex, index) => {
|
||||||
|
normalizedItems[slotIndex] = normalizedSlotItems[index]!;
|
||||||
|
});
|
||||||
|
|
||||||
|
return normalizedItems;
|
||||||
|
}
|
||||||
|
|
||||||
|
export function mergeCandidates(state: LedgerState, candidates: MemoryCandidate[]): LedgerState {
|
||||||
|
let items = [...state.items];
|
||||||
|
const affectedSlots = new Map<string, MemorySlot>();
|
||||||
|
|
||||||
|
for (const candidate of candidates) {
|
||||||
|
const previousIndex = items.findIndex((item) => item.active && sameSlot(item, candidate));
|
||||||
|
const previous = previousIndex === -1 ? undefined : items[previousIndex];
|
||||||
|
const candidateIsNewest = candidateSupersedesPrevious(candidate, previous);
|
||||||
|
|
||||||
|
if (previous && candidateIsNewest) {
|
||||||
|
items[previousIndex] = { ...previous, active: false, freshness: candidate.timestamp };
|
||||||
|
}
|
||||||
|
|
||||||
|
items.push({
|
||||||
|
...candidate,
|
||||||
|
id: ensureUniqueId(items, createId(candidate)),
|
||||||
|
freshness: candidate.timestamp,
|
||||||
|
active: candidateIsNewest,
|
||||||
|
supersedesId: candidateIsNewest ? previous?.id : undefined,
|
||||||
|
});
|
||||||
|
|
||||||
|
affectedSlots.set(createSlotKey(candidate), {
|
||||||
|
kind: candidate.kind,
|
||||||
|
scope: candidate.scope,
|
||||||
|
subject: candidate.subject,
|
||||||
|
});
|
||||||
|
}
|
||||||
|
|
||||||
|
for (const slot of affectedSlots.values()) {
|
||||||
|
items = normalizeSlotItems(items, slot);
|
||||||
|
}
|
||||||
|
|
||||||
|
return { ...state, items };
|
||||||
|
}
|
||||||
|
|
||||||
|
export function getActiveItems(state: LedgerState, kind?: MemoryKind): MemoryItem[] {
|
||||||
|
return state.items.filter((item) => item.active && (kind ? item.kind === kind : true));
|
||||||
|
}
|
||||||
130
.pi/agent/extensions/context-manager/src/packet.test.ts
Normal file
130
.pi/agent/extensions/context-manager/src/packet.test.ts
Normal file
@@ -0,0 +1,130 @@
|
|||||||
|
import test from "node:test";
|
||||||
|
import assert from "node:assert/strict";
|
||||||
|
import { resolvePolicy } from "./config.ts";
|
||||||
|
import { buildContextPacket } from "./packet.ts";
|
||||||
|
import { createEmptyLedger, mergeCandidates, type MemoryCandidate } from "./ledger.ts";
|
||||||
|
|
||||||
|
const baseCandidate: Omit<MemoryCandidate, "kind" | "subject" | "text"> = {
|
||||||
|
scope: "session",
|
||||||
|
sourceEntryId: "seed",
|
||||||
|
sourceType: "user",
|
||||||
|
timestamp: 1,
|
||||||
|
confidence: 1,
|
||||||
|
};
|
||||||
|
|
||||||
|
function estimateTokens(text: string) {
|
||||||
|
return Math.ceil(text.length / 4);
|
||||||
|
}
|
||||||
|
|
||||||
|
function memory(candidate: Pick<MemoryCandidate, "kind" | "subject" | "text"> & Partial<Omit<MemoryCandidate, "kind" | "subject" | "text">>): MemoryCandidate {
|
||||||
|
return {
|
||||||
|
...baseCandidate,
|
||||||
|
...candidate,
|
||||||
|
sourceEntryId: candidate.sourceEntryId ?? candidate.subject,
|
||||||
|
};
|
||||||
|
}
|
||||||
|
|
||||||
|
function buildPolicy(packetTokenCap: number) {
|
||||||
|
return {
|
||||||
|
...resolvePolicy({ mode: "balanced", contextWindow: 200_000 }),
|
||||||
|
packetTokenCap,
|
||||||
|
};
|
||||||
|
}
|
||||||
|
|
||||||
|
test("buildContextPacket keeps top-ranked facts from a section when the cap is tight", () => {
|
||||||
|
const expected = [
|
||||||
|
"## Active goal",
|
||||||
|
"- Keep packets compact.",
|
||||||
|
"",
|
||||||
|
"## Constraints",
|
||||||
|
"- Preserve the highest-priority constraint.",
|
||||||
|
"",
|
||||||
|
"## Key decisions",
|
||||||
|
"- Render selected sections in stable order.",
|
||||||
|
].join("\n");
|
||||||
|
const policy = buildPolicy(estimateTokens(expected));
|
||||||
|
const ledger = mergeCandidates(createEmptyLedger(), [
|
||||||
|
memory({ kind: "goal", subject: "goal", text: "Keep packets compact." }),
|
||||||
|
memory({ kind: "constraint", subject: "constraint-a", text: "Preserve the highest-priority constraint.", confidence: 1, timestamp: 3 }),
|
||||||
|
memory({
|
||||||
|
kind: "constraint",
|
||||||
|
subject: "constraint-b",
|
||||||
|
text: "Avoid dropping every constraint just because one extra bullet is too large for a tight packet cap.",
|
||||||
|
confidence: 0.6,
|
||||||
|
timestamp: 2,
|
||||||
|
}),
|
||||||
|
memory({
|
||||||
|
kind: "decision",
|
||||||
|
subject: "decision-a",
|
||||||
|
text: "Render selected sections in stable order.",
|
||||||
|
confidence: 0.9,
|
||||||
|
timestamp: 4,
|
||||||
|
sourceType: "assistant",
|
||||||
|
}),
|
||||||
|
]);
|
||||||
|
|
||||||
|
const packet = buildContextPacket(ledger, policy);
|
||||||
|
|
||||||
|
assert.equal(packet.text, expected);
|
||||||
|
assert.equal(packet.estimatedTokens, policy.packetTokenCap);
|
||||||
|
});
|
||||||
|
|
||||||
|
test("buildContextPacket uses cross-kind weights when only one lower-priority section can fit", () => {
|
||||||
|
const expected = [
|
||||||
|
"## Active goal",
|
||||||
|
"- Keep the agent moving.",
|
||||||
|
"",
|
||||||
|
"## Current task",
|
||||||
|
"- Fix packet trimming.",
|
||||||
|
].join("\n");
|
||||||
|
const policy = buildPolicy(estimateTokens(expected));
|
||||||
|
const ledger = mergeCandidates(createEmptyLedger(), [
|
||||||
|
memory({ kind: "goal", subject: "goal", text: "Keep the agent moving." }),
|
||||||
|
memory({
|
||||||
|
kind: "decision",
|
||||||
|
subject: "decision-a",
|
||||||
|
text: "Keep logs concise.",
|
||||||
|
confidence: 1,
|
||||||
|
timestamp: 2,
|
||||||
|
sourceType: "assistant",
|
||||||
|
}),
|
||||||
|
memory({
|
||||||
|
kind: "activeTask",
|
||||||
|
subject: "task-a",
|
||||||
|
text: "Fix packet trimming.",
|
||||||
|
confidence: 1,
|
||||||
|
timestamp: 2,
|
||||||
|
sourceType: "assistant",
|
||||||
|
}),
|
||||||
|
]);
|
||||||
|
|
||||||
|
const packet = buildContextPacket(ledger, policy);
|
||||||
|
|
||||||
|
assert.equal(packet.text, expected);
|
||||||
|
assert.equal(packet.estimatedTokens, policy.packetTokenCap);
|
||||||
|
});
|
||||||
|
|
||||||
|
test("buildContextPacket keeps a goal ahead of newer low-priority facts at realistic timestamp scales", () => {
|
||||||
|
const expected = [
|
||||||
|
"## Active goal",
|
||||||
|
"- Keep the agent on track.",
|
||||||
|
].join("\n");
|
||||||
|
const policy = buildPolicy(estimateTokens(expected));
|
||||||
|
const ledger = mergeCandidates(createEmptyLedger(), [
|
||||||
|
memory({ kind: "goal", subject: "goal", text: "Keep the agent on track.", timestamp: 1_000_000 }),
|
||||||
|
memory({
|
||||||
|
kind: "relevantFile",
|
||||||
|
subject: "runtime-file",
|
||||||
|
text: "src/runtime.ts",
|
||||||
|
timestamp: 10_000_000,
|
||||||
|
confidence: 1,
|
||||||
|
sourceType: "assistant",
|
||||||
|
scope: "branch",
|
||||||
|
}),
|
||||||
|
]);
|
||||||
|
|
||||||
|
const packet = buildContextPacket(ledger, policy);
|
||||||
|
|
||||||
|
assert.equal(packet.text, expected);
|
||||||
|
assert.equal(packet.estimatedTokens, policy.packetTokenCap);
|
||||||
|
});
|
||||||
91
.pi/agent/extensions/context-manager/src/packet.ts
Normal file
91
.pi/agent/extensions/context-manager/src/packet.ts
Normal file
@@ -0,0 +1,91 @@
|
|||||||
|
import type { Policy } from "./config.ts";
|
||||||
|
import { getActiveItems, type LedgerState, type MemoryItem, type MemoryKind } from "./ledger.ts";
|
||||||
|
|
||||||
|
const SECTION_ORDER: Array<{ kind: MemoryKind; title: string }> = [
|
||||||
|
{ kind: "goal", title: "Active goal" },
|
||||||
|
{ kind: "constraint", title: "Constraints" },
|
||||||
|
{ kind: "decision", title: "Key decisions" },
|
||||||
|
{ kind: "activeTask", title: "Current task" },
|
||||||
|
{ kind: "relevantFile", title: "Relevant files" },
|
||||||
|
{ kind: "openQuestion", title: "Open questions / blockers" },
|
||||||
|
];
|
||||||
|
|
||||||
|
const WEIGHTS: Record<MemoryKind, number> = {
|
||||||
|
goal: 100,
|
||||||
|
constraint: 90,
|
||||||
|
decision: 80,
|
||||||
|
activeTask: 85,
|
||||||
|
relevantFile: 60,
|
||||||
|
openQuestion: 70,
|
||||||
|
};
|
||||||
|
|
||||||
|
const SECTION_INDEX = new Map(SECTION_ORDER.map((section, index) => [section.kind, index]));
|
||||||
|
|
||||||
|
function estimateTokens(text: string): number {
|
||||||
|
return Math.ceil(text.length / 4);
|
||||||
|
}
|
||||||
|
|
||||||
|
function compareByPriority(left: MemoryItem, right: MemoryItem): number {
|
||||||
|
const weightDifference = WEIGHTS[right.kind] - WEIGHTS[left.kind];
|
||||||
|
if (weightDifference !== 0) {
|
||||||
|
return weightDifference;
|
||||||
|
}
|
||||||
|
|
||||||
|
if (left.confidence !== right.confidence) {
|
||||||
|
return right.confidence - left.confidence;
|
||||||
|
}
|
||||||
|
|
||||||
|
const sectionDifference = SECTION_INDEX.get(left.kind)! - SECTION_INDEX.get(right.kind)!;
|
||||||
|
if (sectionDifference !== 0) {
|
||||||
|
return sectionDifference;
|
||||||
|
}
|
||||||
|
|
||||||
|
if (left.freshness !== right.freshness) {
|
||||||
|
return right.freshness - left.freshness;
|
||||||
|
}
|
||||||
|
|
||||||
|
return left.id.localeCompare(right.id);
|
||||||
|
}
|
||||||
|
|
||||||
|
function sortByPriority(items: MemoryItem[]) {
|
||||||
|
return [...items].sort(compareByPriority);
|
||||||
|
}
|
||||||
|
|
||||||
|
function renderPacket(itemsByKind: Map<MemoryKind, MemoryItem[]>, selectedIds: Set<string>) {
|
||||||
|
const lines: string[] = [];
|
||||||
|
|
||||||
|
for (const section of SECTION_ORDER) {
|
||||||
|
const items = itemsByKind.get(section.kind)?.filter((item) => selectedIds.has(item.id)) ?? [];
|
||||||
|
if (items.length === 0) continue;
|
||||||
|
|
||||||
|
lines.push(`## ${section.title}`, ...items.map((item) => `- ${item.text}`), "");
|
||||||
|
}
|
||||||
|
|
||||||
|
return lines.join("\n").trim();
|
||||||
|
}
|
||||||
|
|
||||||
|
export function buildContextPacket(ledger: LedgerState, policy: Policy): { text: string; estimatedTokens: number } {
|
||||||
|
const itemsByKind = new Map<MemoryKind, MemoryItem[]>();
|
||||||
|
for (const section of SECTION_ORDER) {
|
||||||
|
itemsByKind.set(section.kind, sortByPriority(getActiveItems(ledger, section.kind)));
|
||||||
|
}
|
||||||
|
|
||||||
|
const candidates = sortByPriority(getActiveItems(ledger));
|
||||||
|
const selectedIds = new Set<string>();
|
||||||
|
let text = "";
|
||||||
|
|
||||||
|
for (const item of candidates) {
|
||||||
|
const tentativeSelectedIds = new Set(selectedIds);
|
||||||
|
tentativeSelectedIds.add(item.id);
|
||||||
|
|
||||||
|
const tentative = renderPacket(itemsByKind, tentativeSelectedIds);
|
||||||
|
if (estimateTokens(tentative) > policy.packetTokenCap) {
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
|
||||||
|
selectedIds.add(item.id);
|
||||||
|
text = tentative;
|
||||||
|
}
|
||||||
|
|
||||||
|
return { text, estimatedTokens: estimateTokens(text) };
|
||||||
|
}
|
||||||
67
.pi/agent/extensions/context-manager/src/persist.test.ts
Normal file
67
.pi/agent/extensions/context-manager/src/persist.test.ts
Normal file
@@ -0,0 +1,67 @@
|
|||||||
|
import test from "node:test";
|
||||||
|
import assert from "node:assert/strict";
|
||||||
|
import { deserializeLatestSnapshot, SNAPSHOT_ENTRY_TYPE, serializeSnapshot } from "./persist.ts";
|
||||||
|
|
||||||
|
function createSnapshot(lastZone: "green" | "yellow" | "red" | "compact", lastCompactionSummary: string) {
|
||||||
|
return serializeSnapshot({
|
||||||
|
mode: "balanced",
|
||||||
|
lastZone,
|
||||||
|
lastCompactionSummary,
|
||||||
|
lastBranchSummary: undefined,
|
||||||
|
ledger: {
|
||||||
|
items: [
|
||||||
|
{
|
||||||
|
id: `goal:session:root:${lastCompactionSummary}`,
|
||||||
|
kind: "goal",
|
||||||
|
subject: "root",
|
||||||
|
text: `Goal ${lastCompactionSummary}`,
|
||||||
|
scope: "session",
|
||||||
|
sourceEntryId: "u1",
|
||||||
|
sourceType: "user",
|
||||||
|
timestamp: 1,
|
||||||
|
confidence: 1,
|
||||||
|
freshness: 1,
|
||||||
|
active: true,
|
||||||
|
supersedesId: undefined,
|
||||||
|
},
|
||||||
|
],
|
||||||
|
rollingSummary: "summary",
|
||||||
|
},
|
||||||
|
});
|
||||||
|
}
|
||||||
|
|
||||||
|
test("deserializeLatestSnapshot restores the newest matching custom entry", () => {
|
||||||
|
const first = createSnapshot("yellow", "old");
|
||||||
|
const second = createSnapshot("red", "new");
|
||||||
|
|
||||||
|
const restored = deserializeLatestSnapshot([
|
||||||
|
{ type: "custom", customType: SNAPSHOT_ENTRY_TYPE, data: first },
|
||||||
|
{ type: "custom", customType: SNAPSHOT_ENTRY_TYPE, data: second },
|
||||||
|
]);
|
||||||
|
|
||||||
|
assert.equal(restored?.lastZone, "red");
|
||||||
|
assert.equal(restored?.lastCompactionSummary, "new");
|
||||||
|
});
|
||||||
|
|
||||||
|
test("deserializeLatestSnapshot skips malformed newer entries and clones the accepted snapshot", () => {
|
||||||
|
const valid = createSnapshot("yellow", "valid");
|
||||||
|
|
||||||
|
const restored = deserializeLatestSnapshot([
|
||||||
|
{ type: "custom", customType: SNAPSHOT_ENTRY_TYPE, data: valid },
|
||||||
|
{
|
||||||
|
type: "custom",
|
||||||
|
customType: SNAPSHOT_ENTRY_TYPE,
|
||||||
|
data: {
|
||||||
|
mode: "balanced",
|
||||||
|
lastZone: "red",
|
||||||
|
ledger: { items: "not-an-array", rollingSummary: "broken" },
|
||||||
|
},
|
||||||
|
},
|
||||||
|
]);
|
||||||
|
|
||||||
|
assert.deepEqual(restored, valid);
|
||||||
|
assert.notStrictEqual(restored, valid);
|
||||||
|
|
||||||
|
restored!.ledger.items[0]!.text = "mutated";
|
||||||
|
assert.equal(valid.ledger.items[0]!.text, "Goal valid");
|
||||||
|
});
|
||||||
142
.pi/agent/extensions/context-manager/src/persist.ts
Normal file
142
.pi/agent/extensions/context-manager/src/persist.ts
Normal file
@@ -0,0 +1,142 @@
|
|||||||
|
import type { ContextMode, ContextZone } from "./config.ts";
|
||||||
|
import type { LedgerState, MemoryItem, MemoryKind, MemoryScope, MemorySourceType } from "./ledger.ts";
|
||||||
|
|
||||||
|
export const SNAPSHOT_ENTRY_TYPE = "context-manager.snapshot";
|
||||||
|
|
||||||
|
export interface RuntimeSnapshot {
|
||||||
|
mode: ContextMode;
|
||||||
|
lastZone: ContextZone;
|
||||||
|
lastObservedTokens?: number;
|
||||||
|
lastCompactionSummary?: string;
|
||||||
|
lastBranchSummary?: string;
|
||||||
|
ledger: LedgerState;
|
||||||
|
}
|
||||||
|
|
||||||
|
const CONTEXT_MODES = new Set<ContextMode>(["conservative", "balanced", "aggressive"]);
|
||||||
|
const CONTEXT_ZONES = new Set<ContextZone>(["green", "yellow", "red", "compact"]);
|
||||||
|
const MEMORY_KINDS = new Set<MemoryKind>(["goal", "constraint", "decision", "activeTask", "openQuestion", "relevantFile"]);
|
||||||
|
const MEMORY_SCOPES = new Set<MemoryScope>(["branch", "session"]);
|
||||||
|
const MEMORY_SOURCE_TYPES = new Set<MemorySourceType>(["user", "assistant", "toolResult", "compaction", "branchSummary"]);
|
||||||
|
|
||||||
|
export function serializeSnapshot(snapshot: RuntimeSnapshot): RuntimeSnapshot {
|
||||||
|
return structuredClone(snapshot);
|
||||||
|
}
|
||||||
|
|
||||||
|
function isRecord(value: unknown): value is Record<string, unknown> {
|
||||||
|
return typeof value === "object" && value !== null;
|
||||||
|
}
|
||||||
|
|
||||||
|
function isFiniteNumber(value: unknown): value is number {
|
||||||
|
return typeof value === "number" && Number.isFinite(value);
|
||||||
|
}
|
||||||
|
|
||||||
|
function isOptionalString(value: unknown): value is string | undefined {
|
||||||
|
return value === undefined || typeof value === "string";
|
||||||
|
}
|
||||||
|
|
||||||
|
function parseMemoryItem(value: unknown): MemoryItem | undefined {
|
||||||
|
if (!isRecord(value)) {
|
||||||
|
return undefined;
|
||||||
|
}
|
||||||
|
|
||||||
|
if (
|
||||||
|
typeof value.id !== "string" ||
|
||||||
|
!MEMORY_KINDS.has(value.kind as MemoryKind) ||
|
||||||
|
typeof value.subject !== "string" ||
|
||||||
|
typeof value.text !== "string" ||
|
||||||
|
!MEMORY_SCOPES.has(value.scope as MemoryScope) ||
|
||||||
|
typeof value.sourceEntryId !== "string" ||
|
||||||
|
!MEMORY_SOURCE_TYPES.has(value.sourceType as MemorySourceType) ||
|
||||||
|
!isFiniteNumber(value.timestamp) ||
|
||||||
|
!isFiniteNumber(value.confidence) ||
|
||||||
|
!isFiniteNumber(value.freshness) ||
|
||||||
|
typeof value.active !== "boolean" ||
|
||||||
|
!isOptionalString(value.supersedesId)
|
||||||
|
) {
|
||||||
|
return undefined;
|
||||||
|
}
|
||||||
|
|
||||||
|
return {
|
||||||
|
id: value.id,
|
||||||
|
kind: value.kind as MemoryKind,
|
||||||
|
subject: value.subject,
|
||||||
|
text: value.text,
|
||||||
|
scope: value.scope as MemoryScope,
|
||||||
|
sourceEntryId: value.sourceEntryId,
|
||||||
|
sourceType: value.sourceType as MemorySourceType,
|
||||||
|
timestamp: value.timestamp,
|
||||||
|
confidence: value.confidence,
|
||||||
|
freshness: value.freshness,
|
||||||
|
active: value.active,
|
||||||
|
supersedesId: value.supersedesId,
|
||||||
|
};
|
||||||
|
}
|
||||||
|
|
||||||
|
function parseLedgerState(value: unknown): LedgerState | undefined {
|
||||||
|
if (!isRecord(value) || !Array.isArray(value.items) || typeof value.rollingSummary !== "string") {
|
||||||
|
return undefined;
|
||||||
|
}
|
||||||
|
|
||||||
|
const items: MemoryItem[] = [];
|
||||||
|
for (const item of value.items) {
|
||||||
|
const parsed = parseMemoryItem(item);
|
||||||
|
if (!parsed) {
|
||||||
|
return undefined;
|
||||||
|
}
|
||||||
|
|
||||||
|
items.push(parsed);
|
||||||
|
}
|
||||||
|
|
||||||
|
return {
|
||||||
|
items,
|
||||||
|
rollingSummary: value.rollingSummary,
|
||||||
|
};
|
||||||
|
}
|
||||||
|
|
||||||
|
function parseRuntimeSnapshot(value: unknown): RuntimeSnapshot | undefined {
|
||||||
|
if (
|
||||||
|
!isRecord(value) ||
|
||||||
|
!CONTEXT_MODES.has(value.mode as ContextMode) ||
|
||||||
|
!CONTEXT_ZONES.has(value.lastZone as ContextZone) ||
|
||||||
|
!isOptionalString(value.lastCompactionSummary) ||
|
||||||
|
!isOptionalString(value.lastBranchSummary) ||
|
||||||
|
(value.lastObservedTokens !== undefined && !isFiniteNumber(value.lastObservedTokens))
|
||||||
|
) {
|
||||||
|
return undefined;
|
||||||
|
}
|
||||||
|
|
||||||
|
const ledger = parseLedgerState(value.ledger);
|
||||||
|
if (!ledger) {
|
||||||
|
return undefined;
|
||||||
|
}
|
||||||
|
|
||||||
|
const snapshot: RuntimeSnapshot = {
|
||||||
|
mode: value.mode as ContextMode,
|
||||||
|
lastZone: value.lastZone as ContextZone,
|
||||||
|
lastCompactionSummary: value.lastCompactionSummary,
|
||||||
|
lastBranchSummary: value.lastBranchSummary,
|
||||||
|
ledger,
|
||||||
|
};
|
||||||
|
|
||||||
|
if (value.lastObservedTokens !== undefined) {
|
||||||
|
snapshot.lastObservedTokens = value.lastObservedTokens;
|
||||||
|
}
|
||||||
|
|
||||||
|
return snapshot;
|
||||||
|
}
|
||||||
|
|
||||||
|
export function deserializeLatestSnapshot(entries: Array<{ type: string; customType?: string; data?: unknown }>): RuntimeSnapshot | undefined {
|
||||||
|
for (let index = entries.length - 1; index >= 0; index -= 1) {
|
||||||
|
const entry = entries[index]!;
|
||||||
|
if (entry.type !== "custom" || entry.customType !== SNAPSHOT_ENTRY_TYPE) {
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
|
||||||
|
const snapshot = parseRuntimeSnapshot(entry.data);
|
||||||
|
if (snapshot) {
|
||||||
|
return snapshot;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
return undefined;
|
||||||
|
}
|
||||||
131
.pi/agent/extensions/context-manager/src/prune.test.ts
Normal file
131
.pi/agent/extensions/context-manager/src/prune.test.ts
Normal file
@@ -0,0 +1,131 @@
|
|||||||
|
import test from "node:test";
|
||||||
|
import assert from "node:assert/strict";
|
||||||
|
import { resolvePolicy } from "./config.ts";
|
||||||
|
import { pruneContextMessages } from "./prune.ts";
|
||||||
|
|
||||||
|
const bulky = "line\n".repeat(300);
|
||||||
|
const boundaryBulky = "boundary\n".repeat(300);
|
||||||
|
const thresholdWithTrailingNewline = "threshold\n".repeat(150);
|
||||||
|
|
||||||
|
function buildPolicy(recentUserTurns = 4) {
|
||||||
|
return {
|
||||||
|
...resolvePolicy({ mode: "balanced", contextWindow: 200_000 }),
|
||||||
|
recentUserTurns,
|
||||||
|
};
|
||||||
|
}
|
||||||
|
|
||||||
|
test("pruneContextMessages replaces old bulky tool results with distilled summaries instead of deleting them", () => {
|
||||||
|
const policy = buildPolicy(2);
|
||||||
|
const bulkyFailure = [
|
||||||
|
"Build failed while compiling focus parser",
|
||||||
|
"Error: missing export createFocusMatcher from ./summary-focus.ts",
|
||||||
|
...Array.from({ length: 220 }, () => "stack frame"),
|
||||||
|
].join("\n");
|
||||||
|
const messages = [
|
||||||
|
{ role: "user", content: "turn 1" },
|
||||||
|
{ role: "toolResult", toolName: "bash", content: bulkyFailure },
|
||||||
|
{ role: "assistant", content: "observed turn 1" },
|
||||||
|
{ role: "user", content: "turn 2" },
|
||||||
|
{ role: "assistant", content: "observed turn 2" },
|
||||||
|
{ role: "user", content: "turn 3" },
|
||||||
|
];
|
||||||
|
|
||||||
|
const pruned = pruneContextMessages(messages, policy);
|
||||||
|
|
||||||
|
const distilled = pruned.find((message) => message.role === "toolResult");
|
||||||
|
assert.ok(distilled);
|
||||||
|
assert.match(distilled!.content, /missing export createFocusMatcher/);
|
||||||
|
assert.doesNotMatch(distilled!.content, /stack frame\nstack frame\nstack frame/);
|
||||||
|
});
|
||||||
|
|
||||||
|
test("aggressive mode distills an older bulky tool result sooner than conservative mode", () => {
|
||||||
|
const conservative = resolvePolicy({ mode: "conservative", contextWindow: 200_000 });
|
||||||
|
const aggressive = resolvePolicy({ mode: "aggressive", contextWindow: 200_000 });
|
||||||
|
const messages = [
|
||||||
|
{ role: "user", content: "turn 1" },
|
||||||
|
{ role: "toolResult", toolName: "read", content: bulky },
|
||||||
|
{ role: "assistant", content: "after turn 1" },
|
||||||
|
{ role: "user", content: "turn 2" },
|
||||||
|
{ role: "assistant", content: "after turn 2" },
|
||||||
|
{ role: "user", content: "turn 3" },
|
||||||
|
{ role: "assistant", content: "after turn 3" },
|
||||||
|
{ role: "user", content: "turn 4" },
|
||||||
|
];
|
||||||
|
|
||||||
|
const conservativePruned = pruneContextMessages(messages, conservative);
|
||||||
|
const aggressivePruned = pruneContextMessages(messages, aggressive);
|
||||||
|
|
||||||
|
assert.equal(conservativePruned[1]?.content, bulky);
|
||||||
|
assert.notEqual(aggressivePruned[1]?.content, bulky);
|
||||||
|
assert.match(aggressivePruned[1]?.content ?? "", /^\[distilled read output\]/);
|
||||||
|
});
|
||||||
|
|
||||||
|
test("pruneContextMessages keeps recent bulky tool results inside the recent-turn window", () => {
|
||||||
|
const policy = buildPolicy(2);
|
||||||
|
const messages = [
|
||||||
|
{ role: "user", content: "turn 1" },
|
||||||
|
{ role: "assistant", content: "observed turn 1" },
|
||||||
|
{ role: "user", content: "turn 2" },
|
||||||
|
{ role: "toolResult", toolName: "read", content: bulky },
|
||||||
|
{ role: "assistant", content: "observed turn 2" },
|
||||||
|
{ role: "user", content: "turn 3" },
|
||||||
|
];
|
||||||
|
|
||||||
|
const pruned = pruneContextMessages(messages, policy);
|
||||||
|
|
||||||
|
assert.deepEqual(pruned, messages);
|
||||||
|
});
|
||||||
|
|
||||||
|
test("pruneContextMessages keeps old non-bulky tool results outside the recent-turn window", () => {
|
||||||
|
const policy = buildPolicy(2);
|
||||||
|
const messages = [
|
||||||
|
{ role: "user", content: "turn 1" },
|
||||||
|
{ role: "toolResult", toolName: "read", content: "short output" },
|
||||||
|
{ role: "assistant", content: "observed turn 1" },
|
||||||
|
{ role: "user", content: "turn 2" },
|
||||||
|
{ role: "assistant", content: "observed turn 2" },
|
||||||
|
{ role: "user", content: "turn 3" },
|
||||||
|
];
|
||||||
|
|
||||||
|
const pruned = pruneContextMessages(messages, policy);
|
||||||
|
|
||||||
|
assert.deepEqual(pruned, messages);
|
||||||
|
});
|
||||||
|
|
||||||
|
test("pruneContextMessages keeps exactly-150-line tool results with a trailing newline", () => {
|
||||||
|
const policy = buildPolicy(2);
|
||||||
|
const messages = [
|
||||||
|
{ role: "user", content: "turn 1" },
|
||||||
|
{ role: "toolResult", toolName: "read", content: thresholdWithTrailingNewline },
|
||||||
|
{ role: "assistant", content: "after threshold output" },
|
||||||
|
{ role: "user", content: "turn 2" },
|
||||||
|
{ role: "assistant", content: "after turn 2" },
|
||||||
|
{ role: "user", content: "turn 3" },
|
||||||
|
];
|
||||||
|
|
||||||
|
const pruned = pruneContextMessages(messages, policy);
|
||||||
|
|
||||||
|
assert.deepEqual(pruned, messages);
|
||||||
|
});
|
||||||
|
|
||||||
|
test("pruneContextMessages honors the recent-user-turn boundary", () => {
|
||||||
|
const policy = buildPolicy(2);
|
||||||
|
const messages = [
|
||||||
|
{ role: "user", content: "turn 1" },
|
||||||
|
{ role: "toolResult", toolName: "read", content: bulky },
|
||||||
|
{ role: "assistant", content: "after turn 1" },
|
||||||
|
{ role: "user", content: "turn 2" },
|
||||||
|
{ role: "toolResult", toolName: "read", content: boundaryBulky },
|
||||||
|
{ role: "assistant", content: "after turn 2" },
|
||||||
|
{ role: "user", content: "turn 3" },
|
||||||
|
];
|
||||||
|
|
||||||
|
const pruned = pruneContextMessages(messages, policy);
|
||||||
|
|
||||||
|
assert.equal(pruned[1]?.role, "toolResult");
|
||||||
|
assert.match(pruned[1]?.content ?? "", /^\[distilled read output\]/);
|
||||||
|
assert.deepEqual(
|
||||||
|
pruned.map((message) => message.content),
|
||||||
|
["turn 1", pruned[1]!.content, "after turn 1", "turn 2", boundaryBulky, "after turn 2", "turn 3"]
|
||||||
|
);
|
||||||
|
});
|
||||||
54
.pi/agent/extensions/context-manager/src/prune.ts
Normal file
54
.pi/agent/extensions/context-manager/src/prune.ts
Normal file
@@ -0,0 +1,54 @@
|
|||||||
|
import type { Policy } from "./config.ts";
|
||||||
|
import { distillToolResult } from "./distill.ts";
|
||||||
|
|
||||||
|
export interface ContextMessage {
|
||||||
|
role: string;
|
||||||
|
content: string;
|
||||||
|
toolName?: string;
|
||||||
|
original?: unknown;
|
||||||
|
distilled?: boolean;
|
||||||
|
}
|
||||||
|
|
||||||
|
function isBulky(content: string, policy: Policy) {
|
||||||
|
const bytes = Buffer.byteLength(content, "utf8");
|
||||||
|
const parts = content.split("\n");
|
||||||
|
const lines = content.endsWith("\n") ? parts.length - 1 : parts.length;
|
||||||
|
return bytes > policy.bulkyBytes || lines > policy.bulkyLines;
|
||||||
|
}
|
||||||
|
|
||||||
|
export function pruneContextMessages(messages: ContextMessage[], policy: Policy): ContextMessage[] {
|
||||||
|
let seenUserTurns = 0;
|
||||||
|
const keep = new Set<number>();
|
||||||
|
|
||||||
|
for (let index = messages.length - 1; index >= 0; index -= 1) {
|
||||||
|
const message = messages[index]!;
|
||||||
|
keep.add(index);
|
||||||
|
if (message.role === "user") {
|
||||||
|
seenUserTurns += 1;
|
||||||
|
if (seenUserTurns >= policy.recentUserTurns) {
|
||||||
|
break;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
const next: ContextMessage[] = [];
|
||||||
|
for (const [index, message] of messages.entries()) {
|
||||||
|
if (keep.has(index) || message.role !== "toolResult" || !isBulky(message.content, policy)) {
|
||||||
|
next.push(message);
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
|
||||||
|
const distilled = distillToolResult({ toolName: message.toolName, content: message.content });
|
||||||
|
if (!distilled) {
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
|
||||||
|
next.push({
|
||||||
|
...message,
|
||||||
|
content: distilled,
|
||||||
|
distilled: true,
|
||||||
|
});
|
||||||
|
}
|
||||||
|
|
||||||
|
return next;
|
||||||
|
}
|
||||||
178
.pi/agent/extensions/context-manager/src/runtime.test.ts
Normal file
178
.pi/agent/extensions/context-manager/src/runtime.test.ts
Normal file
@@ -0,0 +1,178 @@
|
|||||||
|
import test from "node:test";
|
||||||
|
import assert from "node:assert/strict";
|
||||||
|
import { SNAPSHOT_ENTRY_TYPE, deserializeLatestSnapshot } from "./persist.ts";
|
||||||
|
import { createContextManagerRuntime } from "./runtime.ts";
|
||||||
|
|
||||||
|
test("runtime ingests transcript slices and updates pressure state", () => {
|
||||||
|
const runtime = createContextManagerRuntime({ mode: "balanced", contextWindow: 200_000 });
|
||||||
|
|
||||||
|
runtime.ingest({ entryId: "u1", role: "user", text: "Goal: Build a pi context manager. Next: wire hooks.", timestamp: 1 });
|
||||||
|
runtime.observeTokens(150_000);
|
||||||
|
|
||||||
|
const packet = runtime.buildPacket();
|
||||||
|
assert.match(packet.text, /Build a pi context manager/);
|
||||||
|
assert.equal(runtime.getSnapshot().lastZone, "red");
|
||||||
|
});
|
||||||
|
|
||||||
|
test("runtime keeps the session root goal while allowing later branch-local goals", () => {
|
||||||
|
const runtime = createContextManagerRuntime({ mode: "balanced", contextWindow: 200_000 });
|
||||||
|
|
||||||
|
runtime.ingest({ entryId: "u-root-goal", role: "user", text: "Goal: Ship the context manager extension.", timestamp: 10 });
|
||||||
|
runtime.ingest({ entryId: "u-branch-goal", role: "user", text: "Goal: prototype a branch-local tree handoff.", timestamp: 11 });
|
||||||
|
|
||||||
|
const packet = runtime.buildPacket();
|
||||||
|
assert.match(packet.text, /Ship the context manager extension/);
|
||||||
|
assert.match(packet.text, /prototype a branch-local tree handoff/);
|
||||||
|
assert.equal(
|
||||||
|
runtime
|
||||||
|
.getSnapshot()
|
||||||
|
.ledger.items.filter((item) => item.active && item.kind === "goal" && item.scope === "session").length,
|
||||||
|
1,
|
||||||
|
);
|
||||||
|
});
|
||||||
|
|
||||||
|
test("recordCompactionSummary and recordBranchSummary update snapshot state and resume output", () => {
|
||||||
|
const runtime = createContextManagerRuntime({ mode: "balanced", contextWindow: 200_000 });
|
||||||
|
|
||||||
|
runtime.ingest({ entryId: "u-artifact-1", role: "user", text: "Goal: Ship the context manager extension.", timestamp: 20 });
|
||||||
|
runtime.recordCompactionSummary(
|
||||||
|
"## Key Decisions\n- Keep summaries deterministic.\n\n## Open questions and blockers\n- Verify /tree replaceInstructions behavior.",
|
||||||
|
);
|
||||||
|
runtime.recordBranchSummary("# Handoff for branch\n\n## Key Decisions\n- Do not leak branch-local goals.");
|
||||||
|
|
||||||
|
const snapshot = runtime.getSnapshot();
|
||||||
|
assert.match(snapshot.lastCompactionSummary ?? "", /Keep summaries deterministic/);
|
||||||
|
assert.match(snapshot.lastBranchSummary ?? "", /Do not leak branch-local goals/);
|
||||||
|
assert.match(runtime.buildResumePacket(), /Verify \/tree replaceInstructions behavior/);
|
||||||
|
assert.match(runtime.buildResumePacket(), /Do not leak branch-local goals/);
|
||||||
|
});
|
||||||
|
|
||||||
|
test("buildPacket tightens the live packet after pressure reaches the compact zone", () => {
|
||||||
|
const runtime = createContextManagerRuntime({ mode: "balanced", contextWindow: 200_000 });
|
||||||
|
|
||||||
|
runtime.restore({
|
||||||
|
mode: "balanced",
|
||||||
|
lastZone: "green",
|
||||||
|
ledger: {
|
||||||
|
rollingSummary: "",
|
||||||
|
items: [
|
||||||
|
{ id: "goal:session:root-goal:1", kind: "goal", subject: "root-goal", text: "Ship the context manager extension with deterministic handoffs and predictable branch-boundary behavior.", scope: "session", sourceEntryId: "u1", sourceType: "user", timestamp: 1, confidence: 1, freshness: 1, active: true },
|
||||||
|
{ id: "constraint:session:must-1:1", kind: "constraint", subject: "must-1", text: "Keep the public API stable while hardening branch-boundary state carryover, fallback summary replay, and resume injection behavior.", scope: "session", sourceEntryId: "u1", sourceType: "user", timestamp: 2, confidence: 0.9, freshness: 2, active: true },
|
||||||
|
{ id: "decision:branch:decision-1:1", kind: "decision", subject: "decision-1", text: "Persist summary artifacts, replay them after the latest snapshot, and surface them through the next hidden resume packet before normal packet injection resumes.", scope: "branch", sourceEntryId: "a1", sourceType: "assistant", timestamp: 3, confidence: 0.9, freshness: 3, active: true },
|
||||||
|
{ id: "activeTask:branch:task-1:1", kind: "activeTask", subject: "task-1", text: "Verify mode-dependent pruning, packet tightening under pressure, and snapshot-less branch rehydration without stale handoff leakage.", scope: "branch", sourceEntryId: "a2", sourceType: "assistant", timestamp: 4, confidence: 0.8, freshness: 4, active: true },
|
||||||
|
{ id: "openQuestion:branch:question-1:1", kind: "openQuestion", subject: "question-1", text: "Confirm whether default pi fallback summaries preserve blockers and active work end to end when custom compaction falls back.", scope: "branch", sourceEntryId: "a3", sourceType: "assistant", timestamp: 5, confidence: 0.8, freshness: 5, active: true },
|
||||||
|
{ id: "relevantFile:branch:file-1:1", kind: "relevantFile", subject: "file-1", text: "src/context-manager/very/long/path/for/runtime/index-and-hook-wiring-reference-file-1.ts", scope: "branch", sourceEntryId: "t1", sourceType: "toolResult", timestamp: 6, confidence: 0.7, freshness: 6, active: true },
|
||||||
|
{ id: "relevantFile:branch:file-2:1", kind: "relevantFile", subject: "file-2", text: "src/context-manager/very/long/path/for/runtime/index-and-hook-wiring-reference-file-2.ts", scope: "branch", sourceEntryId: "t2", sourceType: "toolResult", timestamp: 7, confidence: 0.7, freshness: 7, active: true },
|
||||||
|
{ id: "relevantFile:branch:file-3:1", kind: "relevantFile", subject: "file-3", text: "src/context-manager/very/long/path/for/runtime/index-and-hook-wiring-reference-file-3.ts", scope: "branch", sourceEntryId: "t3", sourceType: "toolResult", timestamp: 8, confidence: 0.7, freshness: 8, active: true },
|
||||||
|
{ id: "relevantFile:branch:file-4:1", kind: "relevantFile", subject: "file-4", text: "src/context-manager/very/long/path/for/runtime/index-and-hook-wiring-reference-file-4.ts", scope: "branch", sourceEntryId: "t4", sourceType: "toolResult", timestamp: 9, confidence: 0.7, freshness: 9, active: true },
|
||||||
|
{ id: "relevantFile:branch:file-5:1", kind: "relevantFile", subject: "file-5", text: "src/context-manager/very/long/path/for/runtime/index-and-hook-wiring-reference-file-5.ts", scope: "branch", sourceEntryId: "t5", sourceType: "toolResult", timestamp: 10, confidence: 0.7, freshness: 10, active: true },
|
||||||
|
{ id: "relevantFile:branch:file-6:1", kind: "relevantFile", subject: "file-6", text: "src/context-manager/very/long/path/for/runtime/index-and-hook-wiring-reference-file-6.ts", scope: "branch", sourceEntryId: "t6", sourceType: "toolResult", timestamp: 11, confidence: 0.7, freshness: 11, active: true },
|
||||||
|
{ id: "relevantFile:branch:file-7:1", kind: "relevantFile", subject: "file-7", text: "src/context-manager/very/long/path/for/runtime/index-and-hook-wiring-reference-file-7.ts", scope: "branch", sourceEntryId: "t7", sourceType: "toolResult", timestamp: 12, confidence: 0.7, freshness: 12, active: true },
|
||||||
|
{ id: "relevantFile:branch:file-8:1", kind: "relevantFile", subject: "file-8", text: "src/context-manager/very/long/path/for/runtime/index-and-hook-wiring-reference-file-8.ts", scope: "branch", sourceEntryId: "t8", sourceType: "toolResult", timestamp: 13, confidence: 0.7, freshness: 13, active: true },
|
||||||
|
{ id: "relevantFile:branch:file-9:1", kind: "relevantFile", subject: "file-9", text: "src/context-manager/very/long/path/for/runtime/index-and-hook-wiring-reference-file-9.ts", scope: "branch", sourceEntryId: "t9", sourceType: "toolResult", timestamp: 14, confidence: 0.7, freshness: 14, active: true },
|
||||||
|
{ id: "relevantFile:branch:file-10:1", kind: "relevantFile", subject: "file-10", text: "src/context-manager/very/long/path/for/runtime/index-and-hook-wiring-reference-file-10.ts", scope: "branch", sourceEntryId: "t10", sourceType: "toolResult", timestamp: 15, confidence: 0.7, freshness: 15, active: true },
|
||||||
|
{ id: "relevantFile:branch:file-11:1", kind: "relevantFile", subject: "file-11", text: "src/context-manager/very/long/path/for/runtime/index-and-hook-wiring-reference-file-11.ts", scope: "branch", sourceEntryId: "t11", sourceType: "toolResult", timestamp: 16, confidence: 0.7, freshness: 16, active: true },
|
||||||
|
{ id: "relevantFile:branch:file-12:1", kind: "relevantFile", subject: "file-12", text: "src/context-manager/very/long/path/for/runtime/index-and-hook-wiring-reference-file-12.ts", scope: "branch", sourceEntryId: "t12", sourceType: "toolResult", timestamp: 17, confidence: 0.7, freshness: 17, active: true },
|
||||||
|
{ id: "relevantFile:branch:file-13:1", kind: "relevantFile", subject: "file-13", text: "src/context-manager/very/long/path/for/runtime/index-and-hook-wiring-reference-file-13.ts", scope: "branch", sourceEntryId: "t13", sourceType: "toolResult", timestamp: 18, confidence: 0.7, freshness: 18, active: true },
|
||||||
|
{ id: "relevantFile:branch:file-14:1", kind: "relevantFile", subject: "file-14", text: "src/context-manager/very/long/path/for/runtime/index-and-hook-wiring-reference-file-14.ts", scope: "branch", sourceEntryId: "t14", sourceType: "toolResult", timestamp: 19, confidence: 0.7, freshness: 19, active: true },
|
||||||
|
{ id: "relevantFile:branch:file-15:1", kind: "relevantFile", subject: "file-15", text: "src/context-manager/very/long/path/for/runtime/index-and-hook-wiring-reference-file-15.ts", scope: "branch", sourceEntryId: "t15", sourceType: "toolResult", timestamp: 20, confidence: 0.7, freshness: 20, active: true },
|
||||||
|
{ id: "relevantFile:branch:file-16:1", kind: "relevantFile", subject: "file-16", text: "src/context-manager/very/long/path/for/runtime/index-and-hook-wiring-reference-file-16.ts", scope: "branch", sourceEntryId: "t16", sourceType: "toolResult", timestamp: 21, confidence: 0.7, freshness: 21, active: true },
|
||||||
|
{ id: "relevantFile:branch:file-17:1", kind: "relevantFile", subject: "file-17", text: "src/context-manager/very/long/path/for/runtime/index-and-hook-wiring-reference-file-17.ts", scope: "branch", sourceEntryId: "t17", sourceType: "toolResult", timestamp: 22, confidence: 0.7, freshness: 22, active: true },
|
||||||
|
{ id: "relevantFile:branch:file-18:1", kind: "relevantFile", subject: "file-18", text: "src/context-manager/very/long/path/for/runtime/index-and-hook-wiring-reference-file-18.ts", scope: "branch", sourceEntryId: "t18", sourceType: "toolResult", timestamp: 23, confidence: 0.7, freshness: 23, active: true },
|
||||||
|
{ id: "relevantFile:branch:file-19:1", kind: "relevantFile", subject: "file-19", text: "src/context-manager/very/long/path/for/runtime/index-and-hook-wiring-reference-file-19.ts", scope: "branch", sourceEntryId: "t19", sourceType: "toolResult", timestamp: 24, confidence: 0.7, freshness: 24, active: true },
|
||||||
|
{ id: "relevantFile:branch:file-20:1", kind: "relevantFile", subject: "file-20", text: "src/context-manager/very/long/path/for/runtime/index-and-hook-wiring-reference-file-20.ts", scope: "branch", sourceEntryId: "t20", sourceType: "toolResult", timestamp: 25, confidence: 0.7, freshness: 25, active: true },
|
||||||
|
{ id: "relevantFile:branch:file-21:1", kind: "relevantFile", subject: "file-21", text: "src/context-manager/very/long/path/for/runtime/index-and-hook-wiring-reference-file-21.ts", scope: "branch", sourceEntryId: "t21", sourceType: "toolResult", timestamp: 26, confidence: 0.7, freshness: 26, active: true },
|
||||||
|
{ id: "relevantFile:branch:file-22:1", kind: "relevantFile", subject: "file-22", text: "src/context-manager/very/long/path/for/runtime/index-and-hook-wiring-reference-file-22.ts", scope: "branch", sourceEntryId: "t22", sourceType: "toolResult", timestamp: 27, confidence: 0.7, freshness: 27, active: true },
|
||||||
|
],
|
||||||
|
},
|
||||||
|
});
|
||||||
|
|
||||||
|
const before = runtime.buildPacket();
|
||||||
|
runtime.observeTokens(170_000);
|
||||||
|
const after = runtime.buildPacket();
|
||||||
|
|
||||||
|
assert.ok(after.estimatedTokens < before.estimatedTokens);
|
||||||
|
});
|
||||||
|
|
||||||
|
test("runtime recomputes lastZone when setContextWindow and setMode change policy", () => {
|
||||||
|
const runtime = createContextManagerRuntime({ mode: "balanced", contextWindow: 200_000 });
|
||||||
|
|
||||||
|
runtime.observeTokens(150_000);
|
||||||
|
assert.equal(runtime.getSnapshot().lastZone, "red");
|
||||||
|
|
||||||
|
runtime.setContextWindow(300_000);
|
||||||
|
assert.equal(runtime.getSnapshot().lastZone, "green");
|
||||||
|
|
||||||
|
runtime.setMode("aggressive");
|
||||||
|
assert.equal(runtime.getSnapshot().lastZone, "yellow");
|
||||||
|
});
|
||||||
|
|
||||||
|
test("restore recomputes lastZone against the receiving runtime policy", () => {
|
||||||
|
const source = createContextManagerRuntime({ mode: "balanced", contextWindow: 200_000 });
|
||||||
|
source.observeTokens(150_000);
|
||||||
|
|
||||||
|
const target = createContextManagerRuntime({ mode: "balanced", contextWindow: 500_000 });
|
||||||
|
target.restore(source.getSnapshot());
|
||||||
|
|
||||||
|
assert.equal(target.getSnapshot().lastZone, "green");
|
||||||
|
});
|
||||||
|
|
||||||
|
test("restore resets legacy lastZone when the snapshot lacks lastObservedTokens", () => {
|
||||||
|
const runtime = createContextManagerRuntime({ mode: "balanced", contextWindow: 500_000 });
|
||||||
|
|
||||||
|
runtime.restore({
|
||||||
|
mode: "balanced",
|
||||||
|
lastZone: "red",
|
||||||
|
ledger: { items: [], rollingSummary: "" },
|
||||||
|
});
|
||||||
|
|
||||||
|
const restored = runtime.getSnapshot();
|
||||||
|
assert.equal(restored.lastZone, "green");
|
||||||
|
assert.equal(restored.lastObservedTokens, undefined);
|
||||||
|
});
|
||||||
|
|
||||||
|
test("legacy snapshot deserialization plus restore clears stale lastZone", () => {
|
||||||
|
const snapshot = deserializeLatestSnapshot([
|
||||||
|
{
|
||||||
|
type: "custom",
|
||||||
|
customType: SNAPSHOT_ENTRY_TYPE,
|
||||||
|
data: {
|
||||||
|
mode: "balanced",
|
||||||
|
lastZone: "red",
|
||||||
|
ledger: { items: [], rollingSummary: "" },
|
||||||
|
},
|
||||||
|
},
|
||||||
|
]);
|
||||||
|
|
||||||
|
assert.ok(snapshot);
|
||||||
|
|
||||||
|
const runtime = createContextManagerRuntime({ mode: "balanced", contextWindow: 500_000 });
|
||||||
|
runtime.restore(snapshot);
|
||||||
|
|
||||||
|
assert.equal(runtime.getSnapshot().lastZone, "green");
|
||||||
|
});
|
||||||
|
|
||||||
|
test("getPolicy returns a clone and restore detaches from external snapshot objects", () => {
|
||||||
|
const runtime = createContextManagerRuntime({ mode: "balanced", contextWindow: 200_000 });
|
||||||
|
|
||||||
|
const policy = runtime.getPolicy();
|
||||||
|
policy.packetTokenCap = 1;
|
||||||
|
policy.redAtTokens = 1;
|
||||||
|
|
||||||
|
const currentPolicy = runtime.getPolicy();
|
||||||
|
assert.equal(currentPolicy.packetTokenCap, 1_200);
|
||||||
|
assert.equal(currentPolicy.redAtTokens, 140_000);
|
||||||
|
|
||||||
|
const snapshot = runtime.getSnapshot();
|
||||||
|
snapshot.mode = "aggressive";
|
||||||
|
snapshot.ledger.rollingSummary = "before restore";
|
||||||
|
|
||||||
|
runtime.restore(snapshot);
|
||||||
|
|
||||||
|
snapshot.mode = "conservative";
|
||||||
|
snapshot.ledger.rollingSummary = "mutated after restore";
|
||||||
|
|
||||||
|
const restored = runtime.getSnapshot();
|
||||||
|
assert.equal(restored.mode, "aggressive");
|
||||||
|
assert.equal(restored.ledger.rollingSummary, "before restore");
|
||||||
|
});
|
||||||
127
.pi/agent/extensions/context-manager/src/runtime.ts
Normal file
127
.pi/agent/extensions/context-manager/src/runtime.ts
Normal file
@@ -0,0 +1,127 @@
|
|||||||
|
import { adjustPolicyForZone, resolvePolicy, zoneForTokens, type ContextMode, type Policy } from "./config.ts";
|
||||||
|
import { extractCandidates, type TranscriptSlice } from "./extract.ts";
|
||||||
|
import { createEmptyLedger, getActiveItems, mergeCandidates } from "./ledger.ts";
|
||||||
|
import { buildContextPacket } from "./packet.ts";
|
||||||
|
import { buildBranchSummary, buildCompactionSummary, buildResumePacket as renderResumePacket } from "./summaries.ts";
|
||||||
|
import type { RuntimeSnapshot } from "./persist.ts";
|
||||||
|
|
||||||
|
function syncSnapshotZone(snapshot: RuntimeSnapshot, policy: Policy): RuntimeSnapshot {
|
||||||
|
if (snapshot.lastObservedTokens === undefined) {
|
||||||
|
return snapshot.lastZone === "green"
|
||||||
|
? snapshot
|
||||||
|
: {
|
||||||
|
...snapshot,
|
||||||
|
lastZone: "green",
|
||||||
|
};
|
||||||
|
}
|
||||||
|
|
||||||
|
return {
|
||||||
|
...snapshot,
|
||||||
|
lastZone: zoneForTokens(snapshot.lastObservedTokens, policy),
|
||||||
|
};
|
||||||
|
}
|
||||||
|
|
||||||
|
export function createContextManagerRuntime(input: { mode?: ContextMode; contextWindow: number }) {
|
||||||
|
let contextWindow = input.contextWindow;
|
||||||
|
let policy = resolvePolicy({ mode: input.mode ?? "balanced", contextWindow });
|
||||||
|
let snapshot: RuntimeSnapshot = {
|
||||||
|
mode: policy.mode,
|
||||||
|
lastZone: "green",
|
||||||
|
ledger: createEmptyLedger(),
|
||||||
|
};
|
||||||
|
|
||||||
|
function applyPolicy(nextPolicy: Policy) {
|
||||||
|
policy = nextPolicy;
|
||||||
|
snapshot = syncSnapshotZone(snapshot, policy);
|
||||||
|
}
|
||||||
|
|
||||||
|
function hasSessionGoal() {
|
||||||
|
return getActiveItems(snapshot.ledger, "goal").some((item) => item.scope === "session" && item.subject === "root-goal");
|
||||||
|
}
|
||||||
|
|
||||||
|
function ingest(slice: TranscriptSlice) {
|
||||||
|
snapshot = {
|
||||||
|
...snapshot,
|
||||||
|
ledger: mergeCandidates(snapshot.ledger, extractCandidates(slice, { hasSessionGoal: hasSessionGoal() })),
|
||||||
|
};
|
||||||
|
}
|
||||||
|
|
||||||
|
function observeTokens(tokens: number) {
|
||||||
|
snapshot = {
|
||||||
|
...snapshot,
|
||||||
|
lastObservedTokens: tokens,
|
||||||
|
lastZone: zoneForTokens(tokens, policy),
|
||||||
|
};
|
||||||
|
}
|
||||||
|
|
||||||
|
function buildPacket() {
|
||||||
|
return buildContextPacket(snapshot.ledger, adjustPolicyForZone(policy, snapshot.lastZone));
|
||||||
|
}
|
||||||
|
|
||||||
|
function mergeArtifact(role: "compaction" | "branchSummary", text: string, entryId: string, timestamp: number) {
|
||||||
|
snapshot = {
|
||||||
|
...snapshot,
|
||||||
|
lastCompactionSummary: role === "compaction" ? text : snapshot.lastCompactionSummary,
|
||||||
|
lastBranchSummary: role === "branchSummary" ? text : snapshot.lastBranchSummary,
|
||||||
|
ledger: mergeCandidates(snapshot.ledger, extractCandidates({ entryId, role, text, timestamp }, { hasSessionGoal: hasSessionGoal() })),
|
||||||
|
};
|
||||||
|
}
|
||||||
|
|
||||||
|
function recordCompactionSummary(text: string, entryId = `compaction-${Date.now()}`, timestamp = Date.now()) {
|
||||||
|
mergeArtifact("compaction", text, entryId, timestamp);
|
||||||
|
}
|
||||||
|
|
||||||
|
function recordBranchSummary(text: string, entryId = `branch-${Date.now()}`, timestamp = Date.now()) {
|
||||||
|
mergeArtifact("branchSummary", text, entryId, timestamp);
|
||||||
|
}
|
||||||
|
|
||||||
|
function buildResumePacket() {
|
||||||
|
const lines: string[] = [];
|
||||||
|
|
||||||
|
if (snapshot.lastCompactionSummary) {
|
||||||
|
lines.push("## Latest compaction handoff", snapshot.lastCompactionSummary, "");
|
||||||
|
}
|
||||||
|
|
||||||
|
if (snapshot.lastBranchSummary) {
|
||||||
|
lines.push("## Latest branch handoff", snapshot.lastBranchSummary, "");
|
||||||
|
}
|
||||||
|
|
||||||
|
const livePacket = renderResumePacket(snapshot.ledger);
|
||||||
|
if (livePacket) {
|
||||||
|
lines.push(livePacket);
|
||||||
|
}
|
||||||
|
|
||||||
|
return lines.join("\n").trim();
|
||||||
|
}
|
||||||
|
|
||||||
|
function setContextWindow(nextContextWindow: number) {
|
||||||
|
contextWindow = Math.max(nextContextWindow, 50_000);
|
||||||
|
applyPolicy(resolvePolicy({ mode: snapshot.mode, contextWindow }));
|
||||||
|
}
|
||||||
|
|
||||||
|
function setMode(mode: ContextMode) {
|
||||||
|
snapshot = { ...snapshot, mode };
|
||||||
|
applyPolicy(resolvePolicy({ mode, contextWindow }));
|
||||||
|
}
|
||||||
|
|
||||||
|
function restore(next: RuntimeSnapshot) {
|
||||||
|
snapshot = structuredClone(next);
|
||||||
|
applyPolicy(resolvePolicy({ mode: snapshot.mode, contextWindow }));
|
||||||
|
}
|
||||||
|
|
||||||
|
return {
|
||||||
|
ingest,
|
||||||
|
observeTokens,
|
||||||
|
buildPacket,
|
||||||
|
buildCompactionSummary: () => buildCompactionSummary(snapshot.ledger),
|
||||||
|
buildBranchSummary: (label: string) => buildBranchSummary(snapshot.ledger, label),
|
||||||
|
buildResumePacket,
|
||||||
|
recordCompactionSummary,
|
||||||
|
recordBranchSummary,
|
||||||
|
setContextWindow,
|
||||||
|
setMode,
|
||||||
|
getPolicy: () => structuredClone(policy),
|
||||||
|
getSnapshot: () => structuredClone(snapshot),
|
||||||
|
restore,
|
||||||
|
};
|
||||||
|
}
|
||||||
139
.pi/agent/extensions/context-manager/src/summaries.test.ts
Normal file
139
.pi/agent/extensions/context-manager/src/summaries.test.ts
Normal file
@@ -0,0 +1,139 @@
|
|||||||
|
import test from "node:test";
|
||||||
|
import assert from "node:assert/strict";
|
||||||
|
import { createEmptyLedger, mergeCandidates } from "./ledger.ts";
|
||||||
|
import {
|
||||||
|
buildBranchSummary,
|
||||||
|
buildBranchSummaryFromEntries,
|
||||||
|
buildCompactionSummary,
|
||||||
|
buildCompactionSummaryFromPreparation,
|
||||||
|
buildResumePacket,
|
||||||
|
} from "./summaries.ts";
|
||||||
|
|
||||||
|
const ledger = mergeCandidates(createEmptyLedger(), [
|
||||||
|
{ kind: "goal", subject: "root-goal", text: "Build a pi context manager", scope: "session", sourceEntryId: "u1", sourceType: "user", timestamp: 1, confidence: 1 },
|
||||||
|
{ kind: "constraint", subject: "must-u1-0", text: "Must adapt to the active model context window.", scope: "session", sourceEntryId: "u1", sourceType: "user", timestamp: 1, confidence: 0.9 },
|
||||||
|
{ kind: "decision", subject: "decision-a1-0", text: "Keep the MVP quiet.", scope: "branch", sourceEntryId: "a1", sourceType: "assistant", timestamp: 2, confidence: 0.9 },
|
||||||
|
{ kind: "activeTask", subject: "next-step-a2-0", text: "Wire hooks into pi.", scope: "branch", sourceEntryId: "a2", sourceType: "assistant", timestamp: 3, confidence: 0.8 },
|
||||||
|
{ kind: "relevantFile", subject: "runtime-ts", text: "src/runtime.ts", scope: "branch", sourceEntryId: "a3", sourceType: "assistant", timestamp: 4, confidence: 0.7 },
|
||||||
|
]);
|
||||||
|
|
||||||
|
test("buildCompactionSummary renders the exact section order and content", () => {
|
||||||
|
const summary = buildCompactionSummary(ledger);
|
||||||
|
|
||||||
|
assert.equal(
|
||||||
|
summary,
|
||||||
|
[
|
||||||
|
"## Goal",
|
||||||
|
"- Build a pi context manager",
|
||||||
|
"",
|
||||||
|
"## Constraints",
|
||||||
|
"- Must adapt to the active model context window.",
|
||||||
|
"",
|
||||||
|
"## Decisions",
|
||||||
|
"- Keep the MVP quiet.",
|
||||||
|
"",
|
||||||
|
"## Active work",
|
||||||
|
"- Wire hooks into pi.",
|
||||||
|
"",
|
||||||
|
"## Relevant files",
|
||||||
|
"- src/runtime.ts",
|
||||||
|
"",
|
||||||
|
"## Next steps",
|
||||||
|
"- Wire hooks into pi.",
|
||||||
|
].join("\n")
|
||||||
|
);
|
||||||
|
});
|
||||||
|
|
||||||
|
test("buildBranchSummary renders the handoff header and sections in order", () => {
|
||||||
|
const summary = buildBranchSummary(ledger, "experimental branch");
|
||||||
|
|
||||||
|
assert.equal(
|
||||||
|
summary,
|
||||||
|
[
|
||||||
|
"# Handoff for experimental branch",
|
||||||
|
"",
|
||||||
|
"## Goal",
|
||||||
|
"- Build a pi context manager",
|
||||||
|
"",
|
||||||
|
"## Decisions",
|
||||||
|
"- Keep the MVP quiet.",
|
||||||
|
"",
|
||||||
|
"## Active work",
|
||||||
|
"- Wire hooks into pi.",
|
||||||
|
"",
|
||||||
|
"## Relevant files",
|
||||||
|
"- src/runtime.ts",
|
||||||
|
].join("\n")
|
||||||
|
);
|
||||||
|
});
|
||||||
|
|
||||||
|
test("buildCompactionSummaryFromPreparation uses preparation messages, previous summary, file ops, and focus text", () => {
|
||||||
|
const previousSummary = [
|
||||||
|
"## Goal",
|
||||||
|
"- Ship the context manager extension",
|
||||||
|
"",
|
||||||
|
"## Key Decisions",
|
||||||
|
"- Keep the public API stable.",
|
||||||
|
].join("\n");
|
||||||
|
|
||||||
|
const summary = buildCompactionSummaryFromPreparation({
|
||||||
|
messagesToSummarize: [
|
||||||
|
{ role: "user", content: "Decision: keep compaction summaries deterministic", timestamp: 1 },
|
||||||
|
{ role: "assistant", content: [{ type: "text", text: "Blocked: verify /tree replaceInstructions behavior" }], timestamp: 2 },
|
||||||
|
],
|
||||||
|
turnPrefixMessages: [
|
||||||
|
{ role: "toolResult", toolName: "read", content: [{ type: "text", text: "Opened .pi/agent/extensions/context-manager/index.ts" }], isError: false, timestamp: 3 },
|
||||||
|
],
|
||||||
|
previousSummary,
|
||||||
|
fileOps: {
|
||||||
|
readFiles: [".pi/agent/extensions/context-manager/index.ts"],
|
||||||
|
modifiedFiles: [".pi/agent/extensions/context-manager/src/summaries.ts"],
|
||||||
|
},
|
||||||
|
customInstructions: "Focus on decisions, blockers, and relevant files.",
|
||||||
|
});
|
||||||
|
|
||||||
|
assert.match(summary, /## Key Decisions/);
|
||||||
|
assert.match(summary, /keep compaction summaries deterministic/);
|
||||||
|
assert.match(summary, /## Open questions and blockers/);
|
||||||
|
assert.match(summary, /verify \/tree replaceInstructions behavior/);
|
||||||
|
assert.match(summary, /<read-files>[\s\S]*index.ts[\s\S]*<\/read-files>/);
|
||||||
|
assert.match(summary, /<modified-files>[\s\S]*src\/summaries.ts[\s\S]*<\/modified-files>/);
|
||||||
|
});
|
||||||
|
|
||||||
|
test("buildBranchSummaryFromEntries uses only the abandoned branch entries and custom focus", () => {
|
||||||
|
const summary = buildBranchSummaryFromEntries({
|
||||||
|
branchLabel: "abandoned branch",
|
||||||
|
entriesToSummarize: [
|
||||||
|
{ type: "message", id: "user-1", parentId: null, timestamp: new Date(1).toISOString(), message: { role: "user", content: "Goal: explore tree handoff" } },
|
||||||
|
{ type: "message", id: "assistant-1", parentId: "user-1", timestamp: new Date(2).toISOString(), message: { role: "assistant", content: [{ type: "text", text: "Decision: do not leak branch-local goals" }] } },
|
||||||
|
],
|
||||||
|
customInstructions: "Focus on goals and decisions.",
|
||||||
|
replaceInstructions: false,
|
||||||
|
commonAncestorId: "root",
|
||||||
|
});
|
||||||
|
|
||||||
|
assert.match(summary, /# Handoff for abandoned branch/);
|
||||||
|
assert.match(summary, /explore tree handoff/);
|
||||||
|
assert.match(summary, /do not leak branch-local goals/);
|
||||||
|
});
|
||||||
|
|
||||||
|
test("buildResumePacket renders restart guidance in the expected order", () => {
|
||||||
|
const summary = buildResumePacket(ledger);
|
||||||
|
|
||||||
|
assert.equal(
|
||||||
|
summary,
|
||||||
|
[
|
||||||
|
"## Goal",
|
||||||
|
"- Build a pi context manager",
|
||||||
|
"",
|
||||||
|
"## Current task",
|
||||||
|
"- Wire hooks into pi.",
|
||||||
|
"",
|
||||||
|
"## Constraints",
|
||||||
|
"- Must adapt to the active model context window.",
|
||||||
|
"",
|
||||||
|
"## Key decisions",
|
||||||
|
"- Keep the MVP quiet.",
|
||||||
|
].join("\n")
|
||||||
|
);
|
||||||
|
});
|
||||||
251
.pi/agent/extensions/context-manager/src/summaries.ts
Normal file
251
.pi/agent/extensions/context-manager/src/summaries.ts
Normal file
@@ -0,0 +1,251 @@
|
|||||||
|
import { createEmptyLedger, getActiveItems, mergeCandidates, type LedgerState } from "./ledger.ts";
|
||||||
|
import { extractCandidates, type TranscriptSlice } from "./extract.ts";
|
||||||
|
|
||||||
|
function lines(title: string, items: string[]) {
|
||||||
|
if (items.length === 0) return [`## ${title}`, "- none", ""];
|
||||||
|
return [`## ${title}`, ...items.map((item) => `- ${item}`), ""];
|
||||||
|
}
|
||||||
|
|
||||||
|
function isTextPart(part: unknown): part is { type: "text"; text?: string } {
|
||||||
|
return typeof part === "object" && part !== null && "type" in part && (part as { type?: unknown }).type === "text";
|
||||||
|
}
|
||||||
|
|
||||||
|
function toText(content: unknown): string {
|
||||||
|
if (typeof content === "string") {
|
||||||
|
return content;
|
||||||
|
}
|
||||||
|
|
||||||
|
if (!Array.isArray(content)) {
|
||||||
|
return "";
|
||||||
|
}
|
||||||
|
|
||||||
|
return content
|
||||||
|
.map((part) => {
|
||||||
|
if (!isTextPart(part)) return "";
|
||||||
|
return typeof part.text === "string" ? part.text : "";
|
||||||
|
})
|
||||||
|
.join("\n")
|
||||||
|
.trim();
|
||||||
|
}
|
||||||
|
|
||||||
|
type FocusKind = "goal" | "constraint" | "decision" | "activeTask" | "openQuestion" | "relevantFile";
|
||||||
|
type SummarySection = { kind: FocusKind; title: string; items: string[] };
|
||||||
|
|
||||||
|
function hasSessionGoal(ledger: LedgerState) {
|
||||||
|
return getActiveItems(ledger, "goal").some((item) => item.scope === "session" && item.subject === "root-goal");
|
||||||
|
}
|
||||||
|
|
||||||
|
function buildLedgerFromSlices(slices: TranscriptSlice[], previousSummary?: string) {
|
||||||
|
let ledger = createEmptyLedger();
|
||||||
|
|
||||||
|
if (previousSummary) {
|
||||||
|
ledger = mergeCandidates(
|
||||||
|
ledger,
|
||||||
|
extractCandidates(
|
||||||
|
{
|
||||||
|
entryId: "previous-summary",
|
||||||
|
role: "compaction",
|
||||||
|
text: previousSummary,
|
||||||
|
timestamp: 0,
|
||||||
|
},
|
||||||
|
{ hasSessionGoal: false },
|
||||||
|
),
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
for (const slice of slices) {
|
||||||
|
ledger = mergeCandidates(ledger, extractCandidates(slice, { hasSessionGoal: hasSessionGoal(ledger) }));
|
||||||
|
}
|
||||||
|
|
||||||
|
return ledger;
|
||||||
|
}
|
||||||
|
|
||||||
|
function parseFocus(customInstructions?: string): Set<FocusKind> {
|
||||||
|
const text = (customInstructions ?? "").toLowerCase();
|
||||||
|
const focus = new Set<FocusKind>();
|
||||||
|
|
||||||
|
if (/\bgoal/.test(text)) focus.add("goal");
|
||||||
|
if (/\bconstraint|preference/.test(text)) focus.add("constraint");
|
||||||
|
if (/\bdecision/.test(text)) focus.add("decision");
|
||||||
|
if (/\btask|next step|progress/.test(text)) focus.add("activeTask");
|
||||||
|
if (/\bblocker|blocked|open question/.test(text)) focus.add("openQuestion");
|
||||||
|
if (/\bfile|read-files|modified-files/.test(text)) focus.add("relevantFile");
|
||||||
|
|
||||||
|
return focus;
|
||||||
|
}
|
||||||
|
|
||||||
|
function buildStructuredSections(ledger: LedgerState): SummarySection[] {
|
||||||
|
return [
|
||||||
|
{ kind: "goal", title: "Goal", items: getActiveItems(ledger, "goal").map((item) => item.text) },
|
||||||
|
{
|
||||||
|
kind: "constraint",
|
||||||
|
title: "Constraints & Preferences",
|
||||||
|
items: getActiveItems(ledger, "constraint").map((item) => item.text),
|
||||||
|
},
|
||||||
|
{
|
||||||
|
kind: "activeTask",
|
||||||
|
title: "Progress",
|
||||||
|
items: getActiveItems(ledger, "activeTask").map((item) => item.text),
|
||||||
|
},
|
||||||
|
{
|
||||||
|
kind: "decision",
|
||||||
|
title: "Key Decisions",
|
||||||
|
items: getActiveItems(ledger, "decision").map((item) => item.text),
|
||||||
|
},
|
||||||
|
{
|
||||||
|
kind: "openQuestion",
|
||||||
|
title: "Open questions and blockers",
|
||||||
|
items: getActiveItems(ledger, "openQuestion").map((item) => item.text),
|
||||||
|
},
|
||||||
|
{
|
||||||
|
kind: "relevantFile",
|
||||||
|
title: "Critical Context",
|
||||||
|
items: getActiveItems(ledger, "relevantFile").map((item) => item.text),
|
||||||
|
},
|
||||||
|
{
|
||||||
|
kind: "activeTask",
|
||||||
|
title: "Next Steps",
|
||||||
|
items: getActiveItems(ledger, "activeTask").map((item) => item.text),
|
||||||
|
},
|
||||||
|
];
|
||||||
|
}
|
||||||
|
|
||||||
|
function sortSectionsForFocus(sections: SummarySection[], focus: Set<FocusKind>): SummarySection[] {
|
||||||
|
if (focus.size === 0) {
|
||||||
|
return sections;
|
||||||
|
}
|
||||||
|
|
||||||
|
return [
|
||||||
|
...sections.filter((section) => focus.has(section.kind)),
|
||||||
|
...sections.filter((section) => !focus.has(section.kind)),
|
||||||
|
];
|
||||||
|
}
|
||||||
|
|
||||||
|
function unique(values: string[]) {
|
||||||
|
return [...new Set(values)];
|
||||||
|
}
|
||||||
|
|
||||||
|
function fileTag(name: "read-files" | "modified-files", values: string[]) {
|
||||||
|
if (values.length === 0) {
|
||||||
|
return [] as string[];
|
||||||
|
}
|
||||||
|
|
||||||
|
return [`<${name}>`, ...values, `</${name}>`, ""];
|
||||||
|
}
|
||||||
|
|
||||||
|
function renderStructuredSummary(
|
||||||
|
ledger: LedgerState,
|
||||||
|
options?: {
|
||||||
|
header?: string;
|
||||||
|
focus?: Set<FocusKind>;
|
||||||
|
readFiles?: string[];
|
||||||
|
modifiedFiles?: string[];
|
||||||
|
},
|
||||||
|
) {
|
||||||
|
const sections = sortSectionsForFocus(buildStructuredSections(ledger), options?.focus ?? new Set());
|
||||||
|
return [
|
||||||
|
...(options?.header ? [options.header, ""] : []),
|
||||||
|
...sections.flatMap((section) => lines(section.title, section.items)),
|
||||||
|
...fileTag("read-files", unique(options?.readFiles ?? [])),
|
||||||
|
...fileTag("modified-files", unique(options?.modifiedFiles ?? [])),
|
||||||
|
].join("\n").trim();
|
||||||
|
}
|
||||||
|
|
||||||
|
function messageToSlice(message: any, entryId: string, timestampFallback: number): TranscriptSlice | undefined {
|
||||||
|
if (!message || typeof message !== "object") {
|
||||||
|
return undefined;
|
||||||
|
}
|
||||||
|
|
||||||
|
if (message.role !== "user" && message.role !== "assistant" && message.role !== "toolResult") {
|
||||||
|
return undefined;
|
||||||
|
}
|
||||||
|
|
||||||
|
return {
|
||||||
|
entryId,
|
||||||
|
role: message.role,
|
||||||
|
text: toText(message.content),
|
||||||
|
timestamp: typeof message.timestamp === "number" ? message.timestamp : timestampFallback,
|
||||||
|
isError: message.role === "toolResult" ? message.isError : undefined,
|
||||||
|
};
|
||||||
|
}
|
||||||
|
|
||||||
|
export function buildCompactionSummary(ledger: LedgerState): string {
|
||||||
|
return [
|
||||||
|
...lines("Goal", getActiveItems(ledger, "goal").map((item) => item.text)),
|
||||||
|
...lines("Constraints", getActiveItems(ledger, "constraint").map((item) => item.text)),
|
||||||
|
...lines("Decisions", getActiveItems(ledger, "decision").map((item) => item.text)),
|
||||||
|
...lines("Active work", getActiveItems(ledger, "activeTask").map((item) => item.text)),
|
||||||
|
...lines("Relevant files", getActiveItems(ledger, "relevantFile").map((item) => item.text)),
|
||||||
|
...lines("Next steps", getActiveItems(ledger, "activeTask").map((item) => item.text)),
|
||||||
|
].join("\n").trim();
|
||||||
|
}
|
||||||
|
|
||||||
|
export function buildBranchSummary(ledger: LedgerState, branchLabel: string): string {
|
||||||
|
return [
|
||||||
|
`# Handoff for ${branchLabel}`,
|
||||||
|
"",
|
||||||
|
...lines("Goal", getActiveItems(ledger, "goal").map((item) => item.text)),
|
||||||
|
...lines("Decisions", getActiveItems(ledger, "decision").map((item) => item.text)),
|
||||||
|
...lines("Active work", getActiveItems(ledger, "activeTask").map((item) => item.text)),
|
||||||
|
...lines("Relevant files", getActiveItems(ledger, "relevantFile").map((item) => item.text)),
|
||||||
|
].join("\n").trim();
|
||||||
|
}
|
||||||
|
|
||||||
|
export function buildResumePacket(ledger: LedgerState): string {
|
||||||
|
return [
|
||||||
|
...lines("Goal", getActiveItems(ledger, "goal").map((item) => item.text)),
|
||||||
|
...lines("Current task", getActiveItems(ledger, "activeTask").map((item) => item.text)),
|
||||||
|
...lines("Constraints", getActiveItems(ledger, "constraint").map((item) => item.text)),
|
||||||
|
...lines("Key decisions", getActiveItems(ledger, "decision").map((item) => item.text)),
|
||||||
|
].join("\n").trim();
|
||||||
|
}
|
||||||
|
|
||||||
|
export function buildCompactionSummaryFromPreparation(input: {
|
||||||
|
messagesToSummarize: any[];
|
||||||
|
turnPrefixMessages: any[];
|
||||||
|
previousSummary?: string;
|
||||||
|
fileOps?: { readFiles?: string[]; modifiedFiles?: string[] };
|
||||||
|
customInstructions?: string;
|
||||||
|
}): string {
|
||||||
|
const slices = [...input.messagesToSummarize, ...input.turnPrefixMessages]
|
||||||
|
.map((message, index) => messageToSlice(message, `compaction-${index}`, index))
|
||||||
|
.filter((slice): slice is TranscriptSlice => Boolean(slice));
|
||||||
|
|
||||||
|
const ledger = buildLedgerFromSlices(slices, input.previousSummary);
|
||||||
|
return renderStructuredSummary(ledger, {
|
||||||
|
focus: parseFocus(input.customInstructions),
|
||||||
|
readFiles: input.fileOps?.readFiles,
|
||||||
|
modifiedFiles: input.fileOps?.modifiedFiles,
|
||||||
|
});
|
||||||
|
}
|
||||||
|
|
||||||
|
export function buildBranchSummaryFromEntries(input: {
|
||||||
|
branchLabel: string;
|
||||||
|
entriesToSummarize: Array<{ type: string; id: string; timestamp: string; message?: any; summary?: string }>;
|
||||||
|
customInstructions?: string;
|
||||||
|
replaceInstructions?: boolean;
|
||||||
|
commonAncestorId?: string | null;
|
||||||
|
}): string {
|
||||||
|
const slices = input.entriesToSummarize.flatMap((entry) => {
|
||||||
|
if (entry.type === "message") {
|
||||||
|
const slice = messageToSlice(entry.message, entry.id, Date.parse(entry.timestamp));
|
||||||
|
return slice ? [slice] : [];
|
||||||
|
}
|
||||||
|
|
||||||
|
if (entry.type === "compaction" && typeof entry.summary === "string") {
|
||||||
|
return [{ entryId: entry.id, role: "compaction", text: entry.summary, timestamp: Date.parse(entry.timestamp) } satisfies TranscriptSlice];
|
||||||
|
}
|
||||||
|
|
||||||
|
if (entry.type === "branch_summary" && typeof entry.summary === "string") {
|
||||||
|
return [{ entryId: entry.id, role: "branchSummary", text: entry.summary, timestamp: Date.parse(entry.timestamp) } satisfies TranscriptSlice];
|
||||||
|
}
|
||||||
|
|
||||||
|
return [];
|
||||||
|
});
|
||||||
|
|
||||||
|
const ledger = buildLedgerFromSlices(slices);
|
||||||
|
return renderStructuredSummary(ledger, {
|
||||||
|
header: `# Handoff for ${input.branchLabel}`,
|
||||||
|
focus: parseFocus(input.customInstructions),
|
||||||
|
});
|
||||||
|
}
|
||||||
50
.pi/agent/extensions/dev-tools/index.ts
Normal file
50
.pi/agent/extensions/dev-tools/index.ts
Normal file
@@ -0,0 +1,50 @@
|
|||||||
|
import type { ExtensionAPI } from "@mariozechner/pi-coding-agent";
|
||||||
|
import { createCommandFormatterRunner } from "./src/formatting/command-runner.ts";
|
||||||
|
import { createCommandDiagnosticsBackend } from "./src/diagnostics/command-backend.ts";
|
||||||
|
import { createLspClientManager } from "./src/diagnostics/lsp-client.ts";
|
||||||
|
import { createSetupSuggestTool } from "./src/tools/setup-suggest.ts";
|
||||||
|
import { probeProject } from "./src/project-probe.ts";
|
||||||
|
import { createFormattedWriteTool } from "./src/tools/write.ts";
|
||||||
|
import { createFormattedEditTool } from "./src/tools/edit.ts";
|
||||||
|
import { createDevToolsRuntime } from "./src/runtime.ts";
|
||||||
|
|
||||||
|
export default function devTools(pi: ExtensionAPI) {
|
||||||
|
const cwd = process.cwd();
|
||||||
|
const agentDir = process.env.PI_CODING_AGENT_DIR ?? `${process.env.HOME}/.pi/agent`;
|
||||||
|
|
||||||
|
const runtime = createDevToolsRuntime({
|
||||||
|
cwd,
|
||||||
|
agentDir,
|
||||||
|
formatterRunner: createCommandFormatterRunner({
|
||||||
|
execCommand: async (command, args, options) => {
|
||||||
|
const result = await pi.exec(command, args, { timeout: options.timeout });
|
||||||
|
return { code: result.code ?? 0, stdout: result.stdout ?? "", stderr: result.stderr ?? "" };
|
||||||
|
},
|
||||||
|
}),
|
||||||
|
commandBackend: createCommandDiagnosticsBackend({
|
||||||
|
execCommand: async (command, args, options) => {
|
||||||
|
const result = await pi.exec(command, args, { timeout: options.timeout });
|
||||||
|
return { code: result.code ?? 0, stdout: result.stdout ?? "", stderr: result.stderr ?? "" };
|
||||||
|
},
|
||||||
|
}),
|
||||||
|
lspBackend: createLspClientManager(),
|
||||||
|
probeProject,
|
||||||
|
});
|
||||||
|
|
||||||
|
pi.registerTool(createFormattedEditTool(cwd, runtime));
|
||||||
|
pi.registerTool(createFormattedWriteTool(cwd, runtime));
|
||||||
|
pi.registerTool(createSetupSuggestTool({
|
||||||
|
suggestSetup: async () => {
|
||||||
|
const probe = await probeProject({ cwd });
|
||||||
|
return probe.summary;
|
||||||
|
},
|
||||||
|
}));
|
||||||
|
|
||||||
|
pi.on("before_agent_start", async (event) => {
|
||||||
|
const block = runtime.getPromptBlock();
|
||||||
|
if (!block) return;
|
||||||
|
return {
|
||||||
|
systemPrompt: `${event.systemPrompt}\n\n${block}`,
|
||||||
|
};
|
||||||
|
});
|
||||||
|
}
|
||||||
4386
.pi/agent/extensions/dev-tools/package-lock.json
generated
Normal file
4386
.pi/agent/extensions/dev-tools/package-lock.json
generated
Normal file
File diff suppressed because it is too large
Load Diff
23
.pi/agent/extensions/dev-tools/package.json
Normal file
23
.pi/agent/extensions/dev-tools/package.json
Normal file
@@ -0,0 +1,23 @@
|
|||||||
|
{
|
||||||
|
"name": "pi-dev-tools-extension",
|
||||||
|
"private": true,
|
||||||
|
"version": "0.0.0",
|
||||||
|
"type": "module",
|
||||||
|
"scripts": {
|
||||||
|
"test": "tsx --test src/*.test.ts src/**/*.test.ts"
|
||||||
|
},
|
||||||
|
"pi": {
|
||||||
|
"extensions": ["./index.ts"]
|
||||||
|
},
|
||||||
|
"dependencies": {
|
||||||
|
"@sinclair/typebox": "^0.34.49",
|
||||||
|
"picomatch": "^4.0.2",
|
||||||
|
"vscode-jsonrpc": "^8.2.1"
|
||||||
|
},
|
||||||
|
"devDependencies": {
|
||||||
|
"@mariozechner/pi-coding-agent": "^0.66.1",
|
||||||
|
"@types/node": "^25.5.2",
|
||||||
|
"tsx": "^4.21.0",
|
||||||
|
"typescript": "^6.0.2"
|
||||||
|
}
|
||||||
|
}
|
||||||
37
.pi/agent/extensions/dev-tools/src/config.test.ts
Normal file
37
.pi/agent/extensions/dev-tools/src/config.test.ts
Normal file
@@ -0,0 +1,37 @@
|
|||||||
|
import test from "node:test";
|
||||||
|
import assert from "node:assert/strict";
|
||||||
|
import { mergeDevToolsConfig } from "./config.ts";
|
||||||
|
|
||||||
|
test("mergeDevToolsConfig lets project defaults override global defaults and replace same-name profiles", () => {
|
||||||
|
const merged = mergeDevToolsConfig(
|
||||||
|
{
|
||||||
|
defaults: { formatTimeoutMs: 8000, maxDiagnosticsPerFile: 10 },
|
||||||
|
profiles: [
|
||||||
|
{
|
||||||
|
name: "typescript",
|
||||||
|
match: ["**/*.ts"],
|
||||||
|
workspaceRootMarkers: ["package.json"],
|
||||||
|
formatter: { kind: "command", command: ["prettier", "--write", "{file}"] },
|
||||||
|
diagnostics: [],
|
||||||
|
},
|
||||||
|
],
|
||||||
|
},
|
||||||
|
{
|
||||||
|
defaults: { formatTimeoutMs: 3000 },
|
||||||
|
profiles: [
|
||||||
|
{
|
||||||
|
name: "typescript",
|
||||||
|
match: ["src/**/*.ts"],
|
||||||
|
workspaceRootMarkers: ["tsconfig.json"],
|
||||||
|
formatter: { kind: "command", command: ["biome", "format", "--write", "{file}"] },
|
||||||
|
diagnostics: [],
|
||||||
|
},
|
||||||
|
],
|
||||||
|
},
|
||||||
|
);
|
||||||
|
|
||||||
|
assert.equal(merged.defaults.formatTimeoutMs, 3000);
|
||||||
|
assert.equal(merged.defaults.maxDiagnosticsPerFile, 10);
|
||||||
|
assert.deepEqual(merged.profiles.map((profile) => profile.name), ["typescript"]);
|
||||||
|
assert.deepEqual(merged.profiles[0]?.match, ["src/**/*.ts"]);
|
||||||
|
});
|
||||||
38
.pi/agent/extensions/dev-tools/src/config.ts
Normal file
38
.pi/agent/extensions/dev-tools/src/config.ts
Normal file
@@ -0,0 +1,38 @@
|
|||||||
|
import { existsSync, readFileSync } from "node:fs";
|
||||||
|
import { resolve } from "node:path";
|
||||||
|
import { Value } from "@sinclair/typebox/value";
|
||||||
|
import { DevToolsConfigSchema, type DevToolsConfig } from "./schema.ts";
|
||||||
|
|
||||||
|
export function mergeDevToolsConfig(globalConfig?: DevToolsConfig, projectConfig?: DevToolsConfig): DevToolsConfig {
|
||||||
|
const defaults = {
|
||||||
|
...(globalConfig?.defaults ?? {}),
|
||||||
|
...(projectConfig?.defaults ?? {}),
|
||||||
|
};
|
||||||
|
|
||||||
|
const globalProfiles = new Map((globalConfig?.profiles ?? []).map((profile) => [profile.name, profile]));
|
||||||
|
const mergedProfiles = [...(projectConfig?.profiles ?? [])];
|
||||||
|
|
||||||
|
for (const profile of globalProfiles.values()) {
|
||||||
|
if (!mergedProfiles.some((candidate) => candidate.name === profile.name)) {
|
||||||
|
mergedProfiles.push(profile);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
return { defaults, profiles: mergedProfiles };
|
||||||
|
}
|
||||||
|
|
||||||
|
function readConfigIfPresent(path: string): DevToolsConfig | undefined {
|
||||||
|
if (!existsSync(path)) return undefined;
|
||||||
|
const parsed = JSON.parse(readFileSync(path, "utf8"));
|
||||||
|
if (!Value.Check(DevToolsConfigSchema, parsed)) {
|
||||||
|
const [firstError] = [...Value.Errors(DevToolsConfigSchema, parsed)];
|
||||||
|
throw new Error(`Invalid dev-tools config at ${path}: ${firstError?.message ?? "validation failed"}`);
|
||||||
|
}
|
||||||
|
return parsed as DevToolsConfig;
|
||||||
|
}
|
||||||
|
|
||||||
|
export function loadDevToolsConfig(cwd: string, agentDir: string): DevToolsConfig | undefined {
|
||||||
|
const globalPath = resolve(agentDir, "dev-tools.json");
|
||||||
|
const projectPath = resolve(cwd, ".pi/dev-tools.json");
|
||||||
|
return mergeDevToolsConfig(readConfigIfPresent(globalPath), readConfigIfPresent(projectPath));
|
||||||
|
}
|
||||||
@@ -0,0 +1,40 @@
|
|||||||
|
import test from "node:test";
|
||||||
|
import assert from "node:assert/strict";
|
||||||
|
import { createCommandDiagnosticsBackend } from "./command-backend.ts";
|
||||||
|
|
||||||
|
test("eslint-json parser returns normalized diagnostics", async () => {
|
||||||
|
const backend = createCommandDiagnosticsBackend({
|
||||||
|
execCommand: async () => ({
|
||||||
|
code: 1,
|
||||||
|
stdout: JSON.stringify([
|
||||||
|
{
|
||||||
|
filePath: "/repo/src/app.ts",
|
||||||
|
messages: [
|
||||||
|
{
|
||||||
|
ruleId: "no-console",
|
||||||
|
severity: 2,
|
||||||
|
message: "Unexpected console statement.",
|
||||||
|
line: 2,
|
||||||
|
column: 3,
|
||||||
|
},
|
||||||
|
],
|
||||||
|
},
|
||||||
|
]),
|
||||||
|
stderr: "",
|
||||||
|
}),
|
||||||
|
});
|
||||||
|
|
||||||
|
const result = await backend.collect({
|
||||||
|
absolutePath: "/repo/src/app.ts",
|
||||||
|
workspaceRoot: "/repo",
|
||||||
|
backend: {
|
||||||
|
kind: "command",
|
||||||
|
parser: "eslint-json",
|
||||||
|
command: ["eslint", "--format", "json", "{file}"],
|
||||||
|
},
|
||||||
|
});
|
||||||
|
|
||||||
|
assert.equal(result.status, "ok");
|
||||||
|
assert.equal(result.items[0]?.severity, "error");
|
||||||
|
assert.equal(result.items[0]?.message, "Unexpected console statement.");
|
||||||
|
});
|
||||||
@@ -0,0 +1,45 @@
|
|||||||
|
import type { DiagnosticsConfig } from "../schema.ts";
|
||||||
|
import type { DiagnosticsState, NormalizedDiagnostic } from "./types.ts";
|
||||||
|
|
||||||
|
function parseEslintJson(stdout: string): NormalizedDiagnostic[] {
|
||||||
|
const parsed = JSON.parse(stdout) as Array<any>;
|
||||||
|
return parsed.flatMap((entry) =>
|
||||||
|
(entry.messages ?? []).map((message: any) => ({
|
||||||
|
severity: message.severity === 2 ? "error" : "warning",
|
||||||
|
message: message.message,
|
||||||
|
line: message.line,
|
||||||
|
column: message.column,
|
||||||
|
source: "eslint",
|
||||||
|
code: message.ruleId ?? undefined,
|
||||||
|
})),
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
export function createCommandDiagnosticsBackend(deps: {
|
||||||
|
execCommand: (
|
||||||
|
command: string,
|
||||||
|
args: string[],
|
||||||
|
options: { cwd: string; timeout?: number },
|
||||||
|
) => Promise<{ code: number; stdout: string; stderr: string }>;
|
||||||
|
}) {
|
||||||
|
return {
|
||||||
|
async collect(input: {
|
||||||
|
absolutePath: string;
|
||||||
|
workspaceRoot: string;
|
||||||
|
backend: Extract<DiagnosticsConfig, { kind: "command" }>;
|
||||||
|
timeoutMs?: number;
|
||||||
|
}): Promise<DiagnosticsState> {
|
||||||
|
const [command, ...args] = input.backend.command.map((part) => part.replaceAll("{file}", input.absolutePath));
|
||||||
|
const result = await deps.execCommand(command, args, { cwd: input.workspaceRoot, timeout: input.timeoutMs });
|
||||||
|
|
||||||
|
try {
|
||||||
|
if (input.backend.parser === "eslint-json") {
|
||||||
|
return { status: "ok", items: parseEslintJson(result.stdout) };
|
||||||
|
}
|
||||||
|
return { status: "unavailable", items: [], message: `Unsupported diagnostics parser: ${input.backend.parser}` };
|
||||||
|
} catch (error) {
|
||||||
|
return { status: "unavailable", items: [], message: (error as Error).message };
|
||||||
|
}
|
||||||
|
},
|
||||||
|
};
|
||||||
|
}
|
||||||
@@ -0,0 +1,41 @@
|
|||||||
|
import test from "node:test";
|
||||||
|
import assert from "node:assert/strict";
|
||||||
|
import { createLspClientManager } from "./lsp-client.ts";
|
||||||
|
|
||||||
|
test("collectForFile sends initialize + didOpen and resolves publishDiagnostics", async () => {
|
||||||
|
const notifications: Array<{ method: string; params: any }> = [];
|
||||||
|
|
||||||
|
const manager = createLspClientManager({
|
||||||
|
createConnection: async () => ({
|
||||||
|
async initialize() {},
|
||||||
|
async openTextDocument(params) {
|
||||||
|
notifications.push({ method: "textDocument/didOpen", params });
|
||||||
|
},
|
||||||
|
async collectDiagnostics() {
|
||||||
|
return [
|
||||||
|
{
|
||||||
|
severity: "error",
|
||||||
|
message: "Type 'number' is not assignable to type 'string'.",
|
||||||
|
line: 1,
|
||||||
|
column: 7,
|
||||||
|
source: "tsserver",
|
||||||
|
},
|
||||||
|
];
|
||||||
|
},
|
||||||
|
async dispose() {},
|
||||||
|
}),
|
||||||
|
});
|
||||||
|
|
||||||
|
const result = await manager.collectForFile({
|
||||||
|
key: "typescript:/repo",
|
||||||
|
absolutePath: "/repo/src/app.ts",
|
||||||
|
workspaceRoot: "/repo",
|
||||||
|
languageId: "typescript",
|
||||||
|
text: "const x: string = 1\n",
|
||||||
|
command: ["typescript-language-server", "--stdio"],
|
||||||
|
});
|
||||||
|
|
||||||
|
assert.equal(result.status, "ok");
|
||||||
|
assert.equal(result.items[0]?.source, "tsserver");
|
||||||
|
assert.equal(notifications[0]?.method, "textDocument/didOpen");
|
||||||
|
});
|
||||||
102
.pi/agent/extensions/dev-tools/src/diagnostics/lsp-client.ts
Normal file
102
.pi/agent/extensions/dev-tools/src/diagnostics/lsp-client.ts
Normal file
@@ -0,0 +1,102 @@
|
|||||||
|
import { spawn } from "node:child_process";
|
||||||
|
import { pathToFileURL } from "node:url";
|
||||||
|
import * as rpc from "vscode-jsonrpc/node";
|
||||||
|
import type { DiagnosticsState } from "./types.ts";
|
||||||
|
|
||||||
|
const INITIALIZE = new rpc.RequestType<any, any, void, void>("initialize");
|
||||||
|
const DID_OPEN = new rpc.NotificationType<any, void>("textDocument/didOpen");
|
||||||
|
const INITIALIZED = new rpc.NotificationType<any, void>("initialized");
|
||||||
|
const PUBLISH_DIAGNOSTICS = new rpc.NotificationType<any, void>("textDocument/publishDiagnostics");
|
||||||
|
|
||||||
|
type LspConnection = {
|
||||||
|
initialize(): Promise<void>;
|
||||||
|
openTextDocument(params: any): Promise<void>;
|
||||||
|
collectDiagnostics(): Promise<DiagnosticsState["items"]>;
|
||||||
|
dispose(): Promise<void>;
|
||||||
|
};
|
||||||
|
|
||||||
|
export function createLspClientManager(deps: {
|
||||||
|
createConnection?: (input: { workspaceRoot: string; command: string[] }) => Promise<LspConnection>;
|
||||||
|
} = {}) {
|
||||||
|
const clients = new Map<string, LspConnection>();
|
||||||
|
|
||||||
|
async function defaultCreateConnection(input: { workspaceRoot: string; command: string[] }): Promise<LspConnection> {
|
||||||
|
const [command, ...args] = input.command;
|
||||||
|
const child = spawn(command, args, {
|
||||||
|
cwd: input.workspaceRoot,
|
||||||
|
stdio: ["pipe", "pipe", "pipe"],
|
||||||
|
});
|
||||||
|
|
||||||
|
const connection = rpc.createMessageConnection(
|
||||||
|
new rpc.StreamMessageReader(child.stdout),
|
||||||
|
new rpc.StreamMessageWriter(child.stdin),
|
||||||
|
);
|
||||||
|
|
||||||
|
let lastDiagnostics: DiagnosticsState["items"] = [];
|
||||||
|
connection.onNotification(PUBLISH_DIAGNOSTICS, (params: any) => {
|
||||||
|
lastDiagnostics = (params.diagnostics ?? []).map((diagnostic: any) => ({
|
||||||
|
severity: diagnostic.severity === 1 ? "error" : diagnostic.severity === 2 ? "warning" : "info",
|
||||||
|
message: diagnostic.message,
|
||||||
|
line: diagnostic.range?.start?.line !== undefined ? diagnostic.range.start.line + 1 : undefined,
|
||||||
|
column: diagnostic.range?.start?.character !== undefined ? diagnostic.range.start.character + 1 : undefined,
|
||||||
|
source: diagnostic.source ?? "lsp",
|
||||||
|
code: diagnostic.code ? String(diagnostic.code) : undefined,
|
||||||
|
}));
|
||||||
|
});
|
||||||
|
|
||||||
|
connection.listen();
|
||||||
|
await connection.sendRequest(INITIALIZE, {
|
||||||
|
processId: process.pid,
|
||||||
|
rootUri: pathToFileURL(input.workspaceRoot).href,
|
||||||
|
capabilities: {},
|
||||||
|
});
|
||||||
|
connection.sendNotification(INITIALIZED, {});
|
||||||
|
|
||||||
|
return {
|
||||||
|
async initialize() {},
|
||||||
|
async openTextDocument(params: any) {
|
||||||
|
connection.sendNotification(DID_OPEN, params);
|
||||||
|
},
|
||||||
|
async collectDiagnostics() {
|
||||||
|
await new Promise((resolve) => setTimeout(resolve, 100));
|
||||||
|
return lastDiagnostics;
|
||||||
|
},
|
||||||
|
async dispose() {
|
||||||
|
connection.dispose();
|
||||||
|
child.kill();
|
||||||
|
},
|
||||||
|
};
|
||||||
|
}
|
||||||
|
|
||||||
|
return {
|
||||||
|
async collectForFile(input: {
|
||||||
|
key: string;
|
||||||
|
absolutePath: string;
|
||||||
|
workspaceRoot: string;
|
||||||
|
languageId: string;
|
||||||
|
text: string;
|
||||||
|
command: string[];
|
||||||
|
}): Promise<DiagnosticsState> {
|
||||||
|
let client = clients.get(input.key);
|
||||||
|
if (!client) {
|
||||||
|
client = await (deps.createConnection ?? defaultCreateConnection)({
|
||||||
|
workspaceRoot: input.workspaceRoot,
|
||||||
|
command: input.command,
|
||||||
|
});
|
||||||
|
clients.set(input.key, client);
|
||||||
|
await client.initialize();
|
||||||
|
}
|
||||||
|
|
||||||
|
await client.openTextDocument({
|
||||||
|
textDocument: {
|
||||||
|
uri: pathToFileURL(input.absolutePath).href,
|
||||||
|
languageId: input.languageId,
|
||||||
|
version: 1,
|
||||||
|
text: input.text,
|
||||||
|
},
|
||||||
|
});
|
||||||
|
|
||||||
|
return { status: "ok", items: await client.collectDiagnostics() };
|
||||||
|
},
|
||||||
|
};
|
||||||
|
}
|
||||||
19
.pi/agent/extensions/dev-tools/src/diagnostics/types.ts
Normal file
19
.pi/agent/extensions/dev-tools/src/diagnostics/types.ts
Normal file
@@ -0,0 +1,19 @@
|
|||||||
|
export interface NormalizedDiagnostic {
|
||||||
|
severity: "error" | "warning" | "info";
|
||||||
|
message: string;
|
||||||
|
line?: number;
|
||||||
|
column?: number;
|
||||||
|
source: string;
|
||||||
|
code?: string;
|
||||||
|
}
|
||||||
|
|
||||||
|
export interface DiagnosticsState {
|
||||||
|
status: "ok" | "unavailable";
|
||||||
|
items: NormalizedDiagnostic[];
|
||||||
|
message?: string;
|
||||||
|
}
|
||||||
|
|
||||||
|
export interface CapabilityGap {
|
||||||
|
path: string;
|
||||||
|
message: string;
|
||||||
|
}
|
||||||
16
.pi/agent/extensions/dev-tools/src/extension.test.ts
Normal file
16
.pi/agent/extensions/dev-tools/src/extension.test.ts
Normal file
@@ -0,0 +1,16 @@
|
|||||||
|
import test from "node:test";
|
||||||
|
import assert from "node:assert/strict";
|
||||||
|
import devToolsExtension from "../index.ts";
|
||||||
|
|
||||||
|
test("the extension entrypoint registers edit, write, and setup suggestion tools", () => {
|
||||||
|
const registeredTools: string[] = [];
|
||||||
|
|
||||||
|
devToolsExtension({
|
||||||
|
registerTool(tool: { name: string }) {
|
||||||
|
registeredTools.push(tool.name);
|
||||||
|
},
|
||||||
|
on() {},
|
||||||
|
} as any);
|
||||||
|
|
||||||
|
assert.deepEqual(registeredTools.sort(), ["dev_tools_suggest_setup", "edit", "write"]);
|
||||||
|
});
|
||||||
@@ -0,0 +1,27 @@
|
|||||||
|
import test from "node:test";
|
||||||
|
import assert from "node:assert/strict";
|
||||||
|
import { createCommandFormatterRunner } from "./command-runner.ts";
|
||||||
|
|
||||||
|
test("formatFile expands {file} and executes in the workspace root", async () => {
|
||||||
|
let captured: { command: string; args: string[]; cwd?: string } | undefined;
|
||||||
|
|
||||||
|
const runner = createCommandFormatterRunner({
|
||||||
|
execCommand: async (command, args, options) => {
|
||||||
|
captured = { command, args, cwd: options.cwd };
|
||||||
|
return { code: 0, stdout: "", stderr: "" };
|
||||||
|
},
|
||||||
|
});
|
||||||
|
|
||||||
|
const result = await runner.formatFile({
|
||||||
|
absolutePath: "/repo/src/app.ts",
|
||||||
|
workspaceRoot: "/repo",
|
||||||
|
formatter: { kind: "command", command: ["biome", "format", "--write", "{file}"] },
|
||||||
|
});
|
||||||
|
|
||||||
|
assert.equal(result.status, "formatted");
|
||||||
|
assert.deepEqual(captured, {
|
||||||
|
command: "biome",
|
||||||
|
args: ["format", "--write", "/repo/src/app.ts"],
|
||||||
|
cwd: "/repo",
|
||||||
|
});
|
||||||
|
});
|
||||||
@@ -0,0 +1,33 @@
|
|||||||
|
import type { FormatterConfig } from "../schema.ts";
|
||||||
|
|
||||||
|
export function createCommandFormatterRunner(deps: {
|
||||||
|
execCommand: (
|
||||||
|
command: string,
|
||||||
|
args: string[],
|
||||||
|
options: { cwd: string; timeout?: number },
|
||||||
|
) => Promise<{ code: number; stdout: string; stderr: string }>;
|
||||||
|
}) {
|
||||||
|
return {
|
||||||
|
async formatFile(input: {
|
||||||
|
absolutePath: string;
|
||||||
|
workspaceRoot: string;
|
||||||
|
formatter: FormatterConfig;
|
||||||
|
timeoutMs?: number;
|
||||||
|
}) {
|
||||||
|
const [command, ...args] = input.formatter.command.map((part) => part.replaceAll("{file}", input.absolutePath));
|
||||||
|
const result = await deps.execCommand(command, args, {
|
||||||
|
cwd: input.workspaceRoot,
|
||||||
|
timeout: input.timeoutMs,
|
||||||
|
});
|
||||||
|
|
||||||
|
if (result.code !== 0) {
|
||||||
|
return {
|
||||||
|
status: "failed" as const,
|
||||||
|
message: (result.stderr || result.stdout || `formatter exited with ${result.code}`).trim(),
|
||||||
|
};
|
||||||
|
}
|
||||||
|
|
||||||
|
return { status: "formatted" as const };
|
||||||
|
},
|
||||||
|
};
|
||||||
|
}
|
||||||
26
.pi/agent/extensions/dev-tools/src/profiles.test.ts
Normal file
26
.pi/agent/extensions/dev-tools/src/profiles.test.ts
Normal file
@@ -0,0 +1,26 @@
|
|||||||
|
import test from "node:test";
|
||||||
|
import assert from "node:assert/strict";
|
||||||
|
import { resolveProfileForPath } from "./profiles.ts";
|
||||||
|
|
||||||
|
test("resolveProfileForPath finds the first matching profile and nearest workspace root", () => {
|
||||||
|
const result = resolveProfileForPath(
|
||||||
|
{
|
||||||
|
defaults: {},
|
||||||
|
profiles: [
|
||||||
|
{
|
||||||
|
name: "typescript",
|
||||||
|
match: ["src/**/*.ts"],
|
||||||
|
workspaceRootMarkers: ["package.json", "tsconfig.json"],
|
||||||
|
formatter: { kind: "command", command: ["biome", "format", "--write", "{file}"] },
|
||||||
|
diagnostics: [],
|
||||||
|
},
|
||||||
|
],
|
||||||
|
},
|
||||||
|
"/repo/src/app.ts",
|
||||||
|
"/repo",
|
||||||
|
["/repo/package.json", "/repo/src/app.ts"],
|
||||||
|
);
|
||||||
|
|
||||||
|
assert.equal(result?.profile.name, "typescript");
|
||||||
|
assert.equal(result?.workspaceRoot, "/repo");
|
||||||
|
});
|
||||||
47
.pi/agent/extensions/dev-tools/src/profiles.ts
Normal file
47
.pi/agent/extensions/dev-tools/src/profiles.ts
Normal file
@@ -0,0 +1,47 @@
|
|||||||
|
import { dirname, relative, resolve } from "node:path";
|
||||||
|
import picomatch from "picomatch";
|
||||||
|
import type { DevToolsConfig, DevToolsProfile } from "./schema.ts";
|
||||||
|
|
||||||
|
export interface ResolvedProfileMatch {
|
||||||
|
profile: DevToolsProfile;
|
||||||
|
workspaceRoot: string;
|
||||||
|
}
|
||||||
|
|
||||||
|
export function resolveProfileForPath(
|
||||||
|
config: DevToolsConfig,
|
||||||
|
absolutePath: string,
|
||||||
|
cwd: string,
|
||||||
|
knownPaths: string[] = [],
|
||||||
|
): ResolvedProfileMatch | undefined {
|
||||||
|
const normalizedPath = resolve(absolutePath);
|
||||||
|
const relativePath = relative(cwd, normalizedPath).replace(/\\/g, "/");
|
||||||
|
|
||||||
|
for (const profile of config.profiles) {
|
||||||
|
if (!profile.match.some((pattern) => picomatch(pattern)(relativePath))) {
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
|
||||||
|
const workspaceRoot = findWorkspaceRoot(normalizedPath, cwd, profile.workspaceRootMarkers, knownPaths);
|
||||||
|
return { profile, workspaceRoot };
|
||||||
|
}
|
||||||
|
|
||||||
|
return undefined;
|
||||||
|
}
|
||||||
|
|
||||||
|
function findWorkspaceRoot(filePath: string, cwd: string, markers: string[], knownPaths: string[]): string {
|
||||||
|
let current = dirname(filePath);
|
||||||
|
const root = resolve(cwd);
|
||||||
|
|
||||||
|
while (current.startsWith(root)) {
|
||||||
|
for (const marker of markers) {
|
||||||
|
if (knownPaths.includes(resolve(current, marker))) {
|
||||||
|
return current;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
const next = dirname(current);
|
||||||
|
if (next === current) break;
|
||||||
|
current = next;
|
||||||
|
}
|
||||||
|
|
||||||
|
return root;
|
||||||
|
}
|
||||||
14
.pi/agent/extensions/dev-tools/src/project-probe.test.ts
Normal file
14
.pi/agent/extensions/dev-tools/src/project-probe.test.ts
Normal file
@@ -0,0 +1,14 @@
|
|||||||
|
import test from "node:test";
|
||||||
|
import assert from "node:assert/strict";
|
||||||
|
import { probeProject } from "./project-probe.ts";
|
||||||
|
|
||||||
|
test("probeProject recognizes a TypeScript workspace and suggests biome + tsserver", async () => {
|
||||||
|
const result = await probeProject({
|
||||||
|
cwd: "/repo",
|
||||||
|
exists: async (path) => ["/repo/package.json", "/repo/tsconfig.json"].includes(path),
|
||||||
|
});
|
||||||
|
|
||||||
|
assert.equal(result.ecosystem, "typescript");
|
||||||
|
assert.match(result.summary, /Biome/);
|
||||||
|
assert.match(result.summary, /typescript-language-server/);
|
||||||
|
});
|
||||||
68
.pi/agent/extensions/dev-tools/src/project-probe.ts
Normal file
68
.pi/agent/extensions/dev-tools/src/project-probe.ts
Normal file
@@ -0,0 +1,68 @@
|
|||||||
|
import { access } from "node:fs/promises";
|
||||||
|
import { resolve } from "node:path";
|
||||||
|
|
||||||
|
export interface ProjectProbeResult {
|
||||||
|
ecosystem: string;
|
||||||
|
summary: string;
|
||||||
|
}
|
||||||
|
|
||||||
|
export async function probeProject(deps: {
|
||||||
|
cwd: string;
|
||||||
|
exists?: (path: string) => Promise<boolean>;
|
||||||
|
}): Promise<ProjectProbeResult> {
|
||||||
|
const exists = deps.exists ?? (async (path: string) => {
|
||||||
|
try {
|
||||||
|
await access(path);
|
||||||
|
return true;
|
||||||
|
} catch {
|
||||||
|
return false;
|
||||||
|
}
|
||||||
|
});
|
||||||
|
const cwd = resolve(deps.cwd);
|
||||||
|
|
||||||
|
const hasPackageJson = await exists(resolve(cwd, "package.json"));
|
||||||
|
const hasTsconfig = await exists(resolve(cwd, "tsconfig.json"));
|
||||||
|
const hasPyproject = await exists(resolve(cwd, "pyproject.toml"));
|
||||||
|
const hasCargo = await exists(resolve(cwd, "Cargo.toml"));
|
||||||
|
const hasGoMod = await exists(resolve(cwd, "go.mod"));
|
||||||
|
|
||||||
|
if (hasPackageJson && hasTsconfig) {
|
||||||
|
return {
|
||||||
|
ecosystem: "typescript",
|
||||||
|
summary: "TypeScript project detected. Recommended: Biome for formatting/linting and typescript-language-server for diagnostics.",
|
||||||
|
};
|
||||||
|
}
|
||||||
|
|
||||||
|
if (hasPyproject) {
|
||||||
|
return {
|
||||||
|
ecosystem: "python",
|
||||||
|
summary: "Python project detected. Recommended: Ruff for formatting/linting and basedpyright or pylsp for diagnostics.",
|
||||||
|
};
|
||||||
|
}
|
||||||
|
|
||||||
|
if (hasCargo) {
|
||||||
|
return {
|
||||||
|
ecosystem: "rust",
|
||||||
|
summary: "Rust project detected. Recommended: rustfmt + cargo clippy and rust-analyzer.",
|
||||||
|
};
|
||||||
|
}
|
||||||
|
|
||||||
|
if (hasGoMod) {
|
||||||
|
return {
|
||||||
|
ecosystem: "go",
|
||||||
|
summary: "Go project detected. Recommended: gofmt/goimports and gopls.",
|
||||||
|
};
|
||||||
|
}
|
||||||
|
|
||||||
|
if (hasPackageJson) {
|
||||||
|
return {
|
||||||
|
ecosystem: "javascript",
|
||||||
|
summary: "JavaScript project detected. Recommended: Biome or Prettier+ESLint, plus TypeScript language tooling if applicable.",
|
||||||
|
};
|
||||||
|
}
|
||||||
|
|
||||||
|
return {
|
||||||
|
ecosystem: "unknown",
|
||||||
|
summary: "No known project toolchain markers detected. Add dev-tools profiles for your formatter, linter, and language server.",
|
||||||
|
};
|
||||||
|
}
|
||||||
44
.pi/agent/extensions/dev-tools/src/runtime.test.ts
Normal file
44
.pi/agent/extensions/dev-tools/src/runtime.test.ts
Normal file
@@ -0,0 +1,44 @@
|
|||||||
|
import test from "node:test";
|
||||||
|
import assert from "node:assert/strict";
|
||||||
|
import { createDevToolsRuntime } from "./runtime.ts";
|
||||||
|
|
||||||
|
test("refreshDiagnostics falls back to command diagnostics when LSP is unavailable", async () => {
|
||||||
|
const runtime = createDevToolsRuntime({
|
||||||
|
cwd: "/repo",
|
||||||
|
agentDir: "/agent",
|
||||||
|
loadConfig: () => ({
|
||||||
|
defaults: { maxDiagnosticsPerFile: 5 },
|
||||||
|
profiles: [
|
||||||
|
{
|
||||||
|
name: "typescript",
|
||||||
|
match: ["src/**/*.ts"],
|
||||||
|
languageId: "typescript",
|
||||||
|
workspaceRootMarkers: ["package.json"],
|
||||||
|
formatter: { kind: "command", command: ["biome", "format", "--write", "{file}"] },
|
||||||
|
diagnostics: [
|
||||||
|
{ kind: "lsp", command: ["typescript-language-server", "--stdio"] },
|
||||||
|
{ kind: "command", parser: "eslint-json", command: ["eslint", "--format", "json", "{file}"] },
|
||||||
|
],
|
||||||
|
},
|
||||||
|
],
|
||||||
|
}),
|
||||||
|
knownPaths: ["/repo/package.json"],
|
||||||
|
formatterRunner: { formatFile: async () => ({ status: "skipped" }) },
|
||||||
|
lspBackend: {
|
||||||
|
collectForFile: async () => ({ status: "unavailable", items: [], message: "spawn ENOENT" }),
|
||||||
|
},
|
||||||
|
commandBackend: {
|
||||||
|
collect: async () => ({
|
||||||
|
status: "ok",
|
||||||
|
items: [{ severity: "error", message: "Unexpected console statement.", line: 2, column: 3, source: "eslint" }],
|
||||||
|
}),
|
||||||
|
},
|
||||||
|
probeProject: async () => ({ ecosystem: "typescript", summary: "TypeScript project detected." }),
|
||||||
|
});
|
||||||
|
|
||||||
|
await runtime.refreshDiagnosticsForPath("/repo/src/app.ts", "console.log('x')\n");
|
||||||
|
|
||||||
|
const promptBlock = runtime.getPromptBlock() ?? "";
|
||||||
|
assert.match(promptBlock, /Unexpected console statement/);
|
||||||
|
assert.match(promptBlock, /spawn ENOENT/);
|
||||||
|
});
|
||||||
134
.pi/agent/extensions/dev-tools/src/runtime.ts
Normal file
134
.pi/agent/extensions/dev-tools/src/runtime.ts
Normal file
@@ -0,0 +1,134 @@
|
|||||||
|
import { readFile } from "node:fs/promises";
|
||||||
|
import { loadDevToolsConfig } from "./config.ts";
|
||||||
|
import type { CapabilityGap, DiagnosticsState } from "./diagnostics/types.ts";
|
||||||
|
import { probeProject } from "./project-probe.ts";
|
||||||
|
import { resolveProfileForPath } from "./profiles.ts";
|
||||||
|
import { buildPromptBlock } from "./summary.ts";
|
||||||
|
|
||||||
|
export interface FormatResult {
|
||||||
|
status: "formatted" | "skipped" | "failed";
|
||||||
|
message?: string;
|
||||||
|
}
|
||||||
|
|
||||||
|
export interface DevToolsRuntime {
|
||||||
|
formatAfterMutation(absolutePath: string): Promise<FormatResult>;
|
||||||
|
noteMutation(absolutePath: string, formatResult: FormatResult): void;
|
||||||
|
setDiagnostics(path: string, state: DiagnosticsState): void;
|
||||||
|
recordCapabilityGap(path: string, message: string): void;
|
||||||
|
getPromptBlock(): string | undefined;
|
||||||
|
refreshDiagnosticsForPath(absolutePath: string, text?: string): Promise<void>;
|
||||||
|
}
|
||||||
|
|
||||||
|
type LoadedConfig = ReturnType<typeof loadDevToolsConfig>;
|
||||||
|
|
||||||
|
export function createDevToolsRuntime(deps: {
|
||||||
|
cwd: string;
|
||||||
|
agentDir: string;
|
||||||
|
loadConfig?: () => LoadedConfig;
|
||||||
|
knownPaths?: string[];
|
||||||
|
formatterRunner: { formatFile: (input: any) => Promise<FormatResult> };
|
||||||
|
lspBackend: { collectForFile: (input: any) => Promise<DiagnosticsState> };
|
||||||
|
commandBackend: { collect: (input: any) => Promise<DiagnosticsState> };
|
||||||
|
probeProject?: typeof probeProject;
|
||||||
|
}): DevToolsRuntime {
|
||||||
|
const diagnosticsByFile = new Map<string, DiagnosticsState>();
|
||||||
|
const capabilityGaps: CapabilityGap[] = [];
|
||||||
|
const getConfig = () => deps.loadConfig?.() ?? loadDevToolsConfig(deps.cwd, deps.agentDir);
|
||||||
|
|
||||||
|
function setDiagnostics(path: string, state: DiagnosticsState) {
|
||||||
|
diagnosticsByFile.set(path, state);
|
||||||
|
}
|
||||||
|
|
||||||
|
function recordCapabilityGap(path: string, message: string) {
|
||||||
|
if (!capabilityGaps.some((gap) => gap.path === path && gap.message === message)) {
|
||||||
|
capabilityGaps.push({ path, message });
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
function getPromptBlock() {
|
||||||
|
const config = getConfig();
|
||||||
|
const maxDiagnosticsPerFile = config?.defaults?.maxDiagnosticsPerFile ?? 10;
|
||||||
|
if (diagnosticsByFile.size === 0 && capabilityGaps.length === 0) return undefined;
|
||||||
|
return buildPromptBlock({
|
||||||
|
maxDiagnosticsPerFile,
|
||||||
|
diagnosticsByFile,
|
||||||
|
capabilityGaps,
|
||||||
|
});
|
||||||
|
}
|
||||||
|
|
||||||
|
async function refreshDiagnosticsForPath(absolutePath: string, text?: string) {
|
||||||
|
const config = getConfig();
|
||||||
|
if (!config) {
|
||||||
|
recordCapabilityGap(absolutePath, "No dev-tools config found.");
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
|
||||||
|
const match = resolveProfileForPath(config, absolutePath, deps.cwd, deps.knownPaths ?? []);
|
||||||
|
if (!match) {
|
||||||
|
const probe = await (deps.probeProject ?? probeProject)({ cwd: deps.cwd });
|
||||||
|
recordCapabilityGap(absolutePath, `No profile matched. ${probe.summary}`);
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
|
||||||
|
const fileText = text ?? await readFile(absolutePath, "utf8");
|
||||||
|
|
||||||
|
for (const backend of match.profile.diagnostics) {
|
||||||
|
if (backend.kind === "lsp") {
|
||||||
|
const lspResult = await deps.lspBackend.collectForFile({
|
||||||
|
key: `${match.profile.languageId ?? "plain"}:${match.workspaceRoot}`,
|
||||||
|
absolutePath,
|
||||||
|
workspaceRoot: match.workspaceRoot,
|
||||||
|
languageId: match.profile.languageId ?? "plaintext",
|
||||||
|
text: fileText,
|
||||||
|
command: backend.command,
|
||||||
|
});
|
||||||
|
|
||||||
|
if (lspResult.status === "ok") {
|
||||||
|
setDiagnostics(absolutePath, lspResult);
|
||||||
|
if (lspResult.message) recordCapabilityGap(absolutePath, lspResult.message);
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
|
||||||
|
recordCapabilityGap(absolutePath, lspResult.message ?? "LSP diagnostics unavailable.");
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
|
||||||
|
const commandResult = await deps.commandBackend.collect({
|
||||||
|
absolutePath,
|
||||||
|
workspaceRoot: match.workspaceRoot,
|
||||||
|
backend,
|
||||||
|
timeoutMs: config.defaults?.diagnosticTimeoutMs,
|
||||||
|
});
|
||||||
|
|
||||||
|
setDiagnostics(absolutePath, commandResult);
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
|
||||||
|
recordCapabilityGap(absolutePath, "No diagnostics backend succeeded.");
|
||||||
|
}
|
||||||
|
|
||||||
|
return {
|
||||||
|
async formatAfterMutation(absolutePath: string) {
|
||||||
|
const config = getConfig();
|
||||||
|
if (!config) return { status: "skipped" as const };
|
||||||
|
const match = resolveProfileForPath(config, absolutePath, deps.cwd, deps.knownPaths ?? []);
|
||||||
|
if (!match?.profile.formatter) return { status: "skipped" as const };
|
||||||
|
return deps.formatterRunner.formatFile({
|
||||||
|
absolutePath,
|
||||||
|
workspaceRoot: match.workspaceRoot,
|
||||||
|
formatter: match.profile.formatter,
|
||||||
|
timeoutMs: config.defaults?.formatTimeoutMs,
|
||||||
|
});
|
||||||
|
},
|
||||||
|
noteMutation(absolutePath: string, formatResult: FormatResult) {
|
||||||
|
if (formatResult.status === "failed") {
|
||||||
|
recordCapabilityGap(absolutePath, formatResult.message ?? "Formatter failed.");
|
||||||
|
}
|
||||||
|
void refreshDiagnosticsForPath(absolutePath);
|
||||||
|
},
|
||||||
|
setDiagnostics,
|
||||||
|
recordCapabilityGap,
|
||||||
|
getPromptBlock,
|
||||||
|
refreshDiagnosticsForPath,
|
||||||
|
};
|
||||||
|
}
|
||||||
40
.pi/agent/extensions/dev-tools/src/schema.ts
Normal file
40
.pi/agent/extensions/dev-tools/src/schema.ts
Normal file
@@ -0,0 +1,40 @@
|
|||||||
|
import { Type, type Static } from "@sinclair/typebox";
|
||||||
|
|
||||||
|
const CommandSchema = Type.Object({
|
||||||
|
kind: Type.Literal("command"),
|
||||||
|
command: Type.Array(Type.String({ minLength: 1 }), { minItems: 1 }),
|
||||||
|
});
|
||||||
|
|
||||||
|
const LspSchema = Type.Object({
|
||||||
|
kind: Type.Literal("lsp"),
|
||||||
|
command: Type.Array(Type.String({ minLength: 1 }), { minItems: 1 }),
|
||||||
|
});
|
||||||
|
|
||||||
|
const CommandDiagnosticsSchema = Type.Object({
|
||||||
|
kind: Type.Literal("command"),
|
||||||
|
parser: Type.String({ minLength: 1 }),
|
||||||
|
command: Type.Array(Type.String({ minLength: 1 }), { minItems: 1 }),
|
||||||
|
});
|
||||||
|
|
||||||
|
export const DevToolsProfileSchema = Type.Object({
|
||||||
|
name: Type.String({ minLength: 1 }),
|
||||||
|
match: Type.Array(Type.String({ minLength: 1 }), { minItems: 1 }),
|
||||||
|
languageId: Type.Optional(Type.String({ minLength: 1 })),
|
||||||
|
workspaceRootMarkers: Type.Array(Type.String({ minLength: 1 }), { minItems: 1 }),
|
||||||
|
formatter: Type.Optional(CommandSchema),
|
||||||
|
diagnostics: Type.Array(Type.Union([LspSchema, CommandDiagnosticsSchema])),
|
||||||
|
});
|
||||||
|
|
||||||
|
export const DevToolsConfigSchema = Type.Object({
|
||||||
|
defaults: Type.Optional(Type.Object({
|
||||||
|
formatTimeoutMs: Type.Optional(Type.Integer({ minimum: 1 })),
|
||||||
|
diagnosticTimeoutMs: Type.Optional(Type.Integer({ minimum: 1 })),
|
||||||
|
maxDiagnosticsPerFile: Type.Optional(Type.Integer({ minimum: 1 })),
|
||||||
|
})),
|
||||||
|
profiles: Type.Array(DevToolsProfileSchema, { minItems: 1 }),
|
||||||
|
});
|
||||||
|
|
||||||
|
export type DevToolsProfile = Static<typeof DevToolsProfileSchema>;
|
||||||
|
export type DevToolsConfig = Static<typeof DevToolsConfigSchema>;
|
||||||
|
export type FormatterConfig = NonNullable<DevToolsProfile["formatter"]>;
|
||||||
|
export type DiagnosticsConfig = DevToolsProfile["diagnostics"][number];
|
||||||
27
.pi/agent/extensions/dev-tools/src/summary.test.ts
Normal file
27
.pi/agent/extensions/dev-tools/src/summary.test.ts
Normal file
@@ -0,0 +1,27 @@
|
|||||||
|
import test from "node:test";
|
||||||
|
import assert from "node:assert/strict";
|
||||||
|
import { buildPromptBlock } from "./summary.ts";
|
||||||
|
|
||||||
|
test("buildPromptBlock caps diagnostics per file and includes capability gaps", () => {
|
||||||
|
const block = buildPromptBlock({
|
||||||
|
maxDiagnosticsPerFile: 1,
|
||||||
|
diagnosticsByFile: new Map([
|
||||||
|
[
|
||||||
|
"/repo/src/app.ts",
|
||||||
|
{
|
||||||
|
status: "ok",
|
||||||
|
items: [
|
||||||
|
{ severity: "error", message: "Unexpected console statement.", line: 2, column: 3, source: "eslint" },
|
||||||
|
{ severity: "warning", message: "Unused variable.", line: 4, column: 9, source: "eslint" },
|
||||||
|
],
|
||||||
|
},
|
||||||
|
],
|
||||||
|
]),
|
||||||
|
capabilityGaps: [{ path: "/repo/src/app.ts", message: "Configured executable `eslint` not found in PATH." }],
|
||||||
|
});
|
||||||
|
|
||||||
|
assert.match(block, /app.ts: 1 error, 1 warning/);
|
||||||
|
assert.match(block, /Unexpected console statement/);
|
||||||
|
assert.doesNotMatch(block, /Unused variable/);
|
||||||
|
assert.match(block, /not found in PATH/);
|
||||||
|
});
|
||||||
31
.pi/agent/extensions/dev-tools/src/summary.ts
Normal file
31
.pi/agent/extensions/dev-tools/src/summary.ts
Normal file
@@ -0,0 +1,31 @@
|
|||||||
|
import type { CapabilityGap, DiagnosticsState } from "./diagnostics/types.ts";
|
||||||
|
|
||||||
|
export function buildPromptBlock(input: {
|
||||||
|
maxDiagnosticsPerFile: number;
|
||||||
|
diagnosticsByFile: Map<string, DiagnosticsState>;
|
||||||
|
capabilityGaps: CapabilityGap[];
|
||||||
|
}) {
|
||||||
|
const lines = ["Current changed-file diagnostics:"];
|
||||||
|
|
||||||
|
for (const [path, state] of input.diagnosticsByFile) {
|
||||||
|
if (state.status === "unavailable") {
|
||||||
|
lines.push(`- ${path}: diagnostics unavailable (${state.message ?? "unknown error"})`);
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
|
||||||
|
const errors = state.items.filter((item) => item.severity === "error");
|
||||||
|
const warnings = state.items.filter((item) => item.severity === "warning");
|
||||||
|
lines.push(`- ${path}: ${errors.length} error${errors.length === 1 ? "" : "s"}, ${warnings.length} warning${warnings.length === 1 ? "" : "s"}`);
|
||||||
|
|
||||||
|
for (const item of state.items.slice(0, input.maxDiagnosticsPerFile)) {
|
||||||
|
const location = item.line ? `:${item.line}${item.column ? `:${item.column}` : ""}` : "";
|
||||||
|
lines.push(` - ${item.severity.toUpperCase()}${location} ${item.message}`);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
for (const gap of input.capabilityGaps) {
|
||||||
|
lines.push(`- setup gap for ${gap.path}: ${gap.message}`);
|
||||||
|
}
|
||||||
|
|
||||||
|
return lines.join("\n");
|
||||||
|
}
|
||||||
54
.pi/agent/extensions/dev-tools/src/tools/edit.test.ts
Normal file
54
.pi/agent/extensions/dev-tools/src/tools/edit.test.ts
Normal file
@@ -0,0 +1,54 @@
|
|||||||
|
import test from "node:test";
|
||||||
|
import assert from "node:assert/strict";
|
||||||
|
import { mkdtemp, readFile, writeFile } from "node:fs/promises";
|
||||||
|
import { join } from "node:path";
|
||||||
|
import { tmpdir } from "node:os";
|
||||||
|
import { createFormattedEditTool } from "./edit.ts";
|
||||||
|
|
||||||
|
test("edit applies the replacement and then formats the file", async () => {
|
||||||
|
const dir = await mkdtemp(join(tmpdir(), "dev-tools-edit-"));
|
||||||
|
const path = join(dir, "app.ts");
|
||||||
|
await writeFile(path, "const x=1\n", "utf8");
|
||||||
|
|
||||||
|
const tool = createFormattedEditTool(dir, {
|
||||||
|
formatAfterMutation: async (absolutePath) => {
|
||||||
|
await writeFile(absolutePath, "const x = 2;\n", "utf8");
|
||||||
|
return { status: "formatted" };
|
||||||
|
},
|
||||||
|
noteMutation() {},
|
||||||
|
});
|
||||||
|
|
||||||
|
await tool.execute(
|
||||||
|
"tool-1",
|
||||||
|
{ path, edits: [{ oldText: "const x=1", newText: "const x=2" }] },
|
||||||
|
undefined,
|
||||||
|
undefined,
|
||||||
|
undefined,
|
||||||
|
);
|
||||||
|
|
||||||
|
assert.equal(await readFile(path, "utf8"), "const x = 2;\n");
|
||||||
|
});
|
||||||
|
|
||||||
|
test("edit preserves the changed file when formatter fails", async () => {
|
||||||
|
const dir = await mkdtemp(join(tmpdir(), "dev-tools-edit-"));
|
||||||
|
const path = join(dir, "app.ts");
|
||||||
|
await writeFile(path, "const x=1\n", "utf8");
|
||||||
|
|
||||||
|
const tool = createFormattedEditTool(dir, {
|
||||||
|
formatAfterMutation: async () => ({ status: "failed", message: "formatter not found" }),
|
||||||
|
noteMutation() {},
|
||||||
|
});
|
||||||
|
|
||||||
|
await assert.rejects(
|
||||||
|
() => tool.execute(
|
||||||
|
"tool-1",
|
||||||
|
{ path, edits: [{ oldText: "const x=1", newText: "const x=2" }] },
|
||||||
|
undefined,
|
||||||
|
undefined,
|
||||||
|
undefined,
|
||||||
|
),
|
||||||
|
/Auto-format failed/,
|
||||||
|
);
|
||||||
|
|
||||||
|
assert.equal(await readFile(path, "utf8"), "const x=2\n");
|
||||||
|
});
|
||||||
21
.pi/agent/extensions/dev-tools/src/tools/edit.ts
Normal file
21
.pi/agent/extensions/dev-tools/src/tools/edit.ts
Normal file
@@ -0,0 +1,21 @@
|
|||||||
|
import { createEditToolDefinition } from "@mariozechner/pi-coding-agent";
|
||||||
|
import { constants } from "node:fs";
|
||||||
|
import { access, readFile, writeFile } from "node:fs/promises";
|
||||||
|
import type { DevToolsRuntime } from "../runtime.ts";
|
||||||
|
|
||||||
|
export function createFormattedEditTool(cwd: string, runtime: DevToolsRuntime) {
|
||||||
|
return createEditToolDefinition(cwd, {
|
||||||
|
operations: {
|
||||||
|
access: (absolutePath) => access(absolutePath, constants.R_OK | constants.W_OK),
|
||||||
|
readFile: (absolutePath) => readFile(absolutePath),
|
||||||
|
writeFile: async (absolutePath, content) => {
|
||||||
|
await writeFile(absolutePath, content, "utf8");
|
||||||
|
const formatResult = await runtime.formatAfterMutation(absolutePath);
|
||||||
|
runtime.noteMutation(absolutePath, formatResult);
|
||||||
|
if (formatResult.status === "failed") {
|
||||||
|
throw new Error(`Auto-format failed for ${absolutePath}: ${formatResult.message}`);
|
||||||
|
}
|
||||||
|
},
|
||||||
|
},
|
||||||
|
});
|
||||||
|
}
|
||||||
@@ -0,0 +1,15 @@
|
|||||||
|
import test from "node:test";
|
||||||
|
import assert from "node:assert/strict";
|
||||||
|
import { createSetupSuggestTool } from "./setup-suggest.ts";
|
||||||
|
|
||||||
|
test("dev_tools_suggest_setup returns a concrete recommendation string", async () => {
|
||||||
|
const tool = createSetupSuggestTool({
|
||||||
|
suggestSetup: async () => "TypeScript project detected. Recommended: bunx biome init and npm i -D typescript-language-server.",
|
||||||
|
});
|
||||||
|
|
||||||
|
const result = await tool.execute("tool-1", {}, undefined, undefined, undefined);
|
||||||
|
const text = result.content[0]?.type === "text" ? result.content[0].text : "";
|
||||||
|
|
||||||
|
assert.match(text, /TypeScript project detected/);
|
||||||
|
assert.match(text, /biome/);
|
||||||
|
});
|
||||||
17
.pi/agent/extensions/dev-tools/src/tools/setup-suggest.ts
Normal file
17
.pi/agent/extensions/dev-tools/src/tools/setup-suggest.ts
Normal file
@@ -0,0 +1,17 @@
|
|||||||
|
import { Type } from "@sinclair/typebox";
|
||||||
|
|
||||||
|
export function createSetupSuggestTool(deps: { suggestSetup: () => Promise<string> }) {
|
||||||
|
return {
|
||||||
|
name: "dev_tools_suggest_setup",
|
||||||
|
label: "Dev Tools Suggest Setup",
|
||||||
|
description: "Suggest formatter/linter/LSP setup for the current project.",
|
||||||
|
parameters: Type.Object({}),
|
||||||
|
async execute() {
|
||||||
|
const text = await deps.suggestSetup();
|
||||||
|
return {
|
||||||
|
content: [{ type: "text" as const, text }],
|
||||||
|
details: { suggestion: text },
|
||||||
|
};
|
||||||
|
},
|
||||||
|
};
|
||||||
|
}
|
||||||
40
.pi/agent/extensions/dev-tools/src/tools/write.test.ts
Normal file
40
.pi/agent/extensions/dev-tools/src/tools/write.test.ts
Normal file
@@ -0,0 +1,40 @@
|
|||||||
|
import test from "node:test";
|
||||||
|
import assert from "node:assert/strict";
|
||||||
|
import { mkdtemp, readFile, writeFile } from "node:fs/promises";
|
||||||
|
import { join } from "node:path";
|
||||||
|
import { tmpdir } from "node:os";
|
||||||
|
import { createFormattedWriteTool } from "./write.ts";
|
||||||
|
|
||||||
|
test("write keeps the file when auto-format fails", async () => {
|
||||||
|
const dir = await mkdtemp(join(tmpdir(), "dev-tools-write-"));
|
||||||
|
const path = join(dir, "app.ts");
|
||||||
|
|
||||||
|
const tool = createFormattedWriteTool(dir, {
|
||||||
|
formatAfterMutation: async () => ({ status: "failed", message: "biome missing" }),
|
||||||
|
noteMutation() {},
|
||||||
|
});
|
||||||
|
|
||||||
|
await assert.rejects(
|
||||||
|
() => tool.execute("tool-1", { path, content: "const x=1\n" }, undefined, undefined, undefined),
|
||||||
|
/Auto-format failed/,
|
||||||
|
);
|
||||||
|
|
||||||
|
assert.equal(await readFile(path, "utf8"), "const x=1\n");
|
||||||
|
});
|
||||||
|
|
||||||
|
test("write calls formatting immediately after writing", async () => {
|
||||||
|
const dir = await mkdtemp(join(tmpdir(), "dev-tools-write-"));
|
||||||
|
const path = join(dir, "app.ts");
|
||||||
|
|
||||||
|
const tool = createFormattedWriteTool(dir, {
|
||||||
|
formatAfterMutation: async (absolutePath) => {
|
||||||
|
await writeFile(absolutePath, "const x = 1;\n", "utf8");
|
||||||
|
return { status: "formatted" };
|
||||||
|
},
|
||||||
|
noteMutation() {},
|
||||||
|
});
|
||||||
|
|
||||||
|
await tool.execute("tool-1", { path, content: "const x=1\n" }, undefined, undefined, undefined);
|
||||||
|
|
||||||
|
assert.equal(await readFile(path, "utf8"), "const x = 1;\n");
|
||||||
|
});
|
||||||
23
.pi/agent/extensions/dev-tools/src/tools/write.ts
Normal file
23
.pi/agent/extensions/dev-tools/src/tools/write.ts
Normal file
@@ -0,0 +1,23 @@
|
|||||||
|
import { createWriteToolDefinition } from "@mariozechner/pi-coding-agent";
|
||||||
|
import { mkdir, writeFile } from "node:fs/promises";
|
||||||
|
import { dirname } from "node:path";
|
||||||
|
import type { DevToolsRuntime } from "../runtime.ts";
|
||||||
|
|
||||||
|
export function createFormattedWriteTool(cwd: string, runtime: DevToolsRuntime) {
|
||||||
|
const original = createWriteToolDefinition(cwd, {
|
||||||
|
operations: {
|
||||||
|
mkdir: (dir) => mkdir(dir, { recursive: true }).then(() => {}),
|
||||||
|
writeFile: async (absolutePath, content) => {
|
||||||
|
await mkdir(dirname(absolutePath), { recursive: true });
|
||||||
|
await writeFile(absolutePath, content, "utf8");
|
||||||
|
const formatResult = await runtime.formatAfterMutation(absolutePath);
|
||||||
|
runtime.noteMutation(absolutePath, formatResult);
|
||||||
|
if (formatResult.status === "failed") {
|
||||||
|
throw new Error(`Auto-format failed for ${absolutePath}: ${formatResult.message}`);
|
||||||
|
}
|
||||||
|
},
|
||||||
|
},
|
||||||
|
});
|
||||||
|
|
||||||
|
return original;
|
||||||
|
}
|
||||||
99
.pi/agent/extensions/tmux-subagent/index.ts
Normal file
99
.pi/agent/extensions/tmux-subagent/index.ts
Normal file
@@ -0,0 +1,99 @@
|
|||||||
|
import type { ExtensionAPI } from "@mariozechner/pi-coding-agent";
|
||||||
|
import { dirname, join } from "node:path";
|
||||||
|
import { fileURLToPath } from "node:url";
|
||||||
|
import { createRunArtifacts } from "./src/artifacts.ts";
|
||||||
|
import { monitorRun } from "./src/monitor.ts";
|
||||||
|
import { listAvailableModelReferences } from "./src/models.ts";
|
||||||
|
import { createTmuxSingleRunner } from "./src/runner.ts";
|
||||||
|
import {
|
||||||
|
buildCurrentWindowArgs,
|
||||||
|
buildKillPaneArgs,
|
||||||
|
buildSplitWindowArgs,
|
||||||
|
buildWrapperShellCommand,
|
||||||
|
isInsideTmux,
|
||||||
|
} from "./src/tmux.ts";
|
||||||
|
import { createSubagentParamsSchema } from "./src/schema.ts";
|
||||||
|
import { createSubagentTool } from "./src/tool.ts";
|
||||||
|
|
||||||
|
const packageRoot = dirname(fileURLToPath(import.meta.url));
|
||||||
|
const wrapperPath = join(packageRoot, "src", "wrapper", "cli.mjs");
|
||||||
|
|
||||||
|
export default function tmuxSubagentExtension(pi: ExtensionAPI) {
|
||||||
|
if (process.env.PI_TMUX_SUBAGENT_GITHUB_COPILOT_INITIATOR === "agent") {
|
||||||
|
pi.registerProvider("github-copilot", {
|
||||||
|
headers: { "X-Initiator": "agent" },
|
||||||
|
});
|
||||||
|
}
|
||||||
|
|
||||||
|
// In wrapper/child sessions spawned by the tmux runner we must not register the
|
||||||
|
// subagent tool (that would cause nested subagent registrations). Skip all
|
||||||
|
// subagent-tool registration logic when PI_TMUX_SUBAGENT_CHILD is set. Provider
|
||||||
|
// overrides (above) are still allowed in child runs, so the guard is placed
|
||||||
|
// after provider registration.
|
||||||
|
if (process.env.PI_TMUX_SUBAGENT_CHILD === "1") {
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
|
||||||
|
let lastRegisteredModelsKey: string | undefined;
|
||||||
|
|
||||||
|
const runSingleTask = createTmuxSingleRunner({
|
||||||
|
assertInsideTmux() {
|
||||||
|
if (!isInsideTmux()) throw new Error("tmux-backed subagents require pi to be running inside tmux.");
|
||||||
|
},
|
||||||
|
async getCurrentWindowId() {
|
||||||
|
const result = await pi.exec("tmux", buildCurrentWindowArgs());
|
||||||
|
return result.stdout.trim();
|
||||||
|
},
|
||||||
|
createArtifacts: createRunArtifacts,
|
||||||
|
buildWrapperCommand(metaPath: string) {
|
||||||
|
return buildWrapperShellCommand({ nodePath: process.execPath, wrapperPath, metaPath });
|
||||||
|
},
|
||||||
|
async createPane(input) {
|
||||||
|
const result = await pi.exec("tmux", buildSplitWindowArgs(input));
|
||||||
|
return result.stdout.trim();
|
||||||
|
},
|
||||||
|
monitorRun,
|
||||||
|
async killPane(paneId: string) {
|
||||||
|
await pi.exec("tmux", buildKillPaneArgs(paneId));
|
||||||
|
},
|
||||||
|
});
|
||||||
|
|
||||||
|
const registerSubagentTool = (availableModels: string[]) => {
|
||||||
|
// Do not register a tool when no models are available. Remember that the
|
||||||
|
// last-registered key is different from the empty sentinel so that a later
|
||||||
|
// non-empty list will still trigger registration.
|
||||||
|
if (!availableModels || availableModels.length === 0) {
|
||||||
|
const emptyKey = "\u0000";
|
||||||
|
if (lastRegisteredModelsKey === emptyKey) return;
|
||||||
|
lastRegisteredModelsKey = emptyKey;
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
|
||||||
|
// Create a deduplication key that is independent of the order of
|
||||||
|
// availableModels by sorting a lowercase copy. Do not mutate
|
||||||
|
// availableModels itself since we want to preserve the original order for
|
||||||
|
// schema enum values.
|
||||||
|
const key = [...availableModels].map((s) => s.toLowerCase()).sort().join("\u0000");
|
||||||
|
if (key === lastRegisteredModelsKey) return;
|
||||||
|
|
||||||
|
lastRegisteredModelsKey = key;
|
||||||
|
pi.registerTool(
|
||||||
|
createSubagentTool({
|
||||||
|
parameters: createSubagentParamsSchema(availableModels),
|
||||||
|
runSingleTask,
|
||||||
|
}),
|
||||||
|
);
|
||||||
|
};
|
||||||
|
|
||||||
|
const syncSubagentTool = (ctx: { modelRegistry: { getAvailable(): Array<{ provider: string; id: string }> } }) => {
|
||||||
|
registerSubagentTool(listAvailableModelReferences(ctx.modelRegistry));
|
||||||
|
};
|
||||||
|
|
||||||
|
pi.on("session_start", (_event, ctx) => {
|
||||||
|
syncSubagentTool(ctx);
|
||||||
|
});
|
||||||
|
|
||||||
|
pi.on("before_agent_start", (_event, ctx) => {
|
||||||
|
syncSubagentTool(ctx);
|
||||||
|
});
|
||||||
|
}
|
||||||
4365
.pi/agent/extensions/tmux-subagent/package-lock.json
generated
Normal file
4365
.pi/agent/extensions/tmux-subagent/package-lock.json
generated
Normal file
File diff suppressed because it is too large
Load Diff
23
.pi/agent/extensions/tmux-subagent/package.json
Normal file
23
.pi/agent/extensions/tmux-subagent/package.json
Normal file
@@ -0,0 +1,23 @@
|
|||||||
|
{
|
||||||
|
"name": "pi-tmux-subagent-extension",
|
||||||
|
"private": true,
|
||||||
|
"version": "0.0.0",
|
||||||
|
"type": "module",
|
||||||
|
"scripts": {
|
||||||
|
"test": "tsx --test src/*.test.ts src/**/*.test.ts"
|
||||||
|
},
|
||||||
|
"pi": {
|
||||||
|
"extensions": ["./index.ts"],
|
||||||
|
"prompts": ["./prompts/*.md"]
|
||||||
|
},
|
||||||
|
"devDependencies": {
|
||||||
|
"@mariozechner/pi-agent-core": "^0.66.1",
|
||||||
|
"@mariozechner/pi-ai": "^0.66.1",
|
||||||
|
"@mariozechner/pi-coding-agent": "^0.66.1",
|
||||||
|
"@mariozechner/pi-tui": "^0.66.1",
|
||||||
|
"@sinclair/typebox": "^0.34.49",
|
||||||
|
"@types/node": "^25.5.2",
|
||||||
|
"tsx": "^4.21.0",
|
||||||
|
"typescript": "^6.0.2"
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,10 @@
|
|||||||
|
---
|
||||||
|
description: Implement, review, then revise using tmux-backed subagents
|
||||||
|
---
|
||||||
|
Use the `subagent` tool in chain mode:
|
||||||
|
|
||||||
|
1. `worker` to implement: $@
|
||||||
|
2. `reviewer` to review `{previous}` and identify issues
|
||||||
|
3. `worker` to revise the implementation using `{previous}`
|
||||||
|
|
||||||
|
User request: $@
|
||||||
10
.pi/agent/extensions/tmux-subagent/prompts/implement.md
Normal file
10
.pi/agent/extensions/tmux-subagent/prompts/implement.md
Normal file
@@ -0,0 +1,10 @@
|
|||||||
|
---
|
||||||
|
description: Scout, plan, and implement using tmux-backed subagents
|
||||||
|
---
|
||||||
|
Use the `subagent` tool to handle this request in three stages:
|
||||||
|
|
||||||
|
1. Run `scout` to inspect the codebase for: $@
|
||||||
|
2. Run `planner` in chain mode, using `{previous}` from the scout output to produce a concrete implementation plan
|
||||||
|
3. Run `worker` in chain mode, using `{previous}` from the planner output to implement the approved plan
|
||||||
|
|
||||||
|
User request: $@
|
||||||
@@ -0,0 +1,9 @@
|
|||||||
|
---
|
||||||
|
description: Scout the codebase, then produce a plan using tmux-backed subagents
|
||||||
|
---
|
||||||
|
Use the `subagent` tool in chain mode:
|
||||||
|
|
||||||
|
1. `scout` to inspect the codebase for: $@
|
||||||
|
2. `planner` to turn `{previous}` into an implementation plan
|
||||||
|
|
||||||
|
User request: $@
|
||||||
54
.pi/agent/extensions/tmux-subagent/src/agents.test.ts
Normal file
54
.pi/agent/extensions/tmux-subagent/src/agents.test.ts
Normal file
@@ -0,0 +1,54 @@
|
|||||||
|
import test from "node:test";
|
||||||
|
import assert from "node:assert/strict";
|
||||||
|
import { mkdir, writeFile, mkdtemp } from "node:fs/promises";
|
||||||
|
import { join } from "node:path";
|
||||||
|
import { tmpdir } from "node:os";
|
||||||
|
import { BUILTIN_AGENTS } from "./builtin-agents.ts";
|
||||||
|
import { discoverAgents } from "./agents.ts";
|
||||||
|
|
||||||
|
test("discoverAgents returns built-ins and lets user markdown override by name", async () => {
|
||||||
|
const root = await mkdtemp(join(tmpdir(), "tmux-subagent-agents-"));
|
||||||
|
const agentDir = join(root, "agent-home");
|
||||||
|
const userAgentsDir = join(agentDir, "agents");
|
||||||
|
await mkdir(userAgentsDir, { recursive: true });
|
||||||
|
await writeFile(
|
||||||
|
join(userAgentsDir, "scout.md"),
|
||||||
|
`---\nname: scout\ndescription: User scout\nmodel: openai/gpt-5\n---\nUser override prompt`,
|
||||||
|
"utf8",
|
||||||
|
);
|
||||||
|
|
||||||
|
const result = discoverAgents(join(root, "repo"), {
|
||||||
|
scope: "user",
|
||||||
|
agentDir,
|
||||||
|
builtins: BUILTIN_AGENTS,
|
||||||
|
});
|
||||||
|
|
||||||
|
const scout = result.agents.find((agent) => agent.name === "scout");
|
||||||
|
assert.equal(scout?.source, "user");
|
||||||
|
assert.equal(scout?.description, "User scout");
|
||||||
|
assert.equal(scout?.model, "openai/gpt-5");
|
||||||
|
});
|
||||||
|
|
||||||
|
test("discoverAgents lets project agents override user agents when scope is both", async () => {
|
||||||
|
const root = await mkdtemp(join(tmpdir(), "tmux-subagent-agents-"));
|
||||||
|
const repo = join(root, "repo");
|
||||||
|
const agentDir = join(root, "agent-home");
|
||||||
|
const userAgentsDir = join(agentDir, "agents");
|
||||||
|
const projectAgentsDir = join(repo, ".pi", "agents");
|
||||||
|
await mkdir(userAgentsDir, { recursive: true });
|
||||||
|
await mkdir(projectAgentsDir, { recursive: true });
|
||||||
|
|
||||||
|
await writeFile(join(userAgentsDir, "worker.md"), `---\nname: worker\ndescription: User worker\n---\nUser worker`, "utf8");
|
||||||
|
await writeFile(join(projectAgentsDir, "worker.md"), `---\nname: worker\ndescription: Project worker\n---\nProject worker`, "utf8");
|
||||||
|
|
||||||
|
const result = discoverAgents(repo, {
|
||||||
|
scope: "both",
|
||||||
|
agentDir,
|
||||||
|
builtins: BUILTIN_AGENTS,
|
||||||
|
});
|
||||||
|
|
||||||
|
const worker = result.agents.find((agent) => agent.name === "worker");
|
||||||
|
assert.equal(worker?.source, "project");
|
||||||
|
assert.equal(worker?.description, "Project worker");
|
||||||
|
assert.equal(result.projectAgentsDir, projectAgentsDir);
|
||||||
|
});
|
||||||
91
.pi/agent/extensions/tmux-subagent/src/agents.ts
Normal file
91
.pi/agent/extensions/tmux-subagent/src/agents.ts
Normal file
@@ -0,0 +1,91 @@
|
|||||||
|
import { existsSync, readdirSync, readFileSync, statSync, type Dirent } from "node:fs";
|
||||||
|
import { dirname, join } from "node:path";
|
||||||
|
import { getAgentDir, parseFrontmatter } from "@mariozechner/pi-coding-agent";
|
||||||
|
import { BUILTIN_AGENTS, type AgentDefinition } from "./builtin-agents.ts";
|
||||||
|
import type { AgentScope } from "./schema.ts";
|
||||||
|
|
||||||
|
export interface AgentDiscoveryOptions {
|
||||||
|
scope?: AgentScope;
|
||||||
|
agentDir?: string;
|
||||||
|
builtins?: AgentDefinition[];
|
||||||
|
}
|
||||||
|
|
||||||
|
export interface AgentDiscoveryResult {
|
||||||
|
agents: AgentDefinition[];
|
||||||
|
projectAgentsDir: string | null;
|
||||||
|
}
|
||||||
|
|
||||||
|
function loadMarkdownAgents(dir: string, source: "user" | "project"): AgentDefinition[] {
|
||||||
|
if (!existsSync(dir)) return [];
|
||||||
|
|
||||||
|
let entries: Dirent[];
|
||||||
|
try {
|
||||||
|
entries = readdirSync(dir, { withFileTypes: true });
|
||||||
|
} catch {
|
||||||
|
return [];
|
||||||
|
}
|
||||||
|
|
||||||
|
const agents: AgentDefinition[] = [];
|
||||||
|
for (const entry of entries) {
|
||||||
|
if (!entry.name.endsWith(".md")) continue;
|
||||||
|
if (!entry.isFile() && !entry.isSymbolicLink()) continue;
|
||||||
|
|
||||||
|
const filePath = join(dir, entry.name);
|
||||||
|
const content = readFileSync(filePath, "utf8");
|
||||||
|
const { frontmatter, body } = parseFrontmatter<Record<string, string>>(content);
|
||||||
|
if (!frontmatter.name || !frontmatter.description) continue;
|
||||||
|
|
||||||
|
agents.push({
|
||||||
|
name: frontmatter.name,
|
||||||
|
description: frontmatter.description,
|
||||||
|
tools: frontmatter.tools?.split(",").map((tool) => tool.trim()).filter(Boolean),
|
||||||
|
model: frontmatter.model,
|
||||||
|
systemPrompt: body.trim(),
|
||||||
|
source,
|
||||||
|
filePath,
|
||||||
|
});
|
||||||
|
}
|
||||||
|
|
||||||
|
return agents;
|
||||||
|
}
|
||||||
|
|
||||||
|
function findNearestProjectAgentsDir(cwd: string): string | null {
|
||||||
|
let current = cwd;
|
||||||
|
while (true) {
|
||||||
|
const candidate = join(current, ".pi", "agents");
|
||||||
|
try {
|
||||||
|
if (statSync(candidate).isDirectory()) return candidate;
|
||||||
|
} catch {}
|
||||||
|
|
||||||
|
const parent = dirname(current);
|
||||||
|
if (parent === current) return null;
|
||||||
|
current = parent;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
export function discoverAgents(cwd: string, options: AgentDiscoveryOptions = {}): AgentDiscoveryResult {
|
||||||
|
const scope = options.scope ?? "user";
|
||||||
|
const builtins = options.builtins ?? BUILTIN_AGENTS;
|
||||||
|
const userAgentDir = join(options.agentDir ?? getAgentDir(), "agents");
|
||||||
|
const projectAgentsDir = findNearestProjectAgentsDir(cwd);
|
||||||
|
|
||||||
|
const sources = new Map<string, AgentDefinition>();
|
||||||
|
for (const agent of builtins) sources.set(agent.name, agent);
|
||||||
|
|
||||||
|
if (scope !== "project") {
|
||||||
|
for (const agent of loadMarkdownAgents(userAgentDir, "user")) {
|
||||||
|
sources.set(agent.name, agent);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
if (scope !== "user" && projectAgentsDir) {
|
||||||
|
for (const agent of loadMarkdownAgents(projectAgentsDir, "project")) {
|
||||||
|
sources.set(agent.name, agent);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
return {
|
||||||
|
agents: [...sources.values()],
|
||||||
|
projectAgentsDir,
|
||||||
|
};
|
||||||
|
}
|
||||||
21
.pi/agent/extensions/tmux-subagent/src/artifacts.test.ts
Normal file
21
.pi/agent/extensions/tmux-subagent/src/artifacts.test.ts
Normal file
@@ -0,0 +1,21 @@
|
|||||||
|
import test from "node:test";
|
||||||
|
import assert from "node:assert/strict";
|
||||||
|
import { mkdtemp, readFile } from "node:fs/promises";
|
||||||
|
import { join } from "node:path";
|
||||||
|
import { tmpdir } from "node:os";
|
||||||
|
import { createRunArtifacts } from "./artifacts.ts";
|
||||||
|
|
||||||
|
test("createRunArtifacts writes metadata and reserves stable artifact paths", async () => {
|
||||||
|
const cwd = await mkdtemp(join(tmpdir(), "tmux-subagent-run-"));
|
||||||
|
|
||||||
|
const artifacts = await createRunArtifacts(cwd, {
|
||||||
|
runId: "run-1",
|
||||||
|
task: "inspect auth",
|
||||||
|
systemPrompt: "You are scout",
|
||||||
|
});
|
||||||
|
|
||||||
|
assert.equal(artifacts.runId, "run-1");
|
||||||
|
assert.match(artifacts.dir, /\.pi\/subagents\/runs\/run-1$/);
|
||||||
|
assert.equal(JSON.parse(await readFile(artifacts.metaPath, "utf8")).task, "inspect auth");
|
||||||
|
assert.equal(await readFile(artifacts.systemPromptPath, "utf8"), "You are scout");
|
||||||
|
});
|
||||||
66
.pi/agent/extensions/tmux-subagent/src/artifacts.ts
Normal file
66
.pi/agent/extensions/tmux-subagent/src/artifacts.ts
Normal file
@@ -0,0 +1,66 @@
|
|||||||
|
import { mkdir, writeFile } from "node:fs/promises";
|
||||||
|
import { randomUUID } from "node:crypto";
|
||||||
|
import { join, resolve } from "node:path";
|
||||||
|
|
||||||
|
export interface RunArtifacts {
|
||||||
|
runId: string;
|
||||||
|
dir: string;
|
||||||
|
metaPath: string;
|
||||||
|
eventsPath: string;
|
||||||
|
resultPath: string;
|
||||||
|
stdoutPath: string;
|
||||||
|
stderrPath: string;
|
||||||
|
transcriptPath: string;
|
||||||
|
sessionPath: string;
|
||||||
|
systemPromptPath: string;
|
||||||
|
}
|
||||||
|
|
||||||
|
export async function createRunArtifacts(
|
||||||
|
cwd: string,
|
||||||
|
meta: Record<string, unknown> & { runId?: string; systemPrompt?: string },
|
||||||
|
): Promise<RunArtifacts> {
|
||||||
|
const runId = meta.runId ?? randomUUID();
|
||||||
|
const dir = resolve(cwd, ".pi", "subagents", "runs", runId);
|
||||||
|
await mkdir(dir, { recursive: true });
|
||||||
|
|
||||||
|
const artifacts: RunArtifacts = {
|
||||||
|
runId,
|
||||||
|
dir,
|
||||||
|
metaPath: join(dir, "meta.json"),
|
||||||
|
eventsPath: join(dir, "events.jsonl"),
|
||||||
|
resultPath: join(dir, "result.json"),
|
||||||
|
stdoutPath: join(dir, "stdout.log"),
|
||||||
|
stderrPath: join(dir, "stderr.log"),
|
||||||
|
transcriptPath: join(dir, "transcript.log"),
|
||||||
|
sessionPath: join(dir, "child-session.jsonl"),
|
||||||
|
systemPromptPath: join(dir, "system-prompt.md"),
|
||||||
|
};
|
||||||
|
|
||||||
|
await writeFile(
|
||||||
|
artifacts.metaPath,
|
||||||
|
JSON.stringify(
|
||||||
|
{
|
||||||
|
...meta,
|
||||||
|
runId,
|
||||||
|
sessionPath: artifacts.sessionPath,
|
||||||
|
eventsPath: artifacts.eventsPath,
|
||||||
|
resultPath: artifacts.resultPath,
|
||||||
|
stdoutPath: artifacts.stdoutPath,
|
||||||
|
stderrPath: artifacts.stderrPath,
|
||||||
|
transcriptPath: artifacts.transcriptPath,
|
||||||
|
systemPromptPath: artifacts.systemPromptPath,
|
||||||
|
},
|
||||||
|
null,
|
||||||
|
2,
|
||||||
|
),
|
||||||
|
"utf8",
|
||||||
|
);
|
||||||
|
|
||||||
|
await writeFile(artifacts.systemPromptPath, typeof meta.systemPrompt === "string" ? meta.systemPrompt : "", "utf8");
|
||||||
|
await writeFile(artifacts.eventsPath, "", "utf8");
|
||||||
|
await writeFile(artifacts.stdoutPath, "", "utf8");
|
||||||
|
await writeFile(artifacts.stderrPath, "", "utf8");
|
||||||
|
await writeFile(artifacts.transcriptPath, "", "utf8");
|
||||||
|
|
||||||
|
return artifacts;
|
||||||
|
}
|
||||||
43
.pi/agent/extensions/tmux-subagent/src/builtin-agents.ts
Normal file
43
.pi/agent/extensions/tmux-subagent/src/builtin-agents.ts
Normal file
@@ -0,0 +1,43 @@
|
|||||||
|
export interface AgentDefinition {
|
||||||
|
name: string;
|
||||||
|
description: string;
|
||||||
|
tools?: string[];
|
||||||
|
model?: string;
|
||||||
|
systemPrompt: string;
|
||||||
|
source: "builtin" | "user" | "project";
|
||||||
|
filePath?: string;
|
||||||
|
}
|
||||||
|
|
||||||
|
export const BUILTIN_AGENTS: AgentDefinition[] = [
|
||||||
|
{
|
||||||
|
name: "scout",
|
||||||
|
description: "Fast codebase recon and compressed context gathering",
|
||||||
|
tools: ["read", "grep", "find", "ls", "bash"],
|
||||||
|
model: "claude-haiku-4-5",
|
||||||
|
systemPrompt: "You are a scout. Explore quickly, summarize clearly, and avoid implementation.",
|
||||||
|
source: "builtin",
|
||||||
|
},
|
||||||
|
{
|
||||||
|
name: "planner",
|
||||||
|
description: "Turns exploration into implementation plans",
|
||||||
|
tools: ["read", "grep", "find", "ls"],
|
||||||
|
model: "claude-sonnet-4-5",
|
||||||
|
systemPrompt: "You are a planner. Produce implementation plans, file lists, and risks.",
|
||||||
|
source: "builtin",
|
||||||
|
},
|
||||||
|
{
|
||||||
|
name: "reviewer",
|
||||||
|
description: "Reviews code and identifies correctness and quality issues",
|
||||||
|
tools: ["read", "grep", "find", "ls", "bash"],
|
||||||
|
model: "claude-sonnet-4-5",
|
||||||
|
systemPrompt: "You are a reviewer. Inspect code critically and report concrete issues.",
|
||||||
|
source: "builtin",
|
||||||
|
},
|
||||||
|
{
|
||||||
|
name: "worker",
|
||||||
|
description: "General-purpose implementation agent",
|
||||||
|
model: "claude-sonnet-4-5",
|
||||||
|
systemPrompt: "You are a worker. Execute the delegated task completely and report final results clearly.",
|
||||||
|
source: "builtin",
|
||||||
|
},
|
||||||
|
];
|
||||||
397
.pi/agent/extensions/tmux-subagent/src/extension.test.ts
Normal file
397
.pi/agent/extensions/tmux-subagent/src/extension.test.ts
Normal file
@@ -0,0 +1,397 @@
|
|||||||
|
import test from "node:test";
|
||||||
|
import assert from "node:assert/strict";
|
||||||
|
import tmuxSubagentExtension from "../index.ts";
|
||||||
|
|
||||||
|
test("the extension entrypoint registers the subagent tool with the currently available models", async () => {
|
||||||
|
const original = process.env.PI_TMUX_SUBAGENT_CHILD;
|
||||||
|
if (original !== undefined) delete process.env.PI_TMUX_SUBAGENT_CHILD;
|
||||||
|
|
||||||
|
try {
|
||||||
|
const registeredTools: any[] = [];
|
||||||
|
const handlers: Record<string, (event: any, ctx: any) => Promise<void> | void> = {};
|
||||||
|
|
||||||
|
tmuxSubagentExtension({
|
||||||
|
on(event: string, handler: (event: any, ctx: any) => Promise<void> | void) {
|
||||||
|
handlers[event] = handler;
|
||||||
|
},
|
||||||
|
registerTool(tool: any) {
|
||||||
|
registeredTools.push(tool);
|
||||||
|
},
|
||||||
|
registerProvider() {},
|
||||||
|
} as any);
|
||||||
|
|
||||||
|
assert.equal(typeof handlers.session_start, "function");
|
||||||
|
|
||||||
|
await handlers.session_start?.(
|
||||||
|
{ reason: "startup" },
|
||||||
|
{
|
||||||
|
modelRegistry: {
|
||||||
|
getAvailable: () => [
|
||||||
|
{ provider: "anthropic", id: "claude-sonnet-4-5" },
|
||||||
|
{ provider: "openai", id: "gpt-5" },
|
||||||
|
],
|
||||||
|
},
|
||||||
|
},
|
||||||
|
);
|
||||||
|
|
||||||
|
assert.equal(registeredTools.length, 1);
|
||||||
|
assert.equal(registeredTools[0]?.name, "subagent");
|
||||||
|
assert.deepEqual(registeredTools[0]?.parameters.required, ["model"]);
|
||||||
|
assert.deepEqual(registeredTools[0]?.parameters.properties.model.enum, [
|
||||||
|
"anthropic/claude-sonnet-4-5",
|
||||||
|
"openai/gpt-5",
|
||||||
|
]);
|
||||||
|
assert.deepEqual(registeredTools[0]?.parameters.properties.tasks.items.properties.model.enum, [
|
||||||
|
"anthropic/claude-sonnet-4-5",
|
||||||
|
"openai/gpt-5",
|
||||||
|
]);
|
||||||
|
} finally {
|
||||||
|
if (original === undefined) delete process.env.PI_TMUX_SUBAGENT_CHILD;
|
||||||
|
else process.env.PI_TMUX_SUBAGENT_CHILD = original;
|
||||||
|
}
|
||||||
|
});
|
||||||
|
|
||||||
|
test("before_agent_start re-applies subagent registration when available models changed", async () => {
|
||||||
|
const original = process.env.PI_TMUX_SUBAGENT_CHILD;
|
||||||
|
if (original !== undefined) delete process.env.PI_TMUX_SUBAGENT_CHILD;
|
||||||
|
|
||||||
|
try {
|
||||||
|
const registeredTools: any[] = [];
|
||||||
|
const handlers: Record<string, (event: any, ctx: any) => Promise<void> | void> = {};
|
||||||
|
|
||||||
|
tmuxSubagentExtension({
|
||||||
|
on(event: string, handler: (event: any, ctx: any) => Promise<void> | void) {
|
||||||
|
handlers[event] = handler;
|
||||||
|
},
|
||||||
|
registerTool(tool: any) {
|
||||||
|
registeredTools.push(tool);
|
||||||
|
},
|
||||||
|
registerProvider() {},
|
||||||
|
} as any);
|
||||||
|
|
||||||
|
assert.equal(typeof handlers.session_start, "function");
|
||||||
|
assert.equal(typeof handlers.before_agent_start, "function");
|
||||||
|
|
||||||
|
// initial registration with two models
|
||||||
|
await handlers.session_start?.(
|
||||||
|
{ reason: "startup" },
|
||||||
|
{
|
||||||
|
modelRegistry: {
|
||||||
|
getAvailable: () => [
|
||||||
|
{ provider: "anthropic", id: "claude-sonnet-4-5" },
|
||||||
|
{ provider: "openai", id: "gpt-5" },
|
||||||
|
],
|
||||||
|
},
|
||||||
|
},
|
||||||
|
);
|
||||||
|
|
||||||
|
assert.equal(registeredTools.length, 1);
|
||||||
|
assert.deepEqual(registeredTools[0]?.parameters.properties.model.enum, [
|
||||||
|
"anthropic/claude-sonnet-4-5",
|
||||||
|
"openai/gpt-5",
|
||||||
|
]);
|
||||||
|
|
||||||
|
// then before agent start with a different model set — should re-register
|
||||||
|
await handlers.before_agent_start?.(
|
||||||
|
{ reason: "about-to-start" },
|
||||||
|
{
|
||||||
|
modelRegistry: {
|
||||||
|
getAvailable: () => [
|
||||||
|
{ provider: "openai", id: "gpt-6" },
|
||||||
|
],
|
||||||
|
},
|
||||||
|
},
|
||||||
|
);
|
||||||
|
|
||||||
|
assert.equal(registeredTools.length, 2);
|
||||||
|
assert.deepEqual(registeredTools[1]?.parameters.properties.model.enum, ["openai/gpt-6"]);
|
||||||
|
assert.deepEqual(registeredTools[1]?.parameters.properties.tasks.items.properties.model.enum, ["openai/gpt-6"]);
|
||||||
|
} finally {
|
||||||
|
if (original === undefined) delete process.env.PI_TMUX_SUBAGENT_CHILD;
|
||||||
|
else process.env.PI_TMUX_SUBAGENT_CHILD = original;
|
||||||
|
}
|
||||||
|
});
|
||||||
|
|
||||||
|
test("child subagent sessions skip registering the subagent tool", async () => {
|
||||||
|
const original = process.env.PI_TMUX_SUBAGENT_CHILD;
|
||||||
|
process.env.PI_TMUX_SUBAGENT_CHILD = "1";
|
||||||
|
|
||||||
|
try {
|
||||||
|
const registeredTools: any[] = [];
|
||||||
|
const handlers: Record<string, (event: any, ctx: any) => Promise<void> | void> = {};
|
||||||
|
|
||||||
|
tmuxSubagentExtension({
|
||||||
|
on(event: string, handler: (event: any, ctx: any) => Promise<void> | void) {
|
||||||
|
handlers[event] = handler;
|
||||||
|
},
|
||||||
|
registerTool(tool: any) {
|
||||||
|
registeredTools.push(tool);
|
||||||
|
},
|
||||||
|
registerProvider() {},
|
||||||
|
} as any);
|
||||||
|
|
||||||
|
assert.equal(typeof handlers.session_start, "undefined");
|
||||||
|
assert.equal(typeof handlers.before_agent_start, "undefined");
|
||||||
|
assert.equal(registeredTools.length, 0);
|
||||||
|
} finally {
|
||||||
|
if (original === undefined) delete process.env.PI_TMUX_SUBAGENT_CHILD;
|
||||||
|
else process.env.PI_TMUX_SUBAGENT_CHILD = original;
|
||||||
|
}
|
||||||
|
});
|
||||||
|
|
||||||
|
test("registers github-copilot provider override when PI_TMUX_SUBAGENT_GITHUB_COPILOT_INITIATOR is set", () => {
|
||||||
|
const registeredProviders: Array<{ name: string; config: any }> = [];
|
||||||
|
const originalInitiator = process.env.PI_TMUX_SUBAGENT_GITHUB_COPILOT_INITIATOR;
|
||||||
|
const originalChild = process.env.PI_TMUX_SUBAGENT_CHILD;
|
||||||
|
// Ensure we exercise the non-child code path for this test
|
||||||
|
if (originalChild !== undefined) delete process.env.PI_TMUX_SUBAGENT_CHILD;
|
||||||
|
process.env.PI_TMUX_SUBAGENT_GITHUB_COPILOT_INITIATOR = "agent";
|
||||||
|
|
||||||
|
try {
|
||||||
|
tmuxSubagentExtension({
|
||||||
|
on() {},
|
||||||
|
registerTool() {},
|
||||||
|
registerProvider(name: string, config: any) {
|
||||||
|
registeredProviders.push({ name, config });
|
||||||
|
},
|
||||||
|
} as any);
|
||||||
|
} finally {
|
||||||
|
if (originalInitiator === undefined) delete process.env.PI_TMUX_SUBAGENT_GITHUB_COPILOT_INITIATOR;
|
||||||
|
else process.env.PI_TMUX_SUBAGENT_GITHUB_COPILOT_INITIATOR = originalInitiator;
|
||||||
|
|
||||||
|
if (originalChild === undefined) delete process.env.PI_TMUX_SUBAGENT_CHILD;
|
||||||
|
else process.env.PI_TMUX_SUBAGENT_CHILD = originalChild;
|
||||||
|
}
|
||||||
|
|
||||||
|
assert.deepEqual(registeredProviders, [
|
||||||
|
{
|
||||||
|
name: "github-copilot",
|
||||||
|
config: { headers: { "X-Initiator": "agent" } },
|
||||||
|
},
|
||||||
|
]);
|
||||||
|
});
|
||||||
|
|
||||||
|
test("combined child+copilot run registers provider but no tools or startup handlers", () => {
|
||||||
|
const registeredProviders: Array<{ name: string; config: any }> = [];
|
||||||
|
const registeredTools: any[] = [];
|
||||||
|
const handlers: Record<string, (event: any, ctx: any) => Promise<void> | void> = {};
|
||||||
|
|
||||||
|
const originalInitiator = process.env.PI_TMUX_SUBAGENT_GITHUB_COPILOT_INITIATOR;
|
||||||
|
const originalChild = process.env.PI_TMUX_SUBAGENT_CHILD;
|
||||||
|
process.env.PI_TMUX_SUBAGENT_GITHUB_COPILOT_INITIATOR = "agent";
|
||||||
|
process.env.PI_TMUX_SUBAGENT_CHILD = "1";
|
||||||
|
|
||||||
|
try {
|
||||||
|
tmuxSubagentExtension({
|
||||||
|
on(event: string, handler: (event: any, ctx: any) => Promise<void> | void) {
|
||||||
|
handlers[event] = handler;
|
||||||
|
},
|
||||||
|
registerTool(tool: any) {
|
||||||
|
registeredTools.push(tool);
|
||||||
|
},
|
||||||
|
registerProvider(name: string, config: any) {
|
||||||
|
registeredProviders.push({ name, config });
|
||||||
|
},
|
||||||
|
} as any);
|
||||||
|
|
||||||
|
assert.deepEqual(registeredProviders, [
|
||||||
|
{
|
||||||
|
name: "github-copilot",
|
||||||
|
config: { headers: { "X-Initiator": "agent" } },
|
||||||
|
},
|
||||||
|
]);
|
||||||
|
|
||||||
|
assert.equal(registeredTools.length, 0);
|
||||||
|
assert.equal(typeof handlers.session_start, "undefined");
|
||||||
|
assert.equal(typeof handlers.before_agent_start, "undefined");
|
||||||
|
} finally {
|
||||||
|
if (originalInitiator === undefined) delete process.env.PI_TMUX_SUBAGENT_GITHUB_COPILOT_INITIATOR;
|
||||||
|
else process.env.PI_TMUX_SUBAGENT_GITHUB_COPILOT_INITIATOR = originalInitiator;
|
||||||
|
|
||||||
|
if (originalChild === undefined) delete process.env.PI_TMUX_SUBAGENT_CHILD;
|
||||||
|
else process.env.PI_TMUX_SUBAGENT_CHILD = originalChild;
|
||||||
|
}
|
||||||
|
});
|
||||||
|
|
||||||
|
test("does not re-register the subagent tool when models list unchanged, but re-registers when changed", async () => {
|
||||||
|
const original = process.env.PI_TMUX_SUBAGENT_CHILD;
|
||||||
|
if (original !== undefined) delete process.env.PI_TMUX_SUBAGENT_CHILD;
|
||||||
|
|
||||||
|
try {
|
||||||
|
let registerToolCalls = 0;
|
||||||
|
const handlers: Record<string, (event: any, ctx: any) => Promise<void> | void> = {};
|
||||||
|
|
||||||
|
tmuxSubagentExtension({
|
||||||
|
on(event: string, handler: (event: any, ctx: any) => Promise<void> | void) {
|
||||||
|
handlers[event] = handler;
|
||||||
|
},
|
||||||
|
registerTool() {
|
||||||
|
registerToolCalls++;
|
||||||
|
},
|
||||||
|
registerProvider() {},
|
||||||
|
} as any);
|
||||||
|
|
||||||
|
assert.equal(typeof handlers.session_start, "function");
|
||||||
|
assert.equal(typeof handlers.before_agent_start, "function");
|
||||||
|
|
||||||
|
// First registration with two models
|
||||||
|
await handlers.session_start?.(
|
||||||
|
{ reason: "startup" },
|
||||||
|
{
|
||||||
|
modelRegistry: {
|
||||||
|
getAvailable: () => [
|
||||||
|
{ provider: "anthropic", id: "claude-sonnet-4-5" },
|
||||||
|
{ provider: "openai", id: "gpt-5" },
|
||||||
|
],
|
||||||
|
},
|
||||||
|
},
|
||||||
|
);
|
||||||
|
|
||||||
|
assert.equal(registerToolCalls, 1);
|
||||||
|
|
||||||
|
// Second registration with the same models — should not increase count
|
||||||
|
await handlers.before_agent_start?.(
|
||||||
|
{ reason: "about-to-start" },
|
||||||
|
{
|
||||||
|
modelRegistry: {
|
||||||
|
getAvailable: () => [
|
||||||
|
{ provider: "anthropic", id: "claude-sonnet-4-5" },
|
||||||
|
{ provider: "openai", id: "gpt-5" },
|
||||||
|
],
|
||||||
|
},
|
||||||
|
},
|
||||||
|
);
|
||||||
|
|
||||||
|
assert.equal(registerToolCalls, 1);
|
||||||
|
|
||||||
|
// Third call with changed model list — should re-register
|
||||||
|
await handlers.session_start?.(
|
||||||
|
{ reason: "startup" },
|
||||||
|
{
|
||||||
|
modelRegistry: {
|
||||||
|
getAvailable: () => [
|
||||||
|
{ provider: "openai", id: "gpt-6" },
|
||||||
|
],
|
||||||
|
},
|
||||||
|
},
|
||||||
|
);
|
||||||
|
|
||||||
|
assert.equal(registerToolCalls, 2);
|
||||||
|
} finally {
|
||||||
|
if (original === undefined) delete process.env.PI_TMUX_SUBAGENT_CHILD;
|
||||||
|
else process.env.PI_TMUX_SUBAGENT_CHILD = original;
|
||||||
|
}
|
||||||
|
});
|
||||||
|
|
||||||
|
|
||||||
|
// New tests for robustness: order-independence and empty model handling
|
||||||
|
|
||||||
|
test("same model set in different orders should NOT trigger re-registration", async () => {
|
||||||
|
const original = process.env.PI_TMUX_SUBAGENT_CHILD;
|
||||||
|
if (original !== undefined) delete process.env.PI_TMUX_SUBAGENT_CHILD;
|
||||||
|
|
||||||
|
try {
|
||||||
|
let registerToolCalls = 0;
|
||||||
|
const handlers: Record<string, (event: any, ctx: any) => Promise<void> | void> = {};
|
||||||
|
|
||||||
|
tmuxSubagentExtension({
|
||||||
|
on(event: string, handler: (event: any, ctx: any) => Promise<void> | void) {
|
||||||
|
handlers[event] = handler;
|
||||||
|
},
|
||||||
|
registerTool() {
|
||||||
|
registerToolCalls++;
|
||||||
|
},
|
||||||
|
registerProvider() {},
|
||||||
|
} as any);
|
||||||
|
|
||||||
|
assert.equal(typeof handlers.session_start, "function");
|
||||||
|
assert.equal(typeof handlers.before_agent_start, "function");
|
||||||
|
|
||||||
|
// First registration with two models in one order
|
||||||
|
await handlers.session_start?.(
|
||||||
|
{ reason: "startup" },
|
||||||
|
{
|
||||||
|
modelRegistry: {
|
||||||
|
getAvailable: () => [
|
||||||
|
{ provider: "anthropic", id: "claude-sonnet-4-5" },
|
||||||
|
{ provider: "openai", id: "gpt-5" },
|
||||||
|
],
|
||||||
|
},
|
||||||
|
},
|
||||||
|
);
|
||||||
|
|
||||||
|
assert.equal(registerToolCalls, 1);
|
||||||
|
|
||||||
|
// Same models but reversed order — should NOT re-register
|
||||||
|
await handlers.before_agent_start?.(
|
||||||
|
{ reason: "about-to-start" },
|
||||||
|
{
|
||||||
|
modelRegistry: {
|
||||||
|
getAvailable: () => [
|
||||||
|
{ provider: "openai", id: "gpt-5" },
|
||||||
|
{ provider: "anthropic", id: "claude-sonnet-4-5" },
|
||||||
|
],
|
||||||
|
},
|
||||||
|
},
|
||||||
|
);
|
||||||
|
|
||||||
|
assert.equal(registerToolCalls, 1);
|
||||||
|
} finally {
|
||||||
|
if (original === undefined) delete process.env.PI_TMUX_SUBAGENT_CHILD;
|
||||||
|
else process.env.PI_TMUX_SUBAGENT_CHILD = original;
|
||||||
|
}
|
||||||
|
});
|
||||||
|
|
||||||
|
|
||||||
|
test("empty model list should NOT register the tool, but a later non-empty list should", async () => {
|
||||||
|
const original = process.env.PI_TMUX_SUBAGENT_CHILD;
|
||||||
|
if (original !== undefined) delete process.env.PI_TMUX_SUBAGENT_CHILD;
|
||||||
|
|
||||||
|
try {
|
||||||
|
let registerToolCalls = 0;
|
||||||
|
const handlers: Record<string, (event: any, ctx: any) => Promise<void> | void> = {};
|
||||||
|
|
||||||
|
tmuxSubagentExtension({
|
||||||
|
on(event: string, handler: (event: any, ctx: any) => Promise<void> | void) {
|
||||||
|
handlers[event] = handler;
|
||||||
|
},
|
||||||
|
registerTool() {
|
||||||
|
registerToolCalls++;
|
||||||
|
},
|
||||||
|
registerProvider() {},
|
||||||
|
} as any);
|
||||||
|
|
||||||
|
assert.equal(typeof handlers.session_start, "function");
|
||||||
|
assert.equal(typeof handlers.before_agent_start, "function");
|
||||||
|
|
||||||
|
// empty list should not register
|
||||||
|
await handlers.session_start?.(
|
||||||
|
{ reason: "startup" },
|
||||||
|
{
|
||||||
|
modelRegistry: {
|
||||||
|
getAvailable: () => [],
|
||||||
|
},
|
||||||
|
},
|
||||||
|
);
|
||||||
|
|
||||||
|
assert.equal(registerToolCalls, 0);
|
||||||
|
|
||||||
|
// later non-empty list should register
|
||||||
|
await handlers.before_agent_start?.(
|
||||||
|
{ reason: "about-to-start" },
|
||||||
|
{
|
||||||
|
modelRegistry: {
|
||||||
|
getAvailable: () => [
|
||||||
|
{ provider: "openai", id: "gpt-6" },
|
||||||
|
],
|
||||||
|
},
|
||||||
|
},
|
||||||
|
);
|
||||||
|
|
||||||
|
assert.equal(registerToolCalls, 1);
|
||||||
|
} finally {
|
||||||
|
if (original === undefined) delete process.env.PI_TMUX_SUBAGENT_CHILD;
|
||||||
|
else process.env.PI_TMUX_SUBAGENT_CHILD = original;
|
||||||
|
}
|
||||||
|
});
|
||||||
|
|
||||||
44
.pi/agent/extensions/tmux-subagent/src/models.test.ts
Normal file
44
.pi/agent/extensions/tmux-subagent/src/models.test.ts
Normal file
@@ -0,0 +1,44 @@
|
|||||||
|
import test from "node:test";
|
||||||
|
import assert from "node:assert/strict";
|
||||||
|
import {
|
||||||
|
formatModelReference,
|
||||||
|
listAvailableModelReferences,
|
||||||
|
normalizeAvailableModelReference,
|
||||||
|
resolveChildModel,
|
||||||
|
} from "./models.ts";
|
||||||
|
|
||||||
|
test("resolveChildModel prefers the per-task override over the required top-level model", () => {
|
||||||
|
const selection = resolveChildModel({
|
||||||
|
taskModel: "openai/gpt-5",
|
||||||
|
topLevelModel: "anthropic/claude-sonnet-4-5",
|
||||||
|
});
|
||||||
|
|
||||||
|
assert.equal(selection.requestedModel, "openai/gpt-5");
|
||||||
|
assert.equal(selection.resolvedModel, "openai/gpt-5");
|
||||||
|
});
|
||||||
|
|
||||||
|
test("formatModelReference returns provider/id", () => {
|
||||||
|
const ref = formatModelReference({ provider: "anthropic", id: "claude-sonnet-4-5" });
|
||||||
|
|
||||||
|
assert.equal(ref, "anthropic/claude-sonnet-4-5");
|
||||||
|
});
|
||||||
|
|
||||||
|
test("listAvailableModelReferences formats all configured available models", () => {
|
||||||
|
const refs = listAvailableModelReferences({
|
||||||
|
getAvailable: () => [
|
||||||
|
{ provider: "anthropic", id: "claude-sonnet-4-5" },
|
||||||
|
{ provider: "openai", id: "gpt-5" },
|
||||||
|
],
|
||||||
|
});
|
||||||
|
|
||||||
|
assert.deepEqual(refs, ["anthropic/claude-sonnet-4-5", "openai/gpt-5"]);
|
||||||
|
});
|
||||||
|
|
||||||
|
test("normalizeAvailableModelReference matches canonical refs case-insensitively", () => {
|
||||||
|
const normalized = normalizeAvailableModelReference("OpenAI/GPT-5", [
|
||||||
|
"anthropic/claude-sonnet-4-5",
|
||||||
|
"openai/gpt-5",
|
||||||
|
]);
|
||||||
|
|
||||||
|
assert.equal(normalized, "openai/gpt-5");
|
||||||
|
});
|
||||||
58
.pi/agent/extensions/tmux-subagent/src/models.ts
Normal file
58
.pi/agent/extensions/tmux-subagent/src/models.ts
Normal file
@@ -0,0 +1,58 @@
|
|||||||
|
export interface ModelLike {
|
||||||
|
provider: string;
|
||||||
|
id: string;
|
||||||
|
}
|
||||||
|
|
||||||
|
export interface AvailableModelRegistryLike {
|
||||||
|
getAvailable(): ModelLike[];
|
||||||
|
}
|
||||||
|
|
||||||
|
export interface ModelSelection {
|
||||||
|
requestedModel?: string;
|
||||||
|
resolvedModel?: string;
|
||||||
|
}
|
||||||
|
|
||||||
|
export function formatModelReference(model: ModelLike): string {
|
||||||
|
return `${model.provider}/${model.id}`;
|
||||||
|
}
|
||||||
|
|
||||||
|
export function listAvailableModelReferences(modelRegistry?: AvailableModelRegistryLike): string[] {
|
||||||
|
if (!modelRegistry) return [];
|
||||||
|
|
||||||
|
const seen = new Set<string>();
|
||||||
|
const refs: string[] = [];
|
||||||
|
for (const model of modelRegistry.getAvailable()) {
|
||||||
|
const ref = formatModelReference(model);
|
||||||
|
const key = ref.toLowerCase();
|
||||||
|
if (seen.has(key)) continue;
|
||||||
|
seen.add(key);
|
||||||
|
refs.push(ref);
|
||||||
|
}
|
||||||
|
|
||||||
|
return refs;
|
||||||
|
}
|
||||||
|
|
||||||
|
export function normalizeAvailableModelReference(
|
||||||
|
requestedModel: string | undefined,
|
||||||
|
availableModels: readonly string[],
|
||||||
|
): string | undefined {
|
||||||
|
if (typeof requestedModel !== "string") return undefined;
|
||||||
|
|
||||||
|
const trimmed = requestedModel.trim();
|
||||||
|
if (!trimmed) return undefined;
|
||||||
|
|
||||||
|
const normalized = trimmed.toLowerCase();
|
||||||
|
return availableModels.find((candidate) => candidate.toLowerCase() === normalized);
|
||||||
|
}
|
||||||
|
|
||||||
|
export function resolveChildModel(input: {
|
||||||
|
taskModel?: string;
|
||||||
|
topLevelModel: string;
|
||||||
|
}): ModelSelection {
|
||||||
|
const requestedModel = input.taskModel ?? input.topLevelModel;
|
||||||
|
|
||||||
|
return {
|
||||||
|
requestedModel,
|
||||||
|
resolvedModel: requestedModel,
|
||||||
|
};
|
||||||
|
}
|
||||||
33
.pi/agent/extensions/tmux-subagent/src/monitor.test.ts
Normal file
33
.pi/agent/extensions/tmux-subagent/src/monitor.test.ts
Normal file
@@ -0,0 +1,33 @@
|
|||||||
|
import test from "node:test";
|
||||||
|
import assert from "node:assert/strict";
|
||||||
|
import { mkdtemp, writeFile, appendFile } from "node:fs/promises";
|
||||||
|
import { join } from "node:path";
|
||||||
|
import { tmpdir } from "node:os";
|
||||||
|
import { monitorRun } from "./monitor.ts";
|
||||||
|
|
||||||
|
test("monitorRun streams normalized events and resolves when result.json appears", async () => {
|
||||||
|
const dir = await mkdtemp(join(tmpdir(), "tmux-subagent-monitor-"));
|
||||||
|
const eventsPath = join(dir, "events.jsonl");
|
||||||
|
const resultPath = join(dir, "result.json");
|
||||||
|
await writeFile(eventsPath, "", "utf8");
|
||||||
|
|
||||||
|
const seen: string[] = [];
|
||||||
|
const waiting = monitorRun({
|
||||||
|
eventsPath,
|
||||||
|
resultPath,
|
||||||
|
onEvent(event) {
|
||||||
|
seen.push(event.type);
|
||||||
|
},
|
||||||
|
});
|
||||||
|
|
||||||
|
await appendFile(eventsPath, `${JSON.stringify({ type: "tool_call", toolName: "read", args: { path: "a.ts" } })}\n`, "utf8");
|
||||||
|
await writeFile(
|
||||||
|
resultPath,
|
||||||
|
JSON.stringify({ runId: "run-1", exitCode: 0, finalText: "done", agent: "scout", task: "inspect auth" }, null, 2),
|
||||||
|
"utf8",
|
||||||
|
);
|
||||||
|
|
||||||
|
const result = await waiting;
|
||||||
|
assert.deepEqual(seen, ["tool_call"]);
|
||||||
|
assert.equal(result.finalText, "done");
|
||||||
|
});
|
||||||
34
.pi/agent/extensions/tmux-subagent/src/monitor.ts
Normal file
34
.pi/agent/extensions/tmux-subagent/src/monitor.ts
Normal file
@@ -0,0 +1,34 @@
|
|||||||
|
import { existsSync } from "node:fs";
|
||||||
|
import { readFile } from "node:fs/promises";
|
||||||
|
|
||||||
|
async function sleep(ms: number) {
|
||||||
|
await new Promise((resolve) => setTimeout(resolve, ms));
|
||||||
|
}
|
||||||
|
|
||||||
|
export async function monitorRun(input: {
|
||||||
|
eventsPath: string;
|
||||||
|
resultPath: string;
|
||||||
|
onEvent?: (event: any) => void;
|
||||||
|
pollMs?: number;
|
||||||
|
}) {
|
||||||
|
const pollMs = input.pollMs ?? 50;
|
||||||
|
let offset = 0;
|
||||||
|
|
||||||
|
while (true) {
|
||||||
|
if (existsSync(input.eventsPath)) {
|
||||||
|
const text = await readFile(input.eventsPath, "utf8");
|
||||||
|
const next = text.slice(offset);
|
||||||
|
offset = text.length;
|
||||||
|
|
||||||
|
for (const line of next.split("\n").filter(Boolean)) {
|
||||||
|
input.onEvent?.(JSON.parse(line));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
if (existsSync(input.resultPath)) {
|
||||||
|
return JSON.parse(await readFile(input.resultPath, "utf8"));
|
||||||
|
}
|
||||||
|
|
||||||
|
await sleep(pollMs);
|
||||||
|
}
|
||||||
|
}
|
||||||
20
.pi/agent/extensions/tmux-subagent/src/prompts.test.ts
Normal file
20
.pi/agent/extensions/tmux-subagent/src/prompts.test.ts
Normal file
@@ -0,0 +1,20 @@
|
|||||||
|
import test from "node:test";
|
||||||
|
import assert from "node:assert/strict";
|
||||||
|
import { readFileSync } from "node:fs";
|
||||||
|
import { dirname, join } from "node:path";
|
||||||
|
import { fileURLToPath } from "node:url";
|
||||||
|
|
||||||
|
const packageRoot = dirname(dirname(fileURLToPath(import.meta.url)));
|
||||||
|
|
||||||
|
test("package.json exposes the extension and workflow prompt templates", () => {
|
||||||
|
const packageJson = JSON.parse(readFileSync(join(packageRoot, "package.json"), "utf8"));
|
||||||
|
|
||||||
|
assert.deepEqual(packageJson.pi.extensions, ["./index.ts"]);
|
||||||
|
assert.deepEqual(packageJson.pi.prompts, ["./prompts/*.md"]);
|
||||||
|
|
||||||
|
for (const name of ["implement.md", "scout-and-plan.md", "implement-and-review.md"]) {
|
||||||
|
const content = readFileSync(join(packageRoot, "prompts", name), "utf8");
|
||||||
|
assert.match(content, /^---\ndescription:/m);
|
||||||
|
assert.match(content, /subagent/);
|
||||||
|
}
|
||||||
|
});
|
||||||
33
.pi/agent/extensions/tmux-subagent/src/runner.test.ts
Normal file
33
.pi/agent/extensions/tmux-subagent/src/runner.test.ts
Normal file
@@ -0,0 +1,33 @@
|
|||||||
|
import test from "node:test";
|
||||||
|
import assert from "node:assert/strict";
|
||||||
|
import { createTmuxSingleRunner } from "./runner.ts";
|
||||||
|
|
||||||
|
test("createTmuxSingleRunner always kills the pane after monitor completion", async () => {
|
||||||
|
const killed: string[] = [];
|
||||||
|
|
||||||
|
const runSingleTask = createTmuxSingleRunner({
|
||||||
|
assertInsideTmux() {},
|
||||||
|
getCurrentWindowId: async () => "@1",
|
||||||
|
createArtifacts: async () => ({
|
||||||
|
metaPath: "/tmp/meta.json",
|
||||||
|
runId: "run-1",
|
||||||
|
eventsPath: "/tmp/events.jsonl",
|
||||||
|
resultPath: "/tmp/result.json",
|
||||||
|
sessionPath: "/tmp/child-session.jsonl",
|
||||||
|
stdoutPath: "/tmp/stdout.log",
|
||||||
|
stderrPath: "/tmp/stderr.log",
|
||||||
|
}),
|
||||||
|
buildWrapperCommand: () => "'node' '/wrapper.mjs' '/tmp/meta.json'",
|
||||||
|
createPane: async () => "%9",
|
||||||
|
monitorRun: async () => ({ finalText: "done", exitCode: 0 }),
|
||||||
|
killPane: async (paneId: string) => {
|
||||||
|
killed.push(paneId);
|
||||||
|
},
|
||||||
|
});
|
||||||
|
|
||||||
|
const result = await runSingleTask({ cwd: "/repo", meta: { task: "inspect auth" } as any });
|
||||||
|
|
||||||
|
assert.equal(result.paneId, "%9");
|
||||||
|
assert.equal(result.finalText, "done");
|
||||||
|
assert.deepEqual(killed, ["%9"]);
|
||||||
|
});
|
||||||
44
.pi/agent/extensions/tmux-subagent/src/runner.ts
Normal file
44
.pi/agent/extensions/tmux-subagent/src/runner.ts
Normal file
@@ -0,0 +1,44 @@
|
|||||||
|
export function createTmuxSingleRunner(deps: {
|
||||||
|
assertInsideTmux(): void;
|
||||||
|
getCurrentWindowId: () => Promise<string>;
|
||||||
|
createArtifacts: (cwd: string, meta: Record<string, unknown>) => Promise<any>;
|
||||||
|
buildWrapperCommand: (metaPath: string) => string;
|
||||||
|
createPane: (input: { windowId: string; cwd: string; command: string }) => Promise<string>;
|
||||||
|
monitorRun: (input: { eventsPath: string; resultPath: string; onEvent?: (event: any) => void }) => Promise<any>;
|
||||||
|
killPane: (paneId: string) => Promise<void>;
|
||||||
|
}) {
|
||||||
|
return async function runSingleTask(input: {
|
||||||
|
cwd: string;
|
||||||
|
meta: Record<string, unknown>;
|
||||||
|
onEvent?: (event: any) => void;
|
||||||
|
}) {
|
||||||
|
deps.assertInsideTmux();
|
||||||
|
|
||||||
|
const artifacts = await deps.createArtifacts(input.cwd, input.meta);
|
||||||
|
const windowId = await deps.getCurrentWindowId();
|
||||||
|
const command = deps.buildWrapperCommand(artifacts.metaPath);
|
||||||
|
const paneId = await deps.createPane({ windowId, cwd: input.cwd, command });
|
||||||
|
|
||||||
|
try {
|
||||||
|
const result = await deps.monitorRun({
|
||||||
|
eventsPath: artifacts.eventsPath,
|
||||||
|
resultPath: artifacts.resultPath,
|
||||||
|
onEvent: input.onEvent,
|
||||||
|
});
|
||||||
|
|
||||||
|
return {
|
||||||
|
...result,
|
||||||
|
runId: result.runId ?? artifacts.runId,
|
||||||
|
paneId,
|
||||||
|
windowId,
|
||||||
|
sessionPath: result.sessionPath ?? artifacts.sessionPath,
|
||||||
|
stdoutPath: result.stdoutPath ?? artifacts.stdoutPath,
|
||||||
|
stderrPath: result.stderrPath ?? artifacts.stderrPath,
|
||||||
|
resultPath: artifacts.resultPath,
|
||||||
|
eventsPath: artifacts.eventsPath,
|
||||||
|
};
|
||||||
|
} finally {
|
||||||
|
await deps.killPane(paneId);
|
||||||
|
}
|
||||||
|
};
|
||||||
|
}
|
||||||
84
.pi/agent/extensions/tmux-subagent/src/schema.ts
Normal file
84
.pi/agent/extensions/tmux-subagent/src/schema.ts
Normal file
@@ -0,0 +1,84 @@
|
|||||||
|
import { StringEnum } from "@mariozechner/pi-ai";
|
||||||
|
import { Type, type Static } from "@sinclair/typebox";
|
||||||
|
|
||||||
|
function createTaskModelSchema(availableModels: readonly string[]) {
|
||||||
|
return Type.Optional(
|
||||||
|
StringEnum(availableModels, {
|
||||||
|
description: "Optional child model override. Must be one of the currently available models.",
|
||||||
|
}),
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
export function createTaskItemSchema(availableModels: readonly string[]) {
|
||||||
|
return Type.Object({
|
||||||
|
agent: Type.String({ description: "Name of the agent to invoke" }),
|
||||||
|
task: Type.String({ description: "Task to delegate to the child agent" }),
|
||||||
|
model: createTaskModelSchema(availableModels),
|
||||||
|
cwd: Type.Optional(Type.String({ description: "Optional working directory override" })),
|
||||||
|
});
|
||||||
|
}
|
||||||
|
|
||||||
|
export function createChainItemSchema(availableModels: readonly string[]) {
|
||||||
|
return Type.Object({
|
||||||
|
agent: Type.String({ description: "Name of the agent to invoke" }),
|
||||||
|
task: Type.String({ description: "Task with optional {previous} placeholder" }),
|
||||||
|
model: createTaskModelSchema(availableModels),
|
||||||
|
cwd: Type.Optional(Type.String({ description: "Optional working directory override" })),
|
||||||
|
});
|
||||||
|
}
|
||||||
|
|
||||||
|
export const TaskItemSchema = createTaskItemSchema([]);
|
||||||
|
export const ChainItemSchema = createChainItemSchema([]);
|
||||||
|
|
||||||
|
export const AgentScopeSchema = StringEnum(["user", "project", "both"] as const, {
|
||||||
|
description: "Which markdown agent sources to use",
|
||||||
|
default: "user",
|
||||||
|
});
|
||||||
|
|
||||||
|
export function createSubagentParamsSchema(availableModels: readonly string[]) {
|
||||||
|
return Type.Object({
|
||||||
|
agent: Type.Optional(Type.String({ description: "Single-mode agent name" })),
|
||||||
|
task: Type.Optional(Type.String({ description: "Single-mode delegated task" })),
|
||||||
|
model: StringEnum(availableModels, {
|
||||||
|
description: "Required top-level child model. Must be one of the currently available models.",
|
||||||
|
}),
|
||||||
|
tasks: Type.Optional(Type.Array(createTaskItemSchema(availableModels), { description: "Parallel tasks" })),
|
||||||
|
chain: Type.Optional(Type.Array(createChainItemSchema(availableModels), { description: "Sequential tasks" })),
|
||||||
|
agentScope: Type.Optional(AgentScopeSchema),
|
||||||
|
confirmProjectAgents: Type.Optional(Type.Boolean({ default: true })),
|
||||||
|
cwd: Type.Optional(Type.String({ description: "Single-mode working directory override" })),
|
||||||
|
});
|
||||||
|
}
|
||||||
|
|
||||||
|
export const SubagentParamsSchema = createSubagentParamsSchema([]);
|
||||||
|
|
||||||
|
export type TaskItem = Static<typeof TaskItemSchema>;
|
||||||
|
export type ChainItem = Static<typeof ChainItemSchema>;
|
||||||
|
export type SubagentParams = Static<typeof SubagentParamsSchema>;
|
||||||
|
export type AgentScope = Static<typeof AgentScopeSchema>;
|
||||||
|
|
||||||
|
export interface SubagentRunResult {
|
||||||
|
runId: string;
|
||||||
|
agent: string;
|
||||||
|
agentSource: "builtin" | "user" | "project" | "unknown";
|
||||||
|
task: string;
|
||||||
|
requestedModel?: string;
|
||||||
|
resolvedModel?: string;
|
||||||
|
paneId?: string;
|
||||||
|
windowId?: string;
|
||||||
|
sessionPath?: string;
|
||||||
|
exitCode: number;
|
||||||
|
stopReason?: string;
|
||||||
|
finalText: string;
|
||||||
|
stdoutPath?: string;
|
||||||
|
stderrPath?: string;
|
||||||
|
resultPath?: string;
|
||||||
|
eventsPath?: string;
|
||||||
|
}
|
||||||
|
|
||||||
|
export interface SubagentToolDetails {
|
||||||
|
mode: "single" | "parallel" | "chain";
|
||||||
|
agentScope: AgentScope;
|
||||||
|
projectAgentsDir: string | null;
|
||||||
|
results: SubagentRunResult[];
|
||||||
|
}
|
||||||
43
.pi/agent/extensions/tmux-subagent/src/tmux.test.ts
Normal file
43
.pi/agent/extensions/tmux-subagent/src/tmux.test.ts
Normal file
@@ -0,0 +1,43 @@
|
|||||||
|
import test from "node:test";
|
||||||
|
import assert from "node:assert/strict";
|
||||||
|
import {
|
||||||
|
buildSplitWindowArgs,
|
||||||
|
buildWrapperShellCommand,
|
||||||
|
isInsideTmux,
|
||||||
|
} from "./tmux.ts";
|
||||||
|
|
||||||
|
test("isInsideTmux reads the TMUX environment variable", () => {
|
||||||
|
assert.equal(isInsideTmux({ TMUX: "/tmp/tmux-1000/default,123,0" } as NodeJS.ProcessEnv), true);
|
||||||
|
assert.equal(isInsideTmux({} as NodeJS.ProcessEnv), false);
|
||||||
|
});
|
||||||
|
|
||||||
|
test("buildWrapperShellCommand single-quotes paths safely", () => {
|
||||||
|
const command = buildWrapperShellCommand({
|
||||||
|
nodePath: "/usr/local/bin/node",
|
||||||
|
wrapperPath: "/repo/tmux-subagent/src/wrapper/cli.mjs",
|
||||||
|
metaPath: "/repo/.pi/subagents/runs/run-1/meta.json",
|
||||||
|
});
|
||||||
|
|
||||||
|
assert.equal(
|
||||||
|
command,
|
||||||
|
"'/usr/local/bin/node' '/repo/tmux-subagent/src/wrapper/cli.mjs' '/repo/.pi/subagents/runs/run-1/meta.json'",
|
||||||
|
);
|
||||||
|
});
|
||||||
|
|
||||||
|
test("buildSplitWindowArgs targets the current window and cwd", () => {
|
||||||
|
assert.deepEqual(buildSplitWindowArgs({
|
||||||
|
windowId: "@7",
|
||||||
|
cwd: "/repo",
|
||||||
|
command: "'node' '/wrapper.mjs' '/meta.json'",
|
||||||
|
}), [
|
||||||
|
"split-window",
|
||||||
|
"-P",
|
||||||
|
"-F",
|
||||||
|
"#{pane_id}",
|
||||||
|
"-t",
|
||||||
|
"@7",
|
||||||
|
"-c",
|
||||||
|
"/repo",
|
||||||
|
"'node' '/wrapper.mjs' '/meta.json'",
|
||||||
|
]);
|
||||||
|
});
|
||||||
41
.pi/agent/extensions/tmux-subagent/src/tmux.ts
Normal file
41
.pi/agent/extensions/tmux-subagent/src/tmux.ts
Normal file
@@ -0,0 +1,41 @@
|
|||||||
|
export function isInsideTmux(env: NodeJS.ProcessEnv = process.env): boolean {
|
||||||
|
return typeof env.TMUX === "string" && env.TMUX.length > 0;
|
||||||
|
}
|
||||||
|
|
||||||
|
function shellEscape(value: string): string {
|
||||||
|
return `'${value.replaceAll("'", "'\\''")}'`;
|
||||||
|
}
|
||||||
|
|
||||||
|
export function buildWrapperShellCommand(input: {
|
||||||
|
nodePath: string;
|
||||||
|
wrapperPath: string;
|
||||||
|
metaPath: string;
|
||||||
|
}): string {
|
||||||
|
return [input.nodePath, input.wrapperPath, input.metaPath].map(shellEscape).join(" ");
|
||||||
|
}
|
||||||
|
|
||||||
|
export function buildSplitWindowArgs(input: {
|
||||||
|
windowId: string;
|
||||||
|
cwd: string;
|
||||||
|
command: string;
|
||||||
|
}): string[] {
|
||||||
|
return [
|
||||||
|
"split-window",
|
||||||
|
"-P",
|
||||||
|
"-F",
|
||||||
|
"#{pane_id}",
|
||||||
|
"-t",
|
||||||
|
input.windowId,
|
||||||
|
"-c",
|
||||||
|
input.cwd,
|
||||||
|
input.command,
|
||||||
|
];
|
||||||
|
}
|
||||||
|
|
||||||
|
export function buildKillPaneArgs(paneId: string): string[] {
|
||||||
|
return ["kill-pane", "-t", paneId];
|
||||||
|
}
|
||||||
|
|
||||||
|
export function buildCurrentWindowArgs(): string[] {
|
||||||
|
return ["display-message", "-p", "#{window_id}"];
|
||||||
|
}
|
||||||
107
.pi/agent/extensions/tmux-subagent/src/tool-chain.test.ts
Normal file
107
.pi/agent/extensions/tmux-subagent/src/tool-chain.test.ts
Normal file
@@ -0,0 +1,107 @@
|
|||||||
|
import test from "node:test";
|
||||||
|
import assert from "node:assert/strict";
|
||||||
|
import { createSubagentTool } from "./tool.ts";
|
||||||
|
|
||||||
|
test("chain mode substitutes {previous} into the next task", async () => {
|
||||||
|
const seenTasks: string[] = [];
|
||||||
|
|
||||||
|
const tool = createSubagentTool({
|
||||||
|
discoverAgents: () => ({
|
||||||
|
agents: [
|
||||||
|
{ name: "scout", description: "Scout", source: "builtin", systemPrompt: "Scout prompt" },
|
||||||
|
{ name: "planner", description: "Planner", source: "builtin", systemPrompt: "Planner prompt" },
|
||||||
|
],
|
||||||
|
projectAgentsDir: null,
|
||||||
|
}),
|
||||||
|
runSingleTask: async ({ meta }: any) => {
|
||||||
|
seenTasks.push(meta.task);
|
||||||
|
return {
|
||||||
|
runId: `${meta.agent}-${seenTasks.length}`,
|
||||||
|
agent: meta.agent,
|
||||||
|
agentSource: meta.agentSource,
|
||||||
|
task: meta.task,
|
||||||
|
exitCode: 0,
|
||||||
|
finalText: meta.agent === "scout" ? "Scout output" : "Plan output",
|
||||||
|
};
|
||||||
|
},
|
||||||
|
} as any);
|
||||||
|
|
||||||
|
const result = await tool.execute(
|
||||||
|
"tool-1",
|
||||||
|
{
|
||||||
|
model: "anthropic/claude-sonnet-4-5",
|
||||||
|
chain: [
|
||||||
|
{ agent: "scout", task: "inspect auth" },
|
||||||
|
{ agent: "planner", task: "use this context: {previous}" },
|
||||||
|
],
|
||||||
|
},
|
||||||
|
undefined,
|
||||||
|
undefined,
|
||||||
|
{
|
||||||
|
cwd: "/repo",
|
||||||
|
modelRegistry: {
|
||||||
|
getAvailable: () => [{ provider: "anthropic", id: "claude-sonnet-4-5" }],
|
||||||
|
},
|
||||||
|
hasUI: false,
|
||||||
|
} as any,
|
||||||
|
);
|
||||||
|
|
||||||
|
assert.deepEqual(seenTasks, ["inspect auth", "use this context: Scout output"]);
|
||||||
|
assert.equal(result.content[0]?.type === "text" ? result.content[0].text : "", "Plan output");
|
||||||
|
});
|
||||||
|
|
||||||
|
test("chain mode stops on the first failed step", async () => {
|
||||||
|
const tool = createSubagentTool({
|
||||||
|
discoverAgents: () => ({
|
||||||
|
agents: [
|
||||||
|
{ name: "scout", description: "Scout", source: "builtin", systemPrompt: "Scout prompt" },
|
||||||
|
{ name: "planner", description: "Planner", source: "builtin", systemPrompt: "Planner prompt" },
|
||||||
|
],
|
||||||
|
projectAgentsDir: null,
|
||||||
|
}),
|
||||||
|
runSingleTask: async ({ meta }: any) => {
|
||||||
|
if (meta.agent === "planner") {
|
||||||
|
return {
|
||||||
|
runId: "planner-2",
|
||||||
|
agent: meta.agent,
|
||||||
|
agentSource: meta.agentSource,
|
||||||
|
task: meta.task,
|
||||||
|
exitCode: 1,
|
||||||
|
finalText: "",
|
||||||
|
stopReason: "error",
|
||||||
|
};
|
||||||
|
}
|
||||||
|
return {
|
||||||
|
runId: "scout-1",
|
||||||
|
agent: meta.agent,
|
||||||
|
agentSource: meta.agentSource,
|
||||||
|
task: meta.task,
|
||||||
|
exitCode: 0,
|
||||||
|
finalText: "Scout output",
|
||||||
|
};
|
||||||
|
},
|
||||||
|
} as any);
|
||||||
|
|
||||||
|
const result = await tool.execute(
|
||||||
|
"tool-1",
|
||||||
|
{
|
||||||
|
model: "anthropic/claude-sonnet-4-5",
|
||||||
|
chain: [
|
||||||
|
{ agent: "scout", task: "inspect auth" },
|
||||||
|
{ agent: "planner", task: "use this context: {previous}" },
|
||||||
|
],
|
||||||
|
},
|
||||||
|
undefined,
|
||||||
|
undefined,
|
||||||
|
{
|
||||||
|
cwd: "/repo",
|
||||||
|
modelRegistry: {
|
||||||
|
getAvailable: () => [{ provider: "anthropic", id: "claude-sonnet-4-5" }],
|
||||||
|
},
|
||||||
|
hasUI: false,
|
||||||
|
} as any,
|
||||||
|
);
|
||||||
|
|
||||||
|
assert.equal(result.isError, true);
|
||||||
|
assert.match(result.content[0]?.type === "text" ? result.content[0].text : "", /Chain stopped at step 2/);
|
||||||
|
});
|
||||||
97
.pi/agent/extensions/tmux-subagent/src/tool-parallel.test.ts
Normal file
97
.pi/agent/extensions/tmux-subagent/src/tool-parallel.test.ts
Normal file
@@ -0,0 +1,97 @@
|
|||||||
|
import test from "node:test";
|
||||||
|
import assert from "node:assert/strict";
|
||||||
|
import { createSubagentTool } from "./tool.ts";
|
||||||
|
|
||||||
|
test("parallel mode runs each task and uses the top-level model unless a task overrides it", async () => {
|
||||||
|
const requestedModels: Array<string | undefined> = [];
|
||||||
|
|
||||||
|
const tool = createSubagentTool({
|
||||||
|
discoverAgents: () => ({
|
||||||
|
agents: [
|
||||||
|
{ name: "scout", description: "Scout", source: "builtin", systemPrompt: "Scout prompt" },
|
||||||
|
{ name: "reviewer", description: "Reviewer", source: "builtin", systemPrompt: "Reviewer prompt" },
|
||||||
|
],
|
||||||
|
projectAgentsDir: null,
|
||||||
|
}),
|
||||||
|
resolveChildModel: ({ taskModel, topLevelModel }: any) => ({
|
||||||
|
requestedModel: taskModel ?? topLevelModel,
|
||||||
|
resolvedModel: taskModel ?? topLevelModel,
|
||||||
|
}),
|
||||||
|
runSingleTask: async ({ meta }: any) => {
|
||||||
|
requestedModels.push(meta.requestedModel);
|
||||||
|
return {
|
||||||
|
runId: `${meta.agent}-${meta.task}`,
|
||||||
|
agent: meta.agent,
|
||||||
|
agentSource: meta.agentSource,
|
||||||
|
task: meta.task,
|
||||||
|
requestedModel: meta.requestedModel,
|
||||||
|
resolvedModel: meta.requestedModel,
|
||||||
|
exitCode: 0,
|
||||||
|
finalText: `${meta.agent}:${meta.task}`,
|
||||||
|
};
|
||||||
|
},
|
||||||
|
} as any);
|
||||||
|
|
||||||
|
const result = await tool.execute(
|
||||||
|
"tool-1",
|
||||||
|
{
|
||||||
|
model: "openai/gpt-5",
|
||||||
|
tasks: [
|
||||||
|
{ agent: "scout", task: "find auth code" },
|
||||||
|
{ agent: "reviewer", task: "review auth code", model: "anthropic/claude-opus-4-5" },
|
||||||
|
],
|
||||||
|
},
|
||||||
|
undefined,
|
||||||
|
undefined,
|
||||||
|
{
|
||||||
|
cwd: "/repo",
|
||||||
|
modelRegistry: {
|
||||||
|
getAvailable: () => [
|
||||||
|
{ provider: "openai", id: "gpt-5" },
|
||||||
|
{ provider: "anthropic", id: "claude-opus-4-5" },
|
||||||
|
],
|
||||||
|
},
|
||||||
|
hasUI: false,
|
||||||
|
} as any,
|
||||||
|
);
|
||||||
|
|
||||||
|
const text = result.content[0]?.type === "text" ? result.content[0].text : "";
|
||||||
|
assert.match(text, /2\/2 succeeded/);
|
||||||
|
assert.deepEqual(requestedModels, ["openai/gpt-5", "anthropic/claude-opus-4-5"]);
|
||||||
|
});
|
||||||
|
|
||||||
|
test("parallel mode rejects per-task model overrides that are not currently available", async () => {
|
||||||
|
let didRun = false;
|
||||||
|
|
||||||
|
const tool = createSubagentTool({
|
||||||
|
discoverAgents: () => ({
|
||||||
|
agents: [{ name: "scout", description: "Scout", source: "builtin", systemPrompt: "Scout prompt" }],
|
||||||
|
projectAgentsDir: null,
|
||||||
|
}),
|
||||||
|
runSingleTask: async () => {
|
||||||
|
didRun = true;
|
||||||
|
throw new Error("should not run");
|
||||||
|
},
|
||||||
|
} as any);
|
||||||
|
|
||||||
|
const result = await tool.execute(
|
||||||
|
"tool-1",
|
||||||
|
{
|
||||||
|
model: "anthropic/claude-sonnet-4-5",
|
||||||
|
tasks: [{ agent: "scout", task: "find auth code", model: "openai/gpt-5" }],
|
||||||
|
},
|
||||||
|
undefined,
|
||||||
|
undefined,
|
||||||
|
{
|
||||||
|
cwd: "/repo",
|
||||||
|
modelRegistry: {
|
||||||
|
getAvailable: () => [{ provider: "anthropic", id: "claude-sonnet-4-5" }],
|
||||||
|
},
|
||||||
|
hasUI: false,
|
||||||
|
} as any,
|
||||||
|
);
|
||||||
|
|
||||||
|
assert.equal(didRun, false);
|
||||||
|
assert.equal(result.isError, true);
|
||||||
|
assert.match(result.content[0]?.type === "text" ? result.content[0].text : "", /parallel task 1/i);
|
||||||
|
});
|
||||||
177
.pi/agent/extensions/tmux-subagent/src/tool.test.ts
Normal file
177
.pi/agent/extensions/tmux-subagent/src/tool.test.ts
Normal file
@@ -0,0 +1,177 @@
|
|||||||
|
import test from "node:test";
|
||||||
|
import assert from "node:assert/strict";
|
||||||
|
import { createSubagentTool } from "./tool.ts";
|
||||||
|
|
||||||
|
test("single-mode subagent uses the required top-level model, emits progress, and returns final text plus metadata", async () => {
|
||||||
|
const updates: string[] = [];
|
||||||
|
|
||||||
|
const tool = createSubagentTool({
|
||||||
|
discoverAgents: () => ({
|
||||||
|
agents: [
|
||||||
|
{
|
||||||
|
name: "scout",
|
||||||
|
description: "Scout",
|
||||||
|
model: "claude-haiku-4-5",
|
||||||
|
systemPrompt: "Scout prompt",
|
||||||
|
source: "builtin",
|
||||||
|
},
|
||||||
|
],
|
||||||
|
projectAgentsDir: null,
|
||||||
|
}),
|
||||||
|
runSingleTask: async ({ onEvent, meta }: any) => {
|
||||||
|
onEvent?.({ type: "tool_call", toolName: "read", args: { path: "src/auth.ts" } });
|
||||||
|
return {
|
||||||
|
runId: "run-1",
|
||||||
|
agent: "scout",
|
||||||
|
agentSource: "builtin",
|
||||||
|
task: "inspect auth",
|
||||||
|
requestedModel: meta.requestedModel,
|
||||||
|
resolvedModel: meta.resolvedModel,
|
||||||
|
paneId: "%3",
|
||||||
|
windowId: "@1",
|
||||||
|
sessionPath: "/repo/.pi/subagents/runs/run-1/child-session.jsonl",
|
||||||
|
exitCode: 0,
|
||||||
|
finalText: "Auth code is in src/auth.ts",
|
||||||
|
};
|
||||||
|
},
|
||||||
|
} as any);
|
||||||
|
|
||||||
|
const result = await tool.execute(
|
||||||
|
"tool-1",
|
||||||
|
{
|
||||||
|
agent: "scout",
|
||||||
|
task: "inspect auth",
|
||||||
|
model: "anthropic/claude-sonnet-4-5",
|
||||||
|
},
|
||||||
|
undefined,
|
||||||
|
(partial: any) => {
|
||||||
|
const first = partial.content?.[0];
|
||||||
|
if (first?.type === "text") updates.push(first.text);
|
||||||
|
},
|
||||||
|
{
|
||||||
|
cwd: "/repo",
|
||||||
|
modelRegistry: {
|
||||||
|
getAvailable: () => [{ provider: "anthropic", id: "claude-sonnet-4-5" }],
|
||||||
|
},
|
||||||
|
hasUI: false,
|
||||||
|
} as any,
|
||||||
|
);
|
||||||
|
|
||||||
|
const text = result.content[0]?.type === "text" ? result.content[0].text : "";
|
||||||
|
assert.equal(text, "Auth code is in src/auth.ts");
|
||||||
|
assert.equal(result.details.results[0]?.paneId, "%3");
|
||||||
|
assert.equal(result.details.results[0]?.requestedModel, "anthropic/claude-sonnet-4-5");
|
||||||
|
assert.match(updates.join("\n"), /Running scout/);
|
||||||
|
});
|
||||||
|
|
||||||
|
test("single-mode subagent requires a top-level model even when execute is called directly", async () => {
|
||||||
|
let didRun = false;
|
||||||
|
|
||||||
|
const tool = createSubagentTool({
|
||||||
|
discoverAgents: () => ({
|
||||||
|
agents: [{ name: "scout", description: "Scout", systemPrompt: "Scout prompt", source: "builtin" }],
|
||||||
|
projectAgentsDir: null,
|
||||||
|
}),
|
||||||
|
runSingleTask: async () => {
|
||||||
|
didRun = true;
|
||||||
|
throw new Error("should not run");
|
||||||
|
},
|
||||||
|
} as any);
|
||||||
|
|
||||||
|
const result = await tool.execute(
|
||||||
|
"tool-1",
|
||||||
|
{ agent: "scout", task: "inspect auth" },
|
||||||
|
undefined,
|
||||||
|
undefined,
|
||||||
|
{
|
||||||
|
cwd: "/repo",
|
||||||
|
modelRegistry: {
|
||||||
|
getAvailable: () => [{ provider: "anthropic", id: "claude-sonnet-4-5" }],
|
||||||
|
},
|
||||||
|
hasUI: false,
|
||||||
|
} as any,
|
||||||
|
);
|
||||||
|
|
||||||
|
assert.equal(didRun, false);
|
||||||
|
assert.equal(result.isError, true);
|
||||||
|
assert.match(result.content[0]?.type === "text" ? result.content[0].text : "", /top-level model/i);
|
||||||
|
});
|
||||||
|
|
||||||
|
test("single-mode subagent rejects models that are not currently available", async () => {
|
||||||
|
let didRun = false;
|
||||||
|
|
||||||
|
const tool = createSubagentTool({
|
||||||
|
discoverAgents: () => ({
|
||||||
|
agents: [{ name: "scout", description: "Scout", systemPrompt: "Scout prompt", source: "builtin" }],
|
||||||
|
projectAgentsDir: null,
|
||||||
|
}),
|
||||||
|
runSingleTask: async () => {
|
||||||
|
didRun = true;
|
||||||
|
throw new Error("should not run");
|
||||||
|
},
|
||||||
|
} as any);
|
||||||
|
|
||||||
|
const result = await tool.execute(
|
||||||
|
"tool-1",
|
||||||
|
{
|
||||||
|
agent: "scout",
|
||||||
|
task: "inspect auth",
|
||||||
|
model: "openai/gpt-5",
|
||||||
|
},
|
||||||
|
undefined,
|
||||||
|
undefined,
|
||||||
|
{
|
||||||
|
cwd: "/repo",
|
||||||
|
modelRegistry: {
|
||||||
|
getAvailable: () => [{ provider: "anthropic", id: "claude-sonnet-4-5" }],
|
||||||
|
},
|
||||||
|
hasUI: false,
|
||||||
|
} as any,
|
||||||
|
);
|
||||||
|
|
||||||
|
assert.equal(didRun, false);
|
||||||
|
assert.equal(result.isError, true);
|
||||||
|
assert.match(result.content[0]?.type === "text" ? result.content[0].text : "", /available models/i);
|
||||||
|
});
|
||||||
|
|
||||||
|
test("single-mode subagent asks before running a project-local agent", async () => {
|
||||||
|
const tool = createSubagentTool({
|
||||||
|
discoverAgents: () => ({
|
||||||
|
agents: [
|
||||||
|
{
|
||||||
|
name: "reviewer",
|
||||||
|
description: "Reviewer",
|
||||||
|
systemPrompt: "Review prompt",
|
||||||
|
source: "project",
|
||||||
|
},
|
||||||
|
],
|
||||||
|
projectAgentsDir: "/repo/.pi/agents",
|
||||||
|
}),
|
||||||
|
runSingleTask: async () => {
|
||||||
|
throw new Error("should not run");
|
||||||
|
},
|
||||||
|
} as any);
|
||||||
|
|
||||||
|
const result = await tool.execute(
|
||||||
|
"tool-1",
|
||||||
|
{
|
||||||
|
agent: "reviewer",
|
||||||
|
task: "review auth",
|
||||||
|
model: "anthropic/claude-sonnet-4-5",
|
||||||
|
agentScope: "both",
|
||||||
|
},
|
||||||
|
undefined,
|
||||||
|
undefined,
|
||||||
|
{
|
||||||
|
cwd: "/repo",
|
||||||
|
modelRegistry: {
|
||||||
|
getAvailable: () => [{ provider: "anthropic", id: "claude-sonnet-4-5" }],
|
||||||
|
},
|
||||||
|
hasUI: true,
|
||||||
|
ui: { confirm: async () => false },
|
||||||
|
} as any,
|
||||||
|
);
|
||||||
|
|
||||||
|
assert.equal(result.isError, true);
|
||||||
|
assert.match(result.content[0]?.type === "text" ? result.content[0].text : "", /not approved/);
|
||||||
|
});
|
||||||
336
.pi/agent/extensions/tmux-subagent/src/tool.ts
Normal file
336
.pi/agent/extensions/tmux-subagent/src/tool.ts
Normal file
@@ -0,0 +1,336 @@
|
|||||||
|
import { Text } from "@mariozechner/pi-tui";
|
||||||
|
import { discoverAgents } from "./agents.ts";
|
||||||
|
import {
|
||||||
|
listAvailableModelReferences,
|
||||||
|
normalizeAvailableModelReference,
|
||||||
|
resolveChildModel,
|
||||||
|
} from "./models.ts";
|
||||||
|
import {
|
||||||
|
SubagentParamsSchema,
|
||||||
|
type AgentScope,
|
||||||
|
type SubagentRunResult,
|
||||||
|
type SubagentToolDetails,
|
||||||
|
} from "./schema.ts";
|
||||||
|
|
||||||
|
const MAX_PARALLEL_TASKS = 8;
|
||||||
|
const MAX_CONCURRENCY = 4;
|
||||||
|
|
||||||
|
async function mapWithConcurrencyLimit<TIn, TOut>(
|
||||||
|
items: TIn[],
|
||||||
|
concurrency: number,
|
||||||
|
fn: (item: TIn, index: number) => Promise<TOut>,
|
||||||
|
): Promise<TOut[]> {
|
||||||
|
const limit = Math.max(1, Math.min(concurrency, items.length || 1));
|
||||||
|
const results = new Array<TOut>(items.length);
|
||||||
|
let nextIndex = 0;
|
||||||
|
|
||||||
|
await Promise.all(
|
||||||
|
Array.from({ length: limit }, async () => {
|
||||||
|
while (nextIndex < items.length) {
|
||||||
|
const index = nextIndex++;
|
||||||
|
results[index] = await fn(items[index], index);
|
||||||
|
}
|
||||||
|
}),
|
||||||
|
);
|
||||||
|
|
||||||
|
return results;
|
||||||
|
}
|
||||||
|
|
||||||
|
function isFailure(result: Pick<SubagentRunResult, "exitCode" | "stopReason">) {
|
||||||
|
return result.exitCode !== 0 || result.stopReason === "error" || result.stopReason === "aborted";
|
||||||
|
}
|
||||||
|
|
||||||
|
function makeDetails(
|
||||||
|
mode: "single" | "parallel" | "chain",
|
||||||
|
agentScope: AgentScope,
|
||||||
|
projectAgentsDir: string | null,
|
||||||
|
results: SubagentRunResult[],
|
||||||
|
): SubagentToolDetails {
|
||||||
|
return { mode, agentScope, projectAgentsDir, results };
|
||||||
|
}
|
||||||
|
|
||||||
|
function makeErrorResult(
|
||||||
|
text: string,
|
||||||
|
mode: "single" | "parallel" | "chain",
|
||||||
|
agentScope: AgentScope,
|
||||||
|
projectAgentsDir: string | null,
|
||||||
|
) {
|
||||||
|
return {
|
||||||
|
content: [{ type: "text" as const, text }],
|
||||||
|
details: makeDetails(mode, agentScope, projectAgentsDir, []),
|
||||||
|
isError: true,
|
||||||
|
};
|
||||||
|
}
|
||||||
|
|
||||||
|
export function createSubagentTool(deps: {
|
||||||
|
discoverAgents?: typeof discoverAgents;
|
||||||
|
listAvailableModelReferences?: typeof listAvailableModelReferences;
|
||||||
|
normalizeAvailableModelReference?: typeof normalizeAvailableModelReference;
|
||||||
|
parameters?: typeof SubagentParamsSchema;
|
||||||
|
resolveChildModel?: typeof resolveChildModel;
|
||||||
|
runSingleTask?: (input: {
|
||||||
|
cwd: string;
|
||||||
|
meta: Record<string, unknown>;
|
||||||
|
onEvent?: (event: any) => void;
|
||||||
|
}) => Promise<SubagentRunResult>;
|
||||||
|
} = {}) {
|
||||||
|
return {
|
||||||
|
name: "subagent",
|
||||||
|
label: "Subagent",
|
||||||
|
description: "Delegate tasks to specialized agents running in tmux panes.",
|
||||||
|
parameters: deps.parameters ?? SubagentParamsSchema,
|
||||||
|
async execute(_toolCallId: string, params: any, _signal: AbortSignal | undefined, onUpdate: any, ctx: any) {
|
||||||
|
const hasSingle = Boolean(params.agent && params.task);
|
||||||
|
const hasParallel = Boolean(params.tasks?.length);
|
||||||
|
const hasChain = Boolean(params.chain?.length);
|
||||||
|
const modeCount = Number(hasSingle) + Number(hasParallel) + Number(hasChain);
|
||||||
|
const mode = hasParallel ? "parallel" : hasChain ? "chain" : "single";
|
||||||
|
const agentScope = (params.agentScope ?? "user") as AgentScope;
|
||||||
|
|
||||||
|
if (modeCount !== 1) {
|
||||||
|
return makeErrorResult("Provide exactly one mode: single, parallel, or chain.", "single", agentScope, null);
|
||||||
|
}
|
||||||
|
|
||||||
|
const discovery = (deps.discoverAgents ?? discoverAgents)(ctx.cwd, { scope: agentScope });
|
||||||
|
const availableModelReferences = (deps.listAvailableModelReferences ?? listAvailableModelReferences)(ctx.modelRegistry);
|
||||||
|
const availableModelsText = availableModelReferences.join(", ") || "(none)";
|
||||||
|
const normalizeModelReference = (requestedModel?: string) =>
|
||||||
|
(deps.normalizeAvailableModelReference ?? normalizeAvailableModelReference)(requestedModel, availableModelReferences);
|
||||||
|
|
||||||
|
if (availableModelReferences.length === 0) {
|
||||||
|
return makeErrorResult(
|
||||||
|
"No available models are configured. Configure at least one model before using subagent.",
|
||||||
|
mode,
|
||||||
|
agentScope,
|
||||||
|
discovery.projectAgentsDir,
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
const topLevelModel = normalizeModelReference(params.model);
|
||||||
|
if (!topLevelModel) {
|
||||||
|
const message =
|
||||||
|
typeof params.model !== "string" || params.model.trim().length === 0
|
||||||
|
? `Subagent requires a top-level model chosen from the available models: ${availableModelsText}`
|
||||||
|
: `Invalid top-level model "${params.model}". Choose one of the available models: ${availableModelsText}`;
|
||||||
|
return makeErrorResult(message, mode, agentScope, discovery.projectAgentsDir);
|
||||||
|
}
|
||||||
|
params.model = topLevelModel;
|
||||||
|
|
||||||
|
for (const [index, task] of (params.tasks ?? []).entries()) {
|
||||||
|
if (task.model === undefined) continue;
|
||||||
|
|
||||||
|
const normalizedTaskModel = normalizeModelReference(task.model);
|
||||||
|
if (!normalizedTaskModel) {
|
||||||
|
return makeErrorResult(
|
||||||
|
`Invalid model for parallel task ${index + 1} (${task.agent}): "${task.model}". Choose one of the available models: ${availableModelsText}`,
|
||||||
|
mode,
|
||||||
|
agentScope,
|
||||||
|
discovery.projectAgentsDir,
|
||||||
|
);
|
||||||
|
}
|
||||||
|
task.model = normalizedTaskModel;
|
||||||
|
}
|
||||||
|
|
||||||
|
for (const [index, step] of (params.chain ?? []).entries()) {
|
||||||
|
if (step.model === undefined) continue;
|
||||||
|
|
||||||
|
const normalizedStepModel = normalizeModelReference(step.model);
|
||||||
|
if (!normalizedStepModel) {
|
||||||
|
return makeErrorResult(
|
||||||
|
`Invalid model for chain step ${index + 1} (${step.agent}): "${step.model}". Choose one of the available models: ${availableModelsText}`,
|
||||||
|
mode,
|
||||||
|
agentScope,
|
||||||
|
discovery.projectAgentsDir,
|
||||||
|
);
|
||||||
|
}
|
||||||
|
step.model = normalizedStepModel;
|
||||||
|
}
|
||||||
|
|
||||||
|
const requestedAgentNames = [
|
||||||
|
...(hasSingle ? [params.agent] : []),
|
||||||
|
...((params.tasks ?? []).map((task: any) => task.agent)),
|
||||||
|
...((params.chain ?? []).map((step: any) => step.agent)),
|
||||||
|
];
|
||||||
|
const projectAgents = requestedAgentNames
|
||||||
|
.map((name) => discovery.agents.find((candidate) => candidate.name === name))
|
||||||
|
.filter((agent): agent is NonNullable<typeof agent> => Boolean(agent && agent.source === "project"));
|
||||||
|
|
||||||
|
if (projectAgents.length > 0 && (params.confirmProjectAgents ?? true) && ctx.hasUI) {
|
||||||
|
const ok = await ctx.ui.confirm(
|
||||||
|
"Run project-local agents?",
|
||||||
|
`Agents: ${projectAgents.map((agent) => agent.name).join(", ")}\nSource: ${
|
||||||
|
discovery.projectAgentsDir ?? "(unknown)"
|
||||||
|
}`,
|
||||||
|
);
|
||||||
|
if (!ok) {
|
||||||
|
return makeErrorResult(
|
||||||
|
"Canceled: project-local agents not approved.",
|
||||||
|
mode,
|
||||||
|
agentScope,
|
||||||
|
discovery.projectAgentsDir,
|
||||||
|
);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
const resolveAgent = (name: string) => {
|
||||||
|
const agent = discovery.agents.find((candidate) => candidate.name === name);
|
||||||
|
if (!agent) throw new Error(`Unknown agent: ${name}`);
|
||||||
|
return agent;
|
||||||
|
};
|
||||||
|
|
||||||
|
const runTask = async (input: {
|
||||||
|
agentName: string;
|
||||||
|
task: string;
|
||||||
|
cwd?: string;
|
||||||
|
taskModel?: string;
|
||||||
|
taskIndex?: number;
|
||||||
|
step?: number;
|
||||||
|
mode: "single" | "parallel" | "chain";
|
||||||
|
}) => {
|
||||||
|
const agent = resolveAgent(input.agentName);
|
||||||
|
const model = (deps.resolveChildModel ?? resolveChildModel)({
|
||||||
|
taskModel: input.taskModel,
|
||||||
|
topLevelModel: params.model,
|
||||||
|
});
|
||||||
|
|
||||||
|
return deps.runSingleTask?.({
|
||||||
|
cwd: input.cwd ?? ctx.cwd,
|
||||||
|
onEvent(event) {
|
||||||
|
onUpdate?.({
|
||||||
|
content: [{ type: "text", text: `Running ${input.agentName}: ${event.type}` }],
|
||||||
|
details: makeDetails(input.mode, agentScope, discovery.projectAgentsDir, []),
|
||||||
|
});
|
||||||
|
},
|
||||||
|
meta: {
|
||||||
|
mode: input.mode,
|
||||||
|
taskIndex: input.taskIndex,
|
||||||
|
step: input.step,
|
||||||
|
agent: agent.name,
|
||||||
|
agentSource: agent.source,
|
||||||
|
task: input.task,
|
||||||
|
cwd: input.cwd ?? ctx.cwd,
|
||||||
|
requestedModel: model.requestedModel,
|
||||||
|
resolvedModel: model.resolvedModel,
|
||||||
|
systemPrompt: agent.systemPrompt,
|
||||||
|
tools: agent.tools,
|
||||||
|
},
|
||||||
|
}) as Promise<SubagentRunResult>;
|
||||||
|
};
|
||||||
|
|
||||||
|
if (hasSingle) {
|
||||||
|
try {
|
||||||
|
const result = await runTask({
|
||||||
|
agentName: params.agent,
|
||||||
|
task: params.task,
|
||||||
|
cwd: params.cwd,
|
||||||
|
mode: "single",
|
||||||
|
});
|
||||||
|
|
||||||
|
return {
|
||||||
|
content: [{ type: "text" as const, text: result.finalText }],
|
||||||
|
details: makeDetails("single", agentScope, discovery.projectAgentsDir, [result]),
|
||||||
|
isError: isFailure(result),
|
||||||
|
};
|
||||||
|
} catch (error) {
|
||||||
|
return {
|
||||||
|
content: [{ type: "text" as const, text: (error as Error).message }],
|
||||||
|
details: makeDetails("single", agentScope, discovery.projectAgentsDir, []),
|
||||||
|
isError: true,
|
||||||
|
};
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
if (hasParallel) {
|
||||||
|
if (params.tasks.length > MAX_PARALLEL_TASKS) {
|
||||||
|
return {
|
||||||
|
content: [
|
||||||
|
{
|
||||||
|
type: "text" as const,
|
||||||
|
text: `Too many parallel tasks (${params.tasks.length}). Max is ${MAX_PARALLEL_TASKS}.`,
|
||||||
|
},
|
||||||
|
],
|
||||||
|
details: makeDetails("parallel", agentScope, discovery.projectAgentsDir, []),
|
||||||
|
isError: true,
|
||||||
|
};
|
||||||
|
}
|
||||||
|
|
||||||
|
const liveResults: SubagentRunResult[] = [];
|
||||||
|
const results = await mapWithConcurrencyLimit(params.tasks, MAX_CONCURRENCY, async (task: any, index) => {
|
||||||
|
const result = await runTask({
|
||||||
|
agentName: task.agent,
|
||||||
|
task: task.task,
|
||||||
|
cwd: task.cwd,
|
||||||
|
taskModel: task.model,
|
||||||
|
taskIndex: index,
|
||||||
|
mode: "parallel",
|
||||||
|
});
|
||||||
|
liveResults[index] = result;
|
||||||
|
onUpdate?.({
|
||||||
|
content: [{ type: "text", text: `Parallel: ${liveResults.filter(Boolean).length}/${params.tasks.length} finished` }],
|
||||||
|
details: makeDetails("parallel", agentScope, discovery.projectAgentsDir, liveResults.filter(Boolean)),
|
||||||
|
});
|
||||||
|
return result;
|
||||||
|
});
|
||||||
|
|
||||||
|
const successCount = results.filter((result) => !isFailure(result)).length;
|
||||||
|
const summary = results
|
||||||
|
.map((result) => `[${result.agent}] ${isFailure(result) ? "failed" : "completed"}: ${result.finalText || "(no output)"}`)
|
||||||
|
.join("\n\n");
|
||||||
|
|
||||||
|
return {
|
||||||
|
content: [{ type: "text" as const, text: `Parallel: ${successCount}/${results.length} succeeded\n\n${summary}` }],
|
||||||
|
details: makeDetails("parallel", agentScope, discovery.projectAgentsDir, results),
|
||||||
|
isError: successCount !== results.length,
|
||||||
|
};
|
||||||
|
}
|
||||||
|
|
||||||
|
const results: SubagentRunResult[] = [];
|
||||||
|
let previous = "";
|
||||||
|
for (let index = 0; index < params.chain.length; index += 1) {
|
||||||
|
const item = params.chain[index];
|
||||||
|
const task = item.task.replaceAll("{previous}", previous);
|
||||||
|
const result = await runTask({
|
||||||
|
agentName: item.agent,
|
||||||
|
task,
|
||||||
|
cwd: item.cwd,
|
||||||
|
taskModel: item.model,
|
||||||
|
step: index + 1,
|
||||||
|
mode: "chain",
|
||||||
|
});
|
||||||
|
onUpdate?.({
|
||||||
|
content: [{ type: "text", text: `Chain: completed step ${index + 1}/${params.chain.length}` }],
|
||||||
|
details: makeDetails("chain", agentScope, discovery.projectAgentsDir, [...results, result]),
|
||||||
|
});
|
||||||
|
results.push(result);
|
||||||
|
if (isFailure(result)) {
|
||||||
|
return {
|
||||||
|
content: [
|
||||||
|
{
|
||||||
|
type: "text" as const,
|
||||||
|
text: `Chain stopped at step ${index + 1} (${item.agent}): ${result.finalText || result.stopReason || "failed"}`,
|
||||||
|
},
|
||||||
|
],
|
||||||
|
details: makeDetails("chain", agentScope, discovery.projectAgentsDir, results),
|
||||||
|
isError: true,
|
||||||
|
};
|
||||||
|
}
|
||||||
|
previous = result.finalText;
|
||||||
|
}
|
||||||
|
|
||||||
|
const finalResult = results[results.length - 1];
|
||||||
|
return {
|
||||||
|
content: [{ type: "text" as const, text: finalResult?.finalText ?? "" }],
|
||||||
|
details: makeDetails("chain", agentScope, discovery.projectAgentsDir, results),
|
||||||
|
};
|
||||||
|
},
|
||||||
|
renderCall(args: any) {
|
||||||
|
if (args.tasks?.length) return new Text(`subagent parallel (${args.tasks.length} tasks)`, 0, 0);
|
||||||
|
if (args.chain?.length) return new Text(`subagent chain (${args.chain.length} steps)`, 0, 0);
|
||||||
|
return new Text(`subagent ${args.agent ?? ""}`.trim(), 0, 0);
|
||||||
|
},
|
||||||
|
renderResult(result: { content: Array<{ type: string; text?: string }> }) {
|
||||||
|
const first = result.content[0];
|
||||||
|
return new Text(first?.type === "text" ? first.text ?? "" : "", 0, 0);
|
||||||
|
},
|
||||||
|
};
|
||||||
|
}
|
||||||
214
.pi/agent/extensions/tmux-subagent/src/wrapper/cli.mjs
Normal file
214
.pi/agent/extensions/tmux-subagent/src/wrapper/cli.mjs
Normal file
@@ -0,0 +1,214 @@
|
|||||||
|
import { appendFile, readFile, writeFile } from "node:fs/promises";
|
||||||
|
import { spawn } from "node:child_process";
|
||||||
|
import { normalizePiEvent } from "./normalize.mjs";
|
||||||
|
import { renderHeader, renderEventLine } from "./render.mjs";
|
||||||
|
|
||||||
|
async function appendJsonLine(path, value) {
|
||||||
|
await appendBestEffort(path, `${JSON.stringify(value)}\n`);
|
||||||
|
}
|
||||||
|
|
||||||
|
async function appendBestEffort(path, text) {
|
||||||
|
try {
|
||||||
|
await appendFile(path, text, "utf8");
|
||||||
|
} catch {
|
||||||
|
// Best-effort artifact logging should never prevent result.json from being written.
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
function makeResult(meta, startedAt, input = {}) {
|
||||||
|
const errorText = typeof input.errorMessage === "string" ? input.errorMessage.trim() : "";
|
||||||
|
const exitCode = typeof input.exitCode === "number" ? input.exitCode : 1;
|
||||||
|
return {
|
||||||
|
runId: meta.runId,
|
||||||
|
mode: meta.mode,
|
||||||
|
taskIndex: meta.taskIndex,
|
||||||
|
step: meta.step,
|
||||||
|
agent: meta.agent,
|
||||||
|
agentSource: meta.agentSource,
|
||||||
|
task: meta.task,
|
||||||
|
cwd: meta.cwd,
|
||||||
|
requestedModel: meta.requestedModel,
|
||||||
|
resolvedModel: input.resolvedModel ?? meta.resolvedModel,
|
||||||
|
sessionPath: meta.sessionPath,
|
||||||
|
startedAt,
|
||||||
|
finishedAt: new Date().toISOString(),
|
||||||
|
exitCode,
|
||||||
|
stopReason: input.stopReason ?? (exitCode === 0 ? undefined : "error"),
|
||||||
|
finalText: input.finalText ?? "",
|
||||||
|
usage: input.usage,
|
||||||
|
stdoutPath: meta.stdoutPath,
|
||||||
|
stderrPath: meta.stderrPath,
|
||||||
|
resultPath: meta.resultPath,
|
||||||
|
eventsPath: meta.eventsPath,
|
||||||
|
transcriptPath: meta.transcriptPath,
|
||||||
|
errorMessage: errorText || undefined,
|
||||||
|
};
|
||||||
|
}
|
||||||
|
|
||||||
|
async function runWrapper(meta, startedAt) {
|
||||||
|
const header = renderHeader(meta);
|
||||||
|
await appendBestEffort(meta.transcriptPath, `${header}\n`);
|
||||||
|
console.log(header);
|
||||||
|
|
||||||
|
const effectiveModel =
|
||||||
|
typeof meta.resolvedModel === "string" && meta.resolvedModel.length > 0
|
||||||
|
? meta.resolvedModel
|
||||||
|
: meta.requestedModel;
|
||||||
|
|
||||||
|
const args = ["--mode", "json", "--session", meta.sessionPath];
|
||||||
|
if (effectiveModel) args.push("--model", effectiveModel);
|
||||||
|
if (Array.isArray(meta.tools) && meta.tools.length > 0) args.push("--tools", meta.tools.join(","));
|
||||||
|
if (meta.systemPromptPath) args.push("--append-system-prompt", meta.systemPromptPath);
|
||||||
|
args.push(meta.task);
|
||||||
|
|
||||||
|
let finalText = "";
|
||||||
|
let resolvedModel = meta.resolvedModel;
|
||||||
|
let stopReason;
|
||||||
|
let usage = undefined;
|
||||||
|
let stdoutBuffer = "";
|
||||||
|
let stderrText = "";
|
||||||
|
let spawnError;
|
||||||
|
let queue = Promise.resolve();
|
||||||
|
|
||||||
|
const enqueue = (work) => {
|
||||||
|
queue = queue.then(work, work);
|
||||||
|
return queue;
|
||||||
|
};
|
||||||
|
|
||||||
|
const handleStdoutLine = async (line) => {
|
||||||
|
if (!line.trim()) return;
|
||||||
|
|
||||||
|
let parsed;
|
||||||
|
try {
|
||||||
|
parsed = JSON.parse(line);
|
||||||
|
} catch {
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
|
||||||
|
const normalized = normalizePiEvent(parsed);
|
||||||
|
if (!normalized) return;
|
||||||
|
|
||||||
|
await appendJsonLine(meta.eventsPath, normalized);
|
||||||
|
const rendered = renderEventLine(normalized);
|
||||||
|
await appendBestEffort(meta.transcriptPath, `${rendered}\n`);
|
||||||
|
console.log(rendered);
|
||||||
|
|
||||||
|
if (normalized.type === "assistant_text") {
|
||||||
|
finalText = normalized.text;
|
||||||
|
resolvedModel = normalized.model ?? resolvedModel;
|
||||||
|
stopReason = normalized.stopReason ?? stopReason;
|
||||||
|
usage = normalized.usage ?? usage;
|
||||||
|
}
|
||||||
|
};
|
||||||
|
|
||||||
|
const childEnv = { ...process.env };
|
||||||
|
// Ensure the copilot initiator flag is not accidentally inherited from the parent
|
||||||
|
// environment; set it only for github-copilot models.
|
||||||
|
delete childEnv.PI_TMUX_SUBAGENT_GITHUB_COPILOT_INITIATOR;
|
||||||
|
// Mark every child run as a nested tmux subagent so it cannot spawn further subagents.
|
||||||
|
childEnv.PI_TMUX_SUBAGENT_CHILD = "1";
|
||||||
|
|
||||||
|
if (typeof effectiveModel === "string" && effectiveModel.startsWith("github-copilot/")) {
|
||||||
|
childEnv.PI_TMUX_SUBAGENT_GITHUB_COPILOT_INITIATOR = "agent";
|
||||||
|
}
|
||||||
|
|
||||||
|
const child = spawn("pi", args, {
|
||||||
|
cwd: meta.cwd,
|
||||||
|
env: childEnv,
|
||||||
|
stdio: ["ignore", "pipe", "pipe"],
|
||||||
|
});
|
||||||
|
|
||||||
|
child.stdout.on("data", (chunk) => {
|
||||||
|
const text = chunk.toString();
|
||||||
|
enqueue(async () => {
|
||||||
|
stdoutBuffer += text;
|
||||||
|
await appendBestEffort(meta.stdoutPath, text);
|
||||||
|
|
||||||
|
const lines = stdoutBuffer.split("\n");
|
||||||
|
stdoutBuffer = lines.pop() ?? "";
|
||||||
|
|
||||||
|
for (const line of lines) {
|
||||||
|
await handleStdoutLine(line);
|
||||||
|
}
|
||||||
|
});
|
||||||
|
});
|
||||||
|
|
||||||
|
child.stderr.on("data", (chunk) => {
|
||||||
|
const text = chunk.toString();
|
||||||
|
enqueue(async () => {
|
||||||
|
stderrText += text;
|
||||||
|
await appendBestEffort(meta.stderrPath, text);
|
||||||
|
});
|
||||||
|
});
|
||||||
|
|
||||||
|
const exitCode = await new Promise((resolve) => {
|
||||||
|
let done = false;
|
||||||
|
const finish = (code) => {
|
||||||
|
if (done) return;
|
||||||
|
done = true;
|
||||||
|
resolve(code);
|
||||||
|
};
|
||||||
|
|
||||||
|
child.on("error", (error) => {
|
||||||
|
spawnError = error;
|
||||||
|
finish(1);
|
||||||
|
});
|
||||||
|
|
||||||
|
child.on("close", (code) => {
|
||||||
|
finish(code ?? (spawnError ? 1 : 0));
|
||||||
|
});
|
||||||
|
});
|
||||||
|
|
||||||
|
await queue;
|
||||||
|
|
||||||
|
if (stdoutBuffer.trim()) {
|
||||||
|
await handleStdoutLine(stdoutBuffer);
|
||||||
|
stdoutBuffer = "";
|
||||||
|
}
|
||||||
|
|
||||||
|
if (spawnError) {
|
||||||
|
const message = spawnError instanceof Error ? spawnError.stack ?? spawnError.message : String(spawnError);
|
||||||
|
if (!stderrText.trim()) {
|
||||||
|
stderrText = message;
|
||||||
|
await appendBestEffort(meta.stderrPath, `${message}\n`);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
return makeResult(meta, startedAt, {
|
||||||
|
exitCode,
|
||||||
|
stopReason,
|
||||||
|
finalText,
|
||||||
|
usage,
|
||||||
|
resolvedModel,
|
||||||
|
errorMessage: stderrText,
|
||||||
|
});
|
||||||
|
}
|
||||||
|
|
||||||
|
async function main() {
|
||||||
|
const metaPath = process.argv[2];
|
||||||
|
if (!metaPath) throw new Error("Expected meta.json path as argv[2]");
|
||||||
|
|
||||||
|
const meta = JSON.parse(await readFile(metaPath, "utf8"));
|
||||||
|
const startedAt = meta.startedAt ?? new Date().toISOString();
|
||||||
|
|
||||||
|
let result;
|
||||||
|
try {
|
||||||
|
result = await runWrapper(meta, startedAt);
|
||||||
|
} catch (error) {
|
||||||
|
const message = error instanceof Error ? error.stack ?? error.message : String(error);
|
||||||
|
await appendBestEffort(meta.stderrPath, `${message}\n`);
|
||||||
|
result = makeResult(meta, startedAt, {
|
||||||
|
exitCode: 1,
|
||||||
|
stopReason: "error",
|
||||||
|
errorMessage: message,
|
||||||
|
});
|
||||||
|
}
|
||||||
|
|
||||||
|
await writeFile(meta.resultPath, JSON.stringify(result, null, 2), "utf8");
|
||||||
|
if (result.exitCode !== 0) process.exitCode = result.exitCode;
|
||||||
|
}
|
||||||
|
|
||||||
|
main().catch((error) => {
|
||||||
|
console.error(error instanceof Error ? error.stack : String(error));
|
||||||
|
process.exitCode = 1;
|
||||||
|
});
|
||||||
192
.pi/agent/extensions/tmux-subagent/src/wrapper/cli.test.ts
Normal file
192
.pi/agent/extensions/tmux-subagent/src/wrapper/cli.test.ts
Normal file
@@ -0,0 +1,192 @@
|
|||||||
|
import test from "node:test";
|
||||||
|
import assert from "node:assert/strict";
|
||||||
|
import { chmod, mkdtemp, readFile, writeFile } from "node:fs/promises";
|
||||||
|
import { spawn } from "node:child_process";
|
||||||
|
import { tmpdir } from "node:os";
|
||||||
|
import { dirname, join } from "node:path";
|
||||||
|
import { fileURLToPath } from "node:url";
|
||||||
|
|
||||||
|
function waitForExit(child: ReturnType<typeof spawn>, timeoutMs = 1500): Promise<number> {
|
||||||
|
return new Promise((resolve, reject) => {
|
||||||
|
const timeout = setTimeout(() => {
|
||||||
|
child.kill("SIGKILL");
|
||||||
|
reject(new Error(`wrapper did not exit within ${timeoutMs}ms`));
|
||||||
|
}, timeoutMs);
|
||||||
|
|
||||||
|
child.on("error", (error) => {
|
||||||
|
clearTimeout(timeout);
|
||||||
|
reject(error);
|
||||||
|
});
|
||||||
|
|
||||||
|
child.on("close", (code) => {
|
||||||
|
clearTimeout(timeout);
|
||||||
|
resolve(code ?? 0);
|
||||||
|
});
|
||||||
|
});
|
||||||
|
}
|
||||||
|
|
||||||
|
async function runWrapperWithFakePi(requestedModel: string, resolvedModel?: string) {
|
||||||
|
const dir = await mkdtemp(join(tmpdir(), "tmux-subagent-wrapper-"));
|
||||||
|
const metaPath = join(dir, "meta.json");
|
||||||
|
const resultPath = join(dir, "result.json");
|
||||||
|
const capturePath = join(dir, "capture.json");
|
||||||
|
const piPath = join(dir, "pi");
|
||||||
|
|
||||||
|
// The fake `pi` is a small Node script that writes a JSON capture file
|
||||||
|
// including relevant PI_* environment variables and the argv it received.
|
||||||
|
const resolved = typeof resolvedModel === "string" ? resolvedModel : requestedModel;
|
||||||
|
await writeFile(
|
||||||
|
piPath,
|
||||||
|
[
|
||||||
|
`#!${process.execPath}`,
|
||||||
|
"const fs = require('fs');",
|
||||||
|
`const capturePath = ${JSON.stringify(capturePath)};`,
|
||||||
|
"const obj = {",
|
||||||
|
" PI_TMUX_SUBAGENT_GITHUB_COPILOT_INITIATOR: process.env.PI_TMUX_SUBAGENT_GITHUB_COPILOT_INITIATOR || '',",
|
||||||
|
" PI_TMUX_SUBAGENT_CHILD: process.env.PI_TMUX_SUBAGENT_CHILD || '',",
|
||||||
|
" argv: process.argv.slice(2)",
|
||||||
|
"};",
|
||||||
|
"fs.writeFileSync(capturePath, JSON.stringify(obj), 'utf8');",
|
||||||
|
"console.log(JSON.stringify({type:'message_end',message:{role:'assistant',content:[{type:'text',text:'done'}],model:'github-copilot/gpt-4o',stopReason:'stop'}}));",
|
||||||
|
].join("\n"),
|
||||||
|
"utf8",
|
||||||
|
);
|
||||||
|
await chmod(piPath, 0o755);
|
||||||
|
|
||||||
|
await writeFile(
|
||||||
|
metaPath,
|
||||||
|
JSON.stringify(
|
||||||
|
{
|
||||||
|
runId: "run-1",
|
||||||
|
mode: "single",
|
||||||
|
agent: "scout",
|
||||||
|
agentSource: "builtin",
|
||||||
|
task: "inspect auth",
|
||||||
|
cwd: dir,
|
||||||
|
requestedModel,
|
||||||
|
resolvedModel: resolved,
|
||||||
|
startedAt: "2026-04-09T00:00:00.000Z",
|
||||||
|
sessionPath: join(dir, "child-session.jsonl"),
|
||||||
|
eventsPath: join(dir, "events.jsonl"),
|
||||||
|
resultPath,
|
||||||
|
stdoutPath: join(dir, "stdout.log"),
|
||||||
|
stderrPath: join(dir, "stderr.log"),
|
||||||
|
transcriptPath: join(dir, "transcript.log"),
|
||||||
|
systemPromptPath: join(dir, "system-prompt.md"),
|
||||||
|
},
|
||||||
|
null,
|
||||||
|
2,
|
||||||
|
),
|
||||||
|
"utf8",
|
||||||
|
);
|
||||||
|
|
||||||
|
const wrapperPath = join(dirname(fileURLToPath(import.meta.url)), "cli.mjs");
|
||||||
|
const child = spawn(process.execPath, [wrapperPath, metaPath], {
|
||||||
|
env: {
|
||||||
|
...process.env,
|
||||||
|
PATH: dir,
|
||||||
|
},
|
||||||
|
stdio: ["ignore", "pipe", "pipe"],
|
||||||
|
});
|
||||||
|
|
||||||
|
const exitCode = await waitForExit(child);
|
||||||
|
assert.equal(exitCode, 0);
|
||||||
|
|
||||||
|
const captureJson = JSON.parse(await readFile(capturePath, "utf8"));
|
||||||
|
return { flags: captureJson };
|
||||||
|
}
|
||||||
|
|
||||||
|
// Dedicated tests: every child run must have PI_TMUX_SUBAGENT_CHILD=1
|
||||||
|
test("wrapper marks github-copilot child run as a tmux subagent child", async () => {
|
||||||
|
const captured = await runWrapperWithFakePi("github-copilot/gpt-4o");
|
||||||
|
assert.equal(captured.flags.PI_TMUX_SUBAGENT_CHILD, "1");
|
||||||
|
});
|
||||||
|
|
||||||
|
test("wrapper marks anthropic child run as a tmux subagent child", async () => {
|
||||||
|
const captured = await runWrapperWithFakePi("anthropic/claude-sonnet-4-5");
|
||||||
|
assert.equal(captured.flags.PI_TMUX_SUBAGENT_CHILD, "1");
|
||||||
|
});
|
||||||
|
|
||||||
|
test("wrapper marks github-copilot child runs as agent-initiated", async () => {
|
||||||
|
const captured = await runWrapperWithFakePi("github-copilot/gpt-4o");
|
||||||
|
assert.equal(captured.flags.PI_TMUX_SUBAGENT_GITHUB_COPILOT_INITIATOR, "agent");
|
||||||
|
assert.equal(captured.flags.PI_TMUX_SUBAGENT_CHILD, "1");
|
||||||
|
});
|
||||||
|
|
||||||
|
test("wrapper leaves non-copilot child runs unchanged", async () => {
|
||||||
|
const captured = await runWrapperWithFakePi("anthropic/claude-sonnet-4-5");
|
||||||
|
// The wrapper should not inject the copilot initiator for non-copilot models.
|
||||||
|
assert.equal(captured.flags.PI_TMUX_SUBAGENT_GITHUB_COPILOT_INITIATOR, "");
|
||||||
|
assert.equal(captured.flags.PI_TMUX_SUBAGENT_CHILD, "1");
|
||||||
|
});
|
||||||
|
|
||||||
|
// Regression test: ensure when requestedModel and resolvedModel differ, the
|
||||||
|
// wrapper uses the same effective model for the child --model arg and the
|
||||||
|
// copilot initiator env flag.
|
||||||
|
test("wrapper uses effective model for both argv and env when requested/resolved differ", async () => {
|
||||||
|
const requested = "anthropic/claude-sonnet-4-5";
|
||||||
|
const resolved = "github-copilot/gpt-4o";
|
||||||
|
|
||||||
|
const captured = await runWrapperWithFakePi(requested, resolved);
|
||||||
|
|
||||||
|
// The effective model should be the resolved model in this case.
|
||||||
|
assert.equal(captured.flags.PI_TMUX_SUBAGENT_GITHUB_COPILOT_INITIATOR, "agent");
|
||||||
|
assert.equal(captured.flags.PI_TMUX_SUBAGENT_CHILD, "1");
|
||||||
|
|
||||||
|
// Verify the child argv contains the effective model after a --model flag.
|
||||||
|
const argv = captured.flags.argv;
|
||||||
|
const modelIndex = argv.indexOf("--model");
|
||||||
|
assert.ok(modelIndex >= 0, "expected --model in argv");
|
||||||
|
assert.equal(argv[modelIndex + 1], resolved);
|
||||||
|
});
|
||||||
|
|
||||||
|
test("wrapper exits and writes result.json when the pi child cannot be spawned", async () => {
|
||||||
|
const dir = await mkdtemp(join(tmpdir(), "tmux-subagent-wrapper-"));
|
||||||
|
const metaPath = join(dir, "meta.json");
|
||||||
|
const resultPath = join(dir, "result.json");
|
||||||
|
|
||||||
|
await writeFile(
|
||||||
|
metaPath,
|
||||||
|
JSON.stringify(
|
||||||
|
{
|
||||||
|
runId: "run-1",
|
||||||
|
mode: "single",
|
||||||
|
agent: "scout",
|
||||||
|
agentSource: "builtin",
|
||||||
|
task: "inspect auth",
|
||||||
|
cwd: dir,
|
||||||
|
requestedModel: "anthropic/claude-sonnet-4-5",
|
||||||
|
resolvedModel: "anthropic/claude-sonnet-4-5",
|
||||||
|
startedAt: "2026-04-09T00:00:00.000Z",
|
||||||
|
sessionPath: join(dir, "child-session.jsonl"),
|
||||||
|
eventsPath: join(dir, "events.jsonl"),
|
||||||
|
resultPath,
|
||||||
|
stdoutPath: join(dir, "stdout.log"),
|
||||||
|
stderrPath: join(dir, "stderr.log"),
|
||||||
|
transcriptPath: join(dir, "transcript.log"),
|
||||||
|
systemPromptPath: join(dir, "system-prompt.md"),
|
||||||
|
},
|
||||||
|
null,
|
||||||
|
2,
|
||||||
|
),
|
||||||
|
"utf8",
|
||||||
|
);
|
||||||
|
|
||||||
|
const wrapperPath = join(dirname(fileURLToPath(import.meta.url)), "cli.mjs");
|
||||||
|
const child = spawn(process.execPath, [wrapperPath, metaPath], {
|
||||||
|
env: {
|
||||||
|
...process.env,
|
||||||
|
PATH: dir,
|
||||||
|
},
|
||||||
|
stdio: ["ignore", "pipe", "pipe"],
|
||||||
|
});
|
||||||
|
|
||||||
|
const exitCode = await waitForExit(child);
|
||||||
|
assert.equal(exitCode, 1);
|
||||||
|
|
||||||
|
const result = JSON.parse(await readFile(resultPath, "utf8"));
|
||||||
|
assert.equal(result.runId, "run-1");
|
||||||
|
assert.equal(result.agent, "scout");
|
||||||
|
assert.equal(result.exitCode, 1);
|
||||||
|
assert.match(result.errorMessage ?? "", /ENOENT|not found|spawn pi/i);
|
||||||
|
});
|
||||||
35
.pi/agent/extensions/tmux-subagent/src/wrapper/normalize.mjs
Normal file
35
.pi/agent/extensions/tmux-subagent/src/wrapper/normalize.mjs
Normal file
@@ -0,0 +1,35 @@
|
|||||||
|
export function normalizePiEvent(event) {
|
||||||
|
if (event?.type === "tool_execution_start") {
|
||||||
|
return {
|
||||||
|
type: "tool_call",
|
||||||
|
toolName: event.toolName,
|
||||||
|
args: event.args ?? {},
|
||||||
|
};
|
||||||
|
}
|
||||||
|
|
||||||
|
if (event?.type === "message_end" && event.message?.role === "assistant") {
|
||||||
|
const text = (event.message.content ?? [])
|
||||||
|
.filter((part) => part.type === "text")
|
||||||
|
.map((part) => part.text)
|
||||||
|
.join("\n")
|
||||||
|
.trim();
|
||||||
|
|
||||||
|
return {
|
||||||
|
type: "assistant_text",
|
||||||
|
text,
|
||||||
|
model: event.message.model,
|
||||||
|
stopReason: event.message.stopReason,
|
||||||
|
usage: event.message.usage,
|
||||||
|
};
|
||||||
|
}
|
||||||
|
|
||||||
|
if (event?.type === "tool_execution_end") {
|
||||||
|
return {
|
||||||
|
type: "tool_result",
|
||||||
|
toolName: event.toolName,
|
||||||
|
isError: Boolean(event.isError),
|
||||||
|
};
|
||||||
|
}
|
||||||
|
|
||||||
|
return null;
|
||||||
|
}
|
||||||
@@ -0,0 +1,38 @@
|
|||||||
|
import test from "node:test";
|
||||||
|
import assert from "node:assert/strict";
|
||||||
|
import { normalizePiEvent } from "./normalize.mjs";
|
||||||
|
|
||||||
|
test("normalizePiEvent converts tool start events into protocol tool-call records", () => {
|
||||||
|
const normalized = normalizePiEvent({
|
||||||
|
type: "tool_execution_start",
|
||||||
|
toolName: "read",
|
||||||
|
args: { path: "src/app.ts", offset: 1, limit: 20 },
|
||||||
|
});
|
||||||
|
|
||||||
|
assert.deepEqual(normalized, {
|
||||||
|
type: "tool_call",
|
||||||
|
toolName: "read",
|
||||||
|
args: { path: "src/app.ts", offset: 1, limit: 20 },
|
||||||
|
});
|
||||||
|
});
|
||||||
|
|
||||||
|
test("normalizePiEvent converts assistant message_end into a final-text record", () => {
|
||||||
|
const normalized = normalizePiEvent({
|
||||||
|
type: "message_end",
|
||||||
|
message: {
|
||||||
|
role: "assistant",
|
||||||
|
model: "anthropic/claude-sonnet-4-5",
|
||||||
|
stopReason: "end_turn",
|
||||||
|
content: [{ type: "text", text: "Final answer" }],
|
||||||
|
usage: { input: 10, output: 5, totalTokens: 15, cost: { total: 0.001 } },
|
||||||
|
},
|
||||||
|
});
|
||||||
|
|
||||||
|
assert.deepEqual(normalized, {
|
||||||
|
type: "assistant_text",
|
||||||
|
text: "Final answer",
|
||||||
|
model: "anthropic/claude-sonnet-4-5",
|
||||||
|
stopReason: "end_turn",
|
||||||
|
usage: { input: 10, output: 5, totalTokens: 15, cost: { total: 0.001 } },
|
||||||
|
});
|
||||||
|
});
|
||||||
33
.pi/agent/extensions/tmux-subagent/src/wrapper/render.mjs
Normal file
33
.pi/agent/extensions/tmux-subagent/src/wrapper/render.mjs
Normal file
@@ -0,0 +1,33 @@
|
|||||||
|
function shortenCommand(command) {
|
||||||
|
return command.length > 100 ? `${command.slice(0, 100)}…` : command;
|
||||||
|
}
|
||||||
|
|
||||||
|
export function renderHeader(meta) {
|
||||||
|
return [
|
||||||
|
"=== tmux subagent ===",
|
||||||
|
`Agent: ${meta.agent}`,
|
||||||
|
`Task: ${meta.task}`,
|
||||||
|
`CWD: ${meta.cwd}`,
|
||||||
|
`Requested model: ${meta.requestedModel ?? "(default)"}`,
|
||||||
|
`Resolved model: ${meta.resolvedModel ?? "(pending)"}`,
|
||||||
|
`Session: ${meta.sessionPath}`,
|
||||||
|
"---------------------",
|
||||||
|
].join("\n");
|
||||||
|
}
|
||||||
|
|
||||||
|
export function renderEventLine(event) {
|
||||||
|
if (event.type === "tool_call") {
|
||||||
|
if (event.toolName === "bash") return `$ ${shortenCommand(event.args.command ?? "")}`;
|
||||||
|
return `→ ${event.toolName} ${JSON.stringify(event.args)}`;
|
||||||
|
}
|
||||||
|
|
||||||
|
if (event.type === "tool_result") {
|
||||||
|
return event.isError ? `✗ ${event.toolName} failed` : `✓ ${event.toolName} done`;
|
||||||
|
}
|
||||||
|
|
||||||
|
if (event.type === "assistant_text") {
|
||||||
|
return event.text || "(no assistant text)";
|
||||||
|
}
|
||||||
|
|
||||||
|
return JSON.stringify(event);
|
||||||
|
}
|
||||||
@@ -0,0 +1,28 @@
|
|||||||
|
import test from "node:test";
|
||||||
|
import assert from "node:assert/strict";
|
||||||
|
import { renderHeader, renderEventLine } from "./render.mjs";
|
||||||
|
|
||||||
|
test("renderHeader prints the key wrapper metadata", () => {
|
||||||
|
const header = renderHeader({
|
||||||
|
agent: "scout",
|
||||||
|
task: "Inspect authentication code",
|
||||||
|
cwd: "/repo",
|
||||||
|
requestedModel: "anthropic/claude-sonnet-4-5",
|
||||||
|
resolvedModel: "anthropic/claude-sonnet-4-5",
|
||||||
|
sessionPath: "/repo/.pi/subagents/runs/run-1/child-session.jsonl",
|
||||||
|
});
|
||||||
|
|
||||||
|
assert.match(header, /Agent: scout/);
|
||||||
|
assert.match(header, /Task: Inspect authentication code/);
|
||||||
|
assert.match(header, /Session: \/repo\/\.pi\/subagents\/runs\/run-1\/child-session\.jsonl/);
|
||||||
|
});
|
||||||
|
|
||||||
|
test("renderEventLine makes tool calls readable for a tmux pane", () => {
|
||||||
|
const line = renderEventLine({
|
||||||
|
type: "tool_call",
|
||||||
|
toolName: "bash",
|
||||||
|
args: { command: "rg -n authentication src" },
|
||||||
|
});
|
||||||
|
|
||||||
|
assert.equal(line, "$ rg -n authentication src");
|
||||||
|
});
|
||||||
@@ -1,30 +1,13 @@
|
|||||||
import type { ExtensionAPI } from "@mariozechner/pi-coding-agent";
|
import type { ExtensionAPI } from "@mariozechner/pi-coding-agent";
|
||||||
import { loadWebSearchConfig } from "./src/config.ts";
|
import { registerWebSearchConfigCommand } from "./src/commands/web-search-config.ts";
|
||||||
import { createExaProvider } from "./src/providers/exa.ts";
|
import { createWebSearchRuntime } from "./src/runtime.ts";
|
||||||
import type { WebProvider } from "./src/providers/types.ts";
|
|
||||||
import { createWebFetchTool } from "./src/tools/web-fetch.ts";
|
import { createWebFetchTool } from "./src/tools/web-fetch.ts";
|
||||||
import { createWebSearchTool } from "./src/tools/web-search.ts";
|
import { createWebSearchTool } from "./src/tools/web-search.ts";
|
||||||
|
|
||||||
async function resolveProvider(providerName?: string): Promise<WebProvider> {
|
|
||||||
const config = await loadWebSearchConfig();
|
|
||||||
const selectedName = providerName ?? config.defaultProviderName;
|
|
||||||
const providerConfig = config.providersByName.get(selectedName);
|
|
||||||
|
|
||||||
if (!providerConfig) {
|
|
||||||
throw new Error(
|
|
||||||
`Unknown web-search provider \"${selectedName}\". Configured providers: ${[...config.providersByName.keys()].join(", ")}`,
|
|
||||||
);
|
|
||||||
}
|
|
||||||
|
|
||||||
switch (providerConfig.type) {
|
|
||||||
case "exa":
|
|
||||||
return createExaProvider(providerConfig);
|
|
||||||
default:
|
|
||||||
throw new Error(`Unsupported web-search provider type: ${(providerConfig as { type: string }).type}`);
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
export default function webSearch(pi: ExtensionAPI) {
|
export default function webSearch(pi: ExtensionAPI) {
|
||||||
pi.registerTool(createWebSearchTool({ resolveProvider }));
|
const runtime = createWebSearchRuntime();
|
||||||
pi.registerTool(createWebFetchTool({ resolveProvider }));
|
|
||||||
|
pi.registerTool(createWebSearchTool({ executeSearch: runtime.search }));
|
||||||
|
pi.registerTool(createWebFetchTool({ executeFetch: runtime.fetch }));
|
||||||
|
registerWebSearchConfigCommand(pi);
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -0,0 +1,65 @@
|
|||||||
|
import test from "node:test";
|
||||||
|
import assert from "node:assert/strict";
|
||||||
|
import {
|
||||||
|
createDefaultWebSearchConfig,
|
||||||
|
removeProviderOrThrow,
|
||||||
|
renameProviderOrThrow,
|
||||||
|
setDefaultProviderOrThrow,
|
||||||
|
updateProviderOrThrow,
|
||||||
|
} from "./web-search-config.ts";
|
||||||
|
|
||||||
|
test("createDefaultWebSearchConfig builds a Tavily-first file", () => {
|
||||||
|
const config = createDefaultWebSearchConfig({
|
||||||
|
tavilyName: "tavily-main",
|
||||||
|
tavilyApiKey: "tvly-test-key",
|
||||||
|
});
|
||||||
|
|
||||||
|
assert.equal(config.defaultProvider, "tavily-main");
|
||||||
|
assert.equal(config.providers[0]?.type, "tavily");
|
||||||
|
});
|
||||||
|
|
||||||
|
test("renameProviderOrThrow updates defaultProvider when renaming the default", () => {
|
||||||
|
const config = createDefaultWebSearchConfig({
|
||||||
|
tavilyName: "tavily-main",
|
||||||
|
tavilyApiKey: "tvly-test-key",
|
||||||
|
});
|
||||||
|
|
||||||
|
const next = renameProviderOrThrow(config, "tavily-main", "tavily-primary");
|
||||||
|
|
||||||
|
assert.equal(next.defaultProvider, "tavily-primary");
|
||||||
|
assert.equal(next.providers[0]?.name, "tavily-primary");
|
||||||
|
});
|
||||||
|
|
||||||
|
test("removeProviderOrThrow rejects removing the last provider", () => {
|
||||||
|
const config = createDefaultWebSearchConfig({
|
||||||
|
tavilyName: "tavily-main",
|
||||||
|
tavilyApiKey: "tvly-test-key",
|
||||||
|
});
|
||||||
|
|
||||||
|
assert.throws(() => removeProviderOrThrow(config, "tavily-main"), /last provider/);
|
||||||
|
});
|
||||||
|
|
||||||
|
test("setDefaultProviderOrThrow requires an existing provider name", () => {
|
||||||
|
const config = createDefaultWebSearchConfig({
|
||||||
|
tavilyName: "tavily-main",
|
||||||
|
tavilyApiKey: "tvly-test-key",
|
||||||
|
});
|
||||||
|
|
||||||
|
assert.throws(() => setDefaultProviderOrThrow(config, "missing"), /Unknown provider/);
|
||||||
|
});
|
||||||
|
|
||||||
|
test("updateProviderOrThrow can change provider-specific options without changing type", () => {
|
||||||
|
const config = createDefaultWebSearchConfig({
|
||||||
|
tavilyName: "tavily-main",
|
||||||
|
tavilyApiKey: "tvly-test-key",
|
||||||
|
});
|
||||||
|
|
||||||
|
const next = updateProviderOrThrow(config, "tavily-main", {
|
||||||
|
apiKey: "tvly-next-key",
|
||||||
|
options: { defaultSearchLimit: 8 },
|
||||||
|
});
|
||||||
|
|
||||||
|
assert.equal(next.providers[0]?.apiKey, "tvly-next-key");
|
||||||
|
assert.equal(next.providers[0]?.options?.defaultSearchLimit, 8);
|
||||||
|
assert.equal(next.providers[0]?.type, "tavily");
|
||||||
|
});
|
||||||
@@ -0,0 +1,229 @@
|
|||||||
|
import type { ExtensionAPI } from "@mariozechner/pi-coding-agent";
|
||||||
|
import {
|
||||||
|
getDefaultWebSearchConfigPath,
|
||||||
|
readRawWebSearchConfig,
|
||||||
|
writeWebSearchConfig,
|
||||||
|
WebSearchConfigError,
|
||||||
|
} from "../config.ts";
|
||||||
|
import type { WebSearchConfig, WebSearchProviderConfig } from "../schema.ts";
|
||||||
|
|
||||||
|
export function createDefaultWebSearchConfig(input: { tavilyName: string; tavilyApiKey: string }): WebSearchConfig {
|
||||||
|
return {
|
||||||
|
defaultProvider: input.tavilyName,
|
||||||
|
providers: [
|
||||||
|
{
|
||||||
|
name: input.tavilyName,
|
||||||
|
type: "tavily",
|
||||||
|
apiKey: input.tavilyApiKey,
|
||||||
|
},
|
||||||
|
],
|
||||||
|
};
|
||||||
|
}
|
||||||
|
|
||||||
|
export function setDefaultProviderOrThrow(config: WebSearchConfig, providerName: string): WebSearchConfig {
|
||||||
|
if (!config.providers.some((provider) => provider.name === providerName)) {
|
||||||
|
throw new Error(`Unknown provider: ${providerName}`);
|
||||||
|
}
|
||||||
|
return { ...config, defaultProvider: providerName };
|
||||||
|
}
|
||||||
|
|
||||||
|
export function renameProviderOrThrow(
|
||||||
|
config: WebSearchConfig,
|
||||||
|
currentName: string,
|
||||||
|
nextName: string,
|
||||||
|
): WebSearchConfig {
|
||||||
|
if (!nextName.trim()) {
|
||||||
|
throw new Error("Provider name cannot be blank.");
|
||||||
|
}
|
||||||
|
if (config.providers.some((provider) => provider.name === nextName && provider.name !== currentName)) {
|
||||||
|
throw new Error(`Duplicate provider name: ${nextName}`);
|
||||||
|
}
|
||||||
|
|
||||||
|
return {
|
||||||
|
defaultProvider: config.defaultProvider === currentName ? nextName : config.defaultProvider,
|
||||||
|
providers: config.providers.map((provider) =>
|
||||||
|
provider.name === currentName ? { ...provider, name: nextName } : provider,
|
||||||
|
),
|
||||||
|
};
|
||||||
|
}
|
||||||
|
|
||||||
|
export function updateProviderOrThrow(
|
||||||
|
config: WebSearchConfig,
|
||||||
|
providerName: string,
|
||||||
|
patch: { apiKey?: string; options?: WebSearchProviderConfig["options"] },
|
||||||
|
): WebSearchConfig {
|
||||||
|
const existing = config.providers.find((provider) => provider.name === providerName);
|
||||||
|
if (!existing) {
|
||||||
|
throw new Error(`Unknown provider: ${providerName}`);
|
||||||
|
}
|
||||||
|
if (patch.apiKey !== undefined && !patch.apiKey.trim()) {
|
||||||
|
throw new Error("Provider apiKey cannot be blank.");
|
||||||
|
}
|
||||||
|
|
||||||
|
return {
|
||||||
|
...config,
|
||||||
|
providers: config.providers.map((provider) =>
|
||||||
|
provider.name === providerName
|
||||||
|
? {
|
||||||
|
...provider,
|
||||||
|
apiKey: patch.apiKey ?? provider.apiKey,
|
||||||
|
options: patch.options ?? provider.options,
|
||||||
|
}
|
||||||
|
: provider,
|
||||||
|
),
|
||||||
|
};
|
||||||
|
}
|
||||||
|
|
||||||
|
export function removeProviderOrThrow(config: WebSearchConfig, providerName: string): WebSearchConfig {
|
||||||
|
if (config.providers.length === 1) {
|
||||||
|
throw new Error("Cannot remove the last provider.");
|
||||||
|
}
|
||||||
|
if (config.defaultProvider === providerName) {
|
||||||
|
throw new Error("Cannot remove the default provider before selecting a new default.");
|
||||||
|
}
|
||||||
|
return {
|
||||||
|
...config,
|
||||||
|
providers: config.providers.filter((provider) => provider.name !== providerName),
|
||||||
|
};
|
||||||
|
}
|
||||||
|
|
||||||
|
function upsertProviderOrThrow(config: WebSearchConfig, nextProvider: WebSearchProviderConfig): WebSearchConfig {
|
||||||
|
if (!nextProvider.name.trim()) {
|
||||||
|
throw new Error("Provider name cannot be blank.");
|
||||||
|
}
|
||||||
|
if (!nextProvider.apiKey.trim()) {
|
||||||
|
throw new Error("Provider apiKey cannot be blank.");
|
||||||
|
}
|
||||||
|
|
||||||
|
const withoutSameName = config.providers.filter((provider) => provider.name !== nextProvider.name);
|
||||||
|
return {
|
||||||
|
...config,
|
||||||
|
providers: [...withoutSameName, nextProvider],
|
||||||
|
};
|
||||||
|
}
|
||||||
|
|
||||||
|
async function promptProviderOptions(ctx: any, provider: WebSearchProviderConfig) {
|
||||||
|
const defaultSearchLimit = await ctx.ui.input(
|
||||||
|
`Default search limit for ${provider.name}`,
|
||||||
|
provider.options?.defaultSearchLimit !== undefined ? String(provider.options.defaultSearchLimit) : "",
|
||||||
|
);
|
||||||
|
const defaultFetchTextMaxCharacters = await ctx.ui.input(
|
||||||
|
`Default fetch text max characters for ${provider.name}`,
|
||||||
|
provider.options?.defaultFetchTextMaxCharacters !== undefined
|
||||||
|
? String(provider.options.defaultFetchTextMaxCharacters)
|
||||||
|
: "",
|
||||||
|
);
|
||||||
|
|
||||||
|
const options = {
|
||||||
|
defaultSearchLimit: defaultSearchLimit ? Number(defaultSearchLimit) : undefined,
|
||||||
|
defaultFetchTextMaxCharacters: defaultFetchTextMaxCharacters
|
||||||
|
? Number(defaultFetchTextMaxCharacters)
|
||||||
|
: undefined,
|
||||||
|
};
|
||||||
|
|
||||||
|
return Object.values(options).some((value) => value !== undefined) ? options : undefined;
|
||||||
|
}
|
||||||
|
|
||||||
|
export function registerWebSearchConfigCommand(pi: ExtensionAPI) {
|
||||||
|
pi.registerCommand("web-search-config", {
|
||||||
|
description: "Configure Tavily/Exa providers for web_search and web_fetch",
|
||||||
|
handler: async (_args, ctx) => {
|
||||||
|
const path = getDefaultWebSearchConfigPath();
|
||||||
|
|
||||||
|
let config: WebSearchConfig;
|
||||||
|
try {
|
||||||
|
config = await readRawWebSearchConfig(path);
|
||||||
|
} catch (error) {
|
||||||
|
if (!(error instanceof WebSearchConfigError)) {
|
||||||
|
throw error;
|
||||||
|
}
|
||||||
|
|
||||||
|
const tavilyName = await ctx.ui.input("Create Tavily provider", "tavily-main");
|
||||||
|
const tavilyApiKey = await ctx.ui.input("Tavily API key", "tvly-...");
|
||||||
|
if (!tavilyName || !tavilyApiKey) {
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
config = createDefaultWebSearchConfig({ tavilyName, tavilyApiKey });
|
||||||
|
}
|
||||||
|
|
||||||
|
const action = await ctx.ui.select("Web search config", [
|
||||||
|
"Set default provider",
|
||||||
|
"Add Tavily provider",
|
||||||
|
"Add Exa provider",
|
||||||
|
"Edit provider",
|
||||||
|
"Remove provider",
|
||||||
|
]);
|
||||||
|
if (!action) {
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
|
||||||
|
if (action === "Set default provider") {
|
||||||
|
const nextDefault = await ctx.ui.select(
|
||||||
|
"Choose default provider",
|
||||||
|
config.providers.map((provider) => provider.name),
|
||||||
|
);
|
||||||
|
if (!nextDefault) {
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
config = setDefaultProviderOrThrow(config, nextDefault);
|
||||||
|
}
|
||||||
|
|
||||||
|
if (action === "Add Tavily provider") {
|
||||||
|
const name = await ctx.ui.input("Provider name", "tavily-main");
|
||||||
|
const apiKey = await ctx.ui.input("Tavily API key", "tvly-...");
|
||||||
|
if (!name || !apiKey) {
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
config = upsertProviderOrThrow(config, { name, type: "tavily", apiKey });
|
||||||
|
}
|
||||||
|
|
||||||
|
if (action === "Add Exa provider") {
|
||||||
|
const name = await ctx.ui.input("Provider name", "exa-fallback");
|
||||||
|
const apiKey = await ctx.ui.input("Exa API key", "exa_...");
|
||||||
|
if (!name || !apiKey) {
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
config = upsertProviderOrThrow(config, { name, type: "exa", apiKey });
|
||||||
|
}
|
||||||
|
|
||||||
|
if (action === "Edit provider") {
|
||||||
|
const providerName = await ctx.ui.select(
|
||||||
|
"Choose provider",
|
||||||
|
config.providers.map((provider) => provider.name),
|
||||||
|
);
|
||||||
|
if (!providerName) {
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
|
||||||
|
const existing = config.providers.find((provider) => provider.name === providerName)!;
|
||||||
|
const nextName = await ctx.ui.input("Provider name", existing.name);
|
||||||
|
const nextApiKey = await ctx.ui.input(`API key for ${existing.name}`, existing.apiKey);
|
||||||
|
if (!nextName || !nextApiKey) {
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
|
||||||
|
config = renameProviderOrThrow(config, existing.name, nextName);
|
||||||
|
const renamed = config.providers.find((provider) => provider.name === nextName)!;
|
||||||
|
const nextOptions = await promptProviderOptions(ctx, renamed);
|
||||||
|
config = updateProviderOrThrow(config, nextName, {
|
||||||
|
apiKey: nextApiKey,
|
||||||
|
options: nextOptions,
|
||||||
|
});
|
||||||
|
}
|
||||||
|
|
||||||
|
if (action === "Remove provider") {
|
||||||
|
const providerName = await ctx.ui.select(
|
||||||
|
"Choose provider to remove",
|
||||||
|
config.providers.map((provider) => provider.name),
|
||||||
|
);
|
||||||
|
if (!providerName) {
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
config = removeProviderOrThrow(config, providerName);
|
||||||
|
}
|
||||||
|
|
||||||
|
await writeWebSearchConfig(path, config);
|
||||||
|
ctx.ui.notify(`Saved web-search config to ${path}`, "info");
|
||||||
|
},
|
||||||
|
});
|
||||||
|
}
|
||||||
@@ -37,6 +37,30 @@ test("loadWebSearchConfig returns a normalized default provider and provider loo
|
|||||||
assert.equal(config.providers[0]?.options?.defaultSearchLimit, 7);
|
assert.equal(config.providers[0]?.options?.defaultSearchLimit, 7);
|
||||||
});
|
});
|
||||||
|
|
||||||
|
test("loadWebSearchConfig normalizes a Tavily default with Exa fallback", async () => {
|
||||||
|
const file = await writeTempConfig({
|
||||||
|
defaultProvider: "tavily-main",
|
||||||
|
providers: [
|
||||||
|
{
|
||||||
|
name: "tavily-main",
|
||||||
|
type: "tavily",
|
||||||
|
apiKey: "tvly-test-key",
|
||||||
|
},
|
||||||
|
{
|
||||||
|
name: "exa-fallback",
|
||||||
|
type: "exa",
|
||||||
|
apiKey: "exa-test-key",
|
||||||
|
},
|
||||||
|
],
|
||||||
|
});
|
||||||
|
|
||||||
|
const config = await loadWebSearchConfig(file);
|
||||||
|
|
||||||
|
assert.equal(config.defaultProviderName, "tavily-main");
|
||||||
|
assert.equal(config.defaultProvider.type, "tavily");
|
||||||
|
assert.equal(config.providersByName.get("exa-fallback")?.type, "exa");
|
||||||
|
});
|
||||||
|
|
||||||
test("loadWebSearchConfig rejects a missing default provider target", async () => {
|
test("loadWebSearchConfig rejects a missing default provider target", async () => {
|
||||||
const file = await writeTempConfig({
|
const file = await writeTempConfig({
|
||||||
defaultProvider: "missing",
|
defaultProvider: "missing",
|
||||||
|
|||||||
@@ -1,15 +1,19 @@
|
|||||||
import { readFile } from "node:fs/promises";
|
import { mkdir, readFile, writeFile } from "node:fs/promises";
|
||||||
import { homedir } from "node:os";
|
import { homedir } from "node:os";
|
||||||
import { join } from "node:path";
|
import { dirname, join } from "node:path";
|
||||||
import { Value } from "@sinclair/typebox/value";
|
import { Value } from "@sinclair/typebox/value";
|
||||||
import { WebSearchConfigSchema, type ExaProviderConfig, type WebSearchConfig } from "./schema.ts";
|
import {
|
||||||
|
WebSearchConfigSchema,
|
||||||
|
type WebSearchConfig,
|
||||||
|
type WebSearchProviderConfig,
|
||||||
|
} from "./schema.ts";
|
||||||
|
|
||||||
export interface ResolvedWebSearchConfig {
|
export interface ResolvedWebSearchConfig {
|
||||||
path: string;
|
path: string;
|
||||||
defaultProviderName: string;
|
defaultProviderName: string;
|
||||||
defaultProvider: ExaProviderConfig;
|
defaultProvider: WebSearchProviderConfig;
|
||||||
providers: ExaProviderConfig[];
|
providers: WebSearchProviderConfig[];
|
||||||
providersByName: Map<string, ExaProviderConfig>;
|
providersByName: Map<string, WebSearchProviderConfig>;
|
||||||
}
|
}
|
||||||
|
|
||||||
export class WebSearchConfigError extends Error {
|
export class WebSearchConfigError extends Error {
|
||||||
@@ -26,10 +30,15 @@ export function getDefaultWebSearchConfigPath() {
|
|||||||
function exampleConfigSnippet() {
|
function exampleConfigSnippet() {
|
||||||
return JSON.stringify(
|
return JSON.stringify(
|
||||||
{
|
{
|
||||||
defaultProvider: "exa-main",
|
defaultProvider: "tavily-main",
|
||||||
providers: [
|
providers: [
|
||||||
{
|
{
|
||||||
name: "exa-main",
|
name: "tavily-main",
|
||||||
|
type: "tavily",
|
||||||
|
apiKey: "tvly-...",
|
||||||
|
},
|
||||||
|
{
|
||||||
|
name: "exa-fallback",
|
||||||
type: "exa",
|
type: "exa",
|
||||||
apiKey: "exa_...",
|
apiKey: "exa_...",
|
||||||
},
|
},
|
||||||
@@ -41,7 +50,7 @@ function exampleConfigSnippet() {
|
|||||||
}
|
}
|
||||||
|
|
||||||
export function normalizeWebSearchConfig(config: WebSearchConfig, path: string): ResolvedWebSearchConfig {
|
export function normalizeWebSearchConfig(config: WebSearchConfig, path: string): ResolvedWebSearchConfig {
|
||||||
const providersByName = new Map<string, ExaProviderConfig>();
|
const providersByName = new Map<string, WebSearchProviderConfig>();
|
||||||
|
|
||||||
for (const provider of config.providers) {
|
for (const provider of config.providers) {
|
||||||
if (!provider.apiKey.trim()) {
|
if (!provider.apiKey.trim()) {
|
||||||
@@ -69,19 +78,7 @@ export function normalizeWebSearchConfig(config: WebSearchConfig, path: string):
|
|||||||
};
|
};
|
||||||
}
|
}
|
||||||
|
|
||||||
export async function loadWebSearchConfig(path = getDefaultWebSearchConfigPath()) {
|
function parseWebSearchConfig(raw: string, path: string) {
|
||||||
let raw: string;
|
|
||||||
try {
|
|
||||||
raw = await readFile(path, "utf8");
|
|
||||||
} catch (error) {
|
|
||||||
if ((error as NodeJS.ErrnoException).code === "ENOENT") {
|
|
||||||
throw new WebSearchConfigError(
|
|
||||||
`Missing web-search config at ${path}.\nCreate ${path} with contents like:\n${exampleConfigSnippet()}`,
|
|
||||||
);
|
|
||||||
}
|
|
||||||
throw error;
|
|
||||||
}
|
|
||||||
|
|
||||||
let parsed: unknown;
|
let parsed: unknown;
|
||||||
try {
|
try {
|
||||||
parsed = JSON.parse(raw);
|
parsed = JSON.parse(raw);
|
||||||
@@ -96,5 +93,35 @@ export async function loadWebSearchConfig(path = getDefaultWebSearchConfigPath()
|
|||||||
);
|
);
|
||||||
}
|
}
|
||||||
|
|
||||||
return normalizeWebSearchConfig(parsed as WebSearchConfig, path);
|
return parsed as WebSearchConfig;
|
||||||
|
}
|
||||||
|
|
||||||
|
export async function readRawWebSearchConfig(path = getDefaultWebSearchConfigPath()): Promise<WebSearchConfig> {
|
||||||
|
let raw: string;
|
||||||
|
try {
|
||||||
|
raw = await readFile(path, "utf8");
|
||||||
|
} catch (error) {
|
||||||
|
if ((error as NodeJS.ErrnoException).code === "ENOENT") {
|
||||||
|
throw new WebSearchConfigError(
|
||||||
|
`Missing web-search config at ${path}.\nCreate ${path} with contents like:\n${exampleConfigSnippet()}`,
|
||||||
|
);
|
||||||
|
}
|
||||||
|
throw error;
|
||||||
|
}
|
||||||
|
|
||||||
|
return parseWebSearchConfig(raw, path);
|
||||||
|
}
|
||||||
|
|
||||||
|
export function stringifyWebSearchConfig(config: WebSearchConfig) {
|
||||||
|
return `${JSON.stringify(config, null, 2)}\n`;
|
||||||
|
}
|
||||||
|
|
||||||
|
export async function writeWebSearchConfig(path: string, config: WebSearchConfig) {
|
||||||
|
await mkdir(dirname(path), { recursive: true });
|
||||||
|
await writeFile(path, stringifyWebSearchConfig(config), "utf8");
|
||||||
|
}
|
||||||
|
|
||||||
|
export async function loadWebSearchConfig(path = getDefaultWebSearchConfigPath()) {
|
||||||
|
const parsed = await readRawWebSearchConfig(path);
|
||||||
|
return normalizeWebSearchConfig(parsed, path);
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -2,14 +2,19 @@ import test from "node:test";
|
|||||||
import assert from "node:assert/strict";
|
import assert from "node:assert/strict";
|
||||||
import webSearchExtension from "../index.ts";
|
import webSearchExtension from "../index.ts";
|
||||||
|
|
||||||
test("the extension entrypoint registers both web_search and web_fetch", () => {
|
test("the extension entrypoint registers both tools and the config command", () => {
|
||||||
const registeredTools: string[] = [];
|
const registeredTools: string[] = [];
|
||||||
|
const registeredCommands: string[] = [];
|
||||||
|
|
||||||
webSearchExtension({
|
webSearchExtension({
|
||||||
registerTool(tool: { name: string }) {
|
registerTool(tool: { name: string }) {
|
||||||
registeredTools.push(tool.name);
|
registeredTools.push(tool.name);
|
||||||
},
|
},
|
||||||
|
registerCommand(name: string) {
|
||||||
|
registeredCommands.push(name);
|
||||||
|
},
|
||||||
} as any);
|
} as any);
|
||||||
|
|
||||||
assert.deepEqual(registeredTools, ["web_search", "web_fetch"]);
|
assert.deepEqual(registeredTools, ["web_search", "web_fetch"]);
|
||||||
|
assert.deepEqual(registeredCommands, ["web-search-config"]);
|
||||||
});
|
});
|
||||||
|
|||||||
@@ -21,6 +21,27 @@ test("formatSearchOutput renders a compact metadata-only list", () => {
|
|||||||
assert.match(output, /https:\/\/exa.ai\/docs/);
|
assert.match(output, /https:\/\/exa.ai\/docs/);
|
||||||
});
|
});
|
||||||
|
|
||||||
|
test("formatSearchOutput shows answer and fallback provider metadata", () => {
|
||||||
|
const output = formatSearchOutput({
|
||||||
|
providerName: "exa-fallback",
|
||||||
|
answer: "pi is a coding agent",
|
||||||
|
execution: {
|
||||||
|
actualProviderName: "exa-fallback",
|
||||||
|
failoverFromProviderName: "tavily-main",
|
||||||
|
},
|
||||||
|
results: [
|
||||||
|
{
|
||||||
|
title: "pi docs",
|
||||||
|
url: "https://pi.dev",
|
||||||
|
rawContent: "Very long raw content body",
|
||||||
|
},
|
||||||
|
],
|
||||||
|
} as any);
|
||||||
|
|
||||||
|
assert.match(output, /Answer: pi is a coding agent/);
|
||||||
|
assert.match(output, /Fallback: tavily-main -> exa-fallback/);
|
||||||
|
});
|
||||||
|
|
||||||
test("truncateText shortens long fetch bodies with an ellipsis", () => {
|
test("truncateText shortens long fetch bodies with an ellipsis", () => {
|
||||||
assert.equal(truncateText("abcdef", 4), "abc…");
|
assert.equal(truncateText("abcdef", 4), "abc…");
|
||||||
assert.equal(truncateText("abc", 10), "abc");
|
assert.equal(truncateText("abc", 10), "abc");
|
||||||
@@ -51,3 +72,26 @@ test("formatFetchOutput includes both successful and failed URLs", () => {
|
|||||||
assert.match(output, /429 rate limited/);
|
assert.match(output, /429 rate limited/);
|
||||||
assert.match(output, /This is a very long…/);
|
assert.match(output, /This is a very long…/);
|
||||||
});
|
});
|
||||||
|
|
||||||
|
test("formatFetchOutput shows fallback metadata and favicon/images when present", () => {
|
||||||
|
const output = formatFetchOutput({
|
||||||
|
providerName: "exa-fallback",
|
||||||
|
execution: {
|
||||||
|
actualProviderName: "exa-fallback",
|
||||||
|
failoverFromProviderName: "tavily-main",
|
||||||
|
},
|
||||||
|
results: [
|
||||||
|
{
|
||||||
|
url: "https://pi.dev",
|
||||||
|
title: "pi",
|
||||||
|
text: "Fetched body",
|
||||||
|
favicon: "https://pi.dev/favicon.ico",
|
||||||
|
images: ["https://pi.dev/logo.png"],
|
||||||
|
},
|
||||||
|
],
|
||||||
|
} as any);
|
||||||
|
|
||||||
|
assert.match(output, /Fallback: tavily-main -> exa-fallback/);
|
||||||
|
assert.match(output, /Favicon: https:\/\/pi.dev\/favicon.ico/);
|
||||||
|
assert.match(output, /Images:/);
|
||||||
|
});
|
||||||
|
|||||||
@@ -1,5 +1,15 @@
|
|||||||
import type { NormalizedFetchResponse, NormalizedSearchResponse } from "./providers/types.ts";
|
import type { NormalizedFetchResponse, NormalizedSearchResponse } from "./providers/types.ts";
|
||||||
|
|
||||||
|
function formatFallbackLine(execution?: {
|
||||||
|
actualProviderName?: string;
|
||||||
|
failoverFromProviderName?: string;
|
||||||
|
}) {
|
||||||
|
if (!execution?.failoverFromProviderName || !execution.actualProviderName) {
|
||||||
|
return undefined;
|
||||||
|
}
|
||||||
|
return `Fallback: ${execution.failoverFromProviderName} -> ${execution.actualProviderName}`;
|
||||||
|
}
|
||||||
|
|
||||||
export function truncateText(text: string, maxCharacters = 4000) {
|
export function truncateText(text: string, maxCharacters = 4000) {
|
||||||
if (text.length <= maxCharacters) {
|
if (text.length <= maxCharacters) {
|
||||||
return text;
|
return text;
|
||||||
@@ -7,14 +17,24 @@ export function truncateText(text: string, maxCharacters = 4000) {
|
|||||||
return `${text.slice(0, Math.max(0, maxCharacters - 1))}…`;
|
return `${text.slice(0, Math.max(0, maxCharacters - 1))}…`;
|
||||||
}
|
}
|
||||||
|
|
||||||
export function formatSearchOutput(response: NormalizedSearchResponse) {
|
export function formatSearchOutput(response: NormalizedSearchResponse & { execution?: any }) {
|
||||||
if (response.results.length === 0) {
|
const lines: string[] = [];
|
||||||
return `No web results via ${response.providerName}.`;
|
const fallbackLine = formatFallbackLine(response.execution);
|
||||||
|
|
||||||
|
if (fallbackLine) {
|
||||||
|
lines.push(fallbackLine, "");
|
||||||
}
|
}
|
||||||
|
|
||||||
const lines = [
|
if (response.answer) {
|
||||||
`Found ${response.results.length} web result${response.results.length === 1 ? "" : "s"} via ${response.providerName}:`,
|
lines.push(`Answer: ${response.answer}`, "");
|
||||||
];
|
}
|
||||||
|
|
||||||
|
if (response.results.length === 0) {
|
||||||
|
lines.push(`No web results via ${response.providerName}.`);
|
||||||
|
return lines.join("\n");
|
||||||
|
}
|
||||||
|
|
||||||
|
lines.push(`Found ${response.results.length} web result${response.results.length === 1 ? "" : "s"} via ${response.providerName}:`);
|
||||||
|
|
||||||
for (const [index, result] of response.results.entries()) {
|
for (const [index, result] of response.results.entries()) {
|
||||||
lines.push(`${index + 1}. ${result.title ?? "(untitled)"}`);
|
lines.push(`${index + 1}. ${result.title ?? "(untitled)"}`);
|
||||||
@@ -28,6 +48,14 @@ export function formatSearchOutput(response: NormalizedSearchResponse) {
|
|||||||
if (typeof result.score === "number") {
|
if (typeof result.score === "number") {
|
||||||
lines.push(` Score: ${result.score}`);
|
lines.push(` Score: ${result.score}`);
|
||||||
}
|
}
|
||||||
|
|
||||||
|
if (result.content) {
|
||||||
|
lines.push(` Snippet: ${truncateText(result.content, 500)}`);
|
||||||
|
}
|
||||||
|
|
||||||
|
if (result.rawContent) {
|
||||||
|
lines.push(` Raw content: ${truncateText(result.rawContent, 700)}`);
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
return lines.join("\n");
|
return lines.join("\n");
|
||||||
@@ -37,11 +65,16 @@ export interface FetchFormatOptions {
|
|||||||
maxCharactersPerResult?: number;
|
maxCharactersPerResult?: number;
|
||||||
}
|
}
|
||||||
|
|
||||||
export function formatFetchOutput(response: NormalizedFetchResponse, options: FetchFormatOptions = {}) {
|
export function formatFetchOutput(response: NormalizedFetchResponse & { execution?: any }, options: FetchFormatOptions = {}) {
|
||||||
const maxCharactersPerResult = options.maxCharactersPerResult ?? 4000;
|
const maxCharactersPerResult = options.maxCharactersPerResult ?? 4000;
|
||||||
const lines = [
|
const lines: string[] = [];
|
||||||
`Fetched ${response.results.length} URL${response.results.length === 1 ? "" : "s"} via ${response.providerName}:`,
|
const fallbackLine = formatFallbackLine(response.execution);
|
||||||
];
|
|
||||||
|
if (fallbackLine) {
|
||||||
|
lines.push(fallbackLine, "");
|
||||||
|
}
|
||||||
|
|
||||||
|
lines.push(`Fetched ${response.results.length} URL${response.results.length === 1 ? "" : "s"} via ${response.providerName}:`);
|
||||||
|
|
||||||
for (const result of response.results) {
|
for (const result of response.results) {
|
||||||
lines.push("");
|
lines.push("");
|
||||||
@@ -66,6 +99,15 @@ export function formatFetchOutput(response: NormalizedFetchResponse, options: Fe
|
|||||||
lines.push(`- ${highlight}`);
|
lines.push(`- ${highlight}`);
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
if (result.favicon) {
|
||||||
|
lines.push(`Favicon: ${result.favicon}`);
|
||||||
|
}
|
||||||
|
if (result.images?.length) {
|
||||||
|
lines.push("Images:");
|
||||||
|
for (const image of result.images) {
|
||||||
|
lines.push(`- ${image}`);
|
||||||
|
}
|
||||||
|
}
|
||||||
if (result.text) {
|
if (result.text) {
|
||||||
lines.push("Text:");
|
lines.push("Text:");
|
||||||
lines.push(truncateText(result.text, maxCharactersPerResult));
|
lines.push(truncateText(result.text, maxCharactersPerResult));
|
||||||
|
|||||||
84
.pi/agent/extensions/web-search/src/providers/tavily.test.ts
Normal file
84
.pi/agent/extensions/web-search/src/providers/tavily.test.ts
Normal file
@@ -0,0 +1,84 @@
|
|||||||
|
import test from "node:test";
|
||||||
|
import assert from "node:assert/strict";
|
||||||
|
import { createTavilyProvider } from "./tavily.ts";
|
||||||
|
|
||||||
|
const baseConfig = {
|
||||||
|
name: "tavily-main",
|
||||||
|
type: "tavily" as const,
|
||||||
|
apiKey: "tvly-test-key",
|
||||||
|
options: {
|
||||||
|
defaultSearchLimit: 6,
|
||||||
|
defaultFetchTextMaxCharacters: 8000,
|
||||||
|
},
|
||||||
|
};
|
||||||
|
|
||||||
|
test("createTavilyProvider maps search requests to Tavily REST params", async () => {
|
||||||
|
let captured: RequestInit | undefined;
|
||||||
|
|
||||||
|
const provider = createTavilyProvider(baseConfig, async (_url, init) => {
|
||||||
|
captured = init;
|
||||||
|
return new Response(
|
||||||
|
JSON.stringify({
|
||||||
|
answer: "pi is a coding agent",
|
||||||
|
results: [
|
||||||
|
{
|
||||||
|
title: "pi docs",
|
||||||
|
url: "https://pi.dev",
|
||||||
|
content: "pi docs summary",
|
||||||
|
raw_content: "long raw body",
|
||||||
|
},
|
||||||
|
],
|
||||||
|
}),
|
||||||
|
{ status: 200 },
|
||||||
|
);
|
||||||
|
});
|
||||||
|
|
||||||
|
const result = await provider.search({
|
||||||
|
query: "pi docs",
|
||||||
|
limit: 4,
|
||||||
|
tavily: {
|
||||||
|
includeAnswer: true,
|
||||||
|
includeRawContent: true,
|
||||||
|
searchDepth: "advanced",
|
||||||
|
},
|
||||||
|
});
|
||||||
|
|
||||||
|
const body = JSON.parse(String(captured?.body));
|
||||||
|
assert.equal(body.max_results, 4);
|
||||||
|
assert.equal(body.include_answer, true);
|
||||||
|
assert.equal(body.include_raw_content, true);
|
||||||
|
assert.equal(body.search_depth, "advanced");
|
||||||
|
assert.equal(result.answer, "pi is a coding agent");
|
||||||
|
assert.equal(result.results[0]?.rawContent, "long raw body");
|
||||||
|
});
|
||||||
|
|
||||||
|
test("createTavilyProvider maps extract responses into normalized fetch results", async () => {
|
||||||
|
const provider = createTavilyProvider(baseConfig, async () => {
|
||||||
|
return new Response(
|
||||||
|
JSON.stringify({
|
||||||
|
results: [
|
||||||
|
{
|
||||||
|
url: "https://pi.dev",
|
||||||
|
title: "pi",
|
||||||
|
raw_content: "Fetched body",
|
||||||
|
images: ["https://pi.dev/logo.png"],
|
||||||
|
favicon: "https://pi.dev/favicon.ico",
|
||||||
|
},
|
||||||
|
],
|
||||||
|
}),
|
||||||
|
{ status: 200 },
|
||||||
|
);
|
||||||
|
});
|
||||||
|
|
||||||
|
const result = await provider.fetch({
|
||||||
|
urls: ["https://pi.dev"],
|
||||||
|
tavily: {
|
||||||
|
includeImages: true,
|
||||||
|
includeFavicon: true,
|
||||||
|
},
|
||||||
|
});
|
||||||
|
|
||||||
|
assert.equal(result.results[0]?.text, "Fetched body");
|
||||||
|
assert.deepEqual(result.results[0]?.images, ["https://pi.dev/logo.png"]);
|
||||||
|
assert.equal(result.results[0]?.favicon, "https://pi.dev/favicon.ico");
|
||||||
|
});
|
||||||
107
.pi/agent/extensions/web-search/src/providers/tavily.ts
Normal file
107
.pi/agent/extensions/web-search/src/providers/tavily.ts
Normal file
@@ -0,0 +1,107 @@
|
|||||||
|
import type { TavilyProviderConfig } from "../schema.ts";
|
||||||
|
import type {
|
||||||
|
NormalizedFetchRequest,
|
||||||
|
NormalizedFetchResponse,
|
||||||
|
NormalizedSearchRequest,
|
||||||
|
NormalizedSearchResponse,
|
||||||
|
WebProvider,
|
||||||
|
} from "./types.ts";
|
||||||
|
|
||||||
|
export type TavilyFetchLike = (input: string, init?: RequestInit) => Promise<Response>;
|
||||||
|
|
||||||
|
async function readError(response: Response) {
|
||||||
|
const text = await response.text();
|
||||||
|
throw new Error(`Tavily ${response.status} ${response.statusText}: ${text.slice(0, 300)}`);
|
||||||
|
}
|
||||||
|
|
||||||
|
export function createTavilyProvider(
|
||||||
|
config: TavilyProviderConfig,
|
||||||
|
fetchImpl: TavilyFetchLike = fetch,
|
||||||
|
): WebProvider {
|
||||||
|
return {
|
||||||
|
name: config.name,
|
||||||
|
type: config.type,
|
||||||
|
|
||||||
|
async search(request: NormalizedSearchRequest): Promise<NormalizedSearchResponse> {
|
||||||
|
const response = await fetchImpl("https://api.tavily.com/search", {
|
||||||
|
method: "POST",
|
||||||
|
headers: {
|
||||||
|
"content-type": "application/json",
|
||||||
|
authorization: `Bearer ${config.apiKey}`,
|
||||||
|
},
|
||||||
|
body: JSON.stringify({
|
||||||
|
query: request.query,
|
||||||
|
max_results: request.limit ?? config.options?.defaultSearchLimit ?? 5,
|
||||||
|
include_domains: request.includeDomains,
|
||||||
|
exclude_domains: request.excludeDomains,
|
||||||
|
start_date: request.startPublishedDate,
|
||||||
|
end_date: request.endPublishedDate,
|
||||||
|
topic: request.tavily?.topic,
|
||||||
|
search_depth: request.tavily?.searchDepth,
|
||||||
|
time_range: request.tavily?.timeRange,
|
||||||
|
days: request.tavily?.days,
|
||||||
|
chunks_per_source: request.tavily?.chunksPerSource,
|
||||||
|
include_answer: request.tavily?.includeAnswer,
|
||||||
|
include_raw_content: request.tavily?.includeRawContent,
|
||||||
|
include_images: request.tavily?.includeImages,
|
||||||
|
}),
|
||||||
|
});
|
||||||
|
|
||||||
|
if (!response.ok) {
|
||||||
|
await readError(response);
|
||||||
|
}
|
||||||
|
|
||||||
|
const data = (await response.json()) as any;
|
||||||
|
return {
|
||||||
|
providerName: config.name,
|
||||||
|
requestId: data.request_id,
|
||||||
|
answer: typeof data.answer === "string" ? data.answer : undefined,
|
||||||
|
results: (data.results ?? []).map((item: any) => ({
|
||||||
|
title: item.title ?? null,
|
||||||
|
url: item.url,
|
||||||
|
content: typeof item.content === "string" ? item.content : undefined,
|
||||||
|
rawContent: typeof item.raw_content === "string" ? item.raw_content : undefined,
|
||||||
|
images: Array.isArray(item.images) ? item.images : undefined,
|
||||||
|
score: item.score,
|
||||||
|
publishedDate: item.published_date,
|
||||||
|
})),
|
||||||
|
};
|
||||||
|
},
|
||||||
|
|
||||||
|
async fetch(request: NormalizedFetchRequest): Promise<NormalizedFetchResponse> {
|
||||||
|
const response = await fetchImpl("https://api.tavily.com/extract", {
|
||||||
|
method: "POST",
|
||||||
|
headers: {
|
||||||
|
"content-type": "application/json",
|
||||||
|
authorization: `Bearer ${config.apiKey}`,
|
||||||
|
},
|
||||||
|
body: JSON.stringify({
|
||||||
|
urls: request.urls,
|
||||||
|
query: request.tavily?.query,
|
||||||
|
extract_depth: request.tavily?.extractDepth,
|
||||||
|
chunks_per_source: request.tavily?.chunksPerSource,
|
||||||
|
include_images: request.tavily?.includeImages,
|
||||||
|
include_favicon: request.tavily?.includeFavicon,
|
||||||
|
format: request.tavily?.format,
|
||||||
|
}),
|
||||||
|
});
|
||||||
|
|
||||||
|
if (!response.ok) {
|
||||||
|
await readError(response);
|
||||||
|
}
|
||||||
|
|
||||||
|
const data = (await response.json()) as any;
|
||||||
|
return {
|
||||||
|
providerName: config.name,
|
||||||
|
requestIds: data.request_id ? [data.request_id] : [],
|
||||||
|
results: (data.results ?? []).map((item: any) => ({
|
||||||
|
url: item.url,
|
||||||
|
title: item.title ?? null,
|
||||||
|
text: typeof item.raw_content === "string" ? item.raw_content : undefined,
|
||||||
|
images: Array.isArray(item.images) ? item.images : undefined,
|
||||||
|
favicon: typeof item.favicon === "string" ? item.favicon : undefined,
|
||||||
|
})),
|
||||||
|
};
|
||||||
|
},
|
||||||
|
};
|
||||||
|
}
|
||||||
@@ -1,3 +1,23 @@
|
|||||||
|
export interface TavilySearchOptions {
|
||||||
|
searchDepth?: "advanced" | "basic" | "fast" | "ultra-fast";
|
||||||
|
topic?: "general" | "news" | "finance";
|
||||||
|
timeRange?: string;
|
||||||
|
days?: number;
|
||||||
|
chunksPerSource?: number;
|
||||||
|
includeAnswer?: boolean;
|
||||||
|
includeRawContent?: boolean;
|
||||||
|
includeImages?: boolean;
|
||||||
|
}
|
||||||
|
|
||||||
|
export interface TavilyFetchOptions {
|
||||||
|
query?: string;
|
||||||
|
extractDepth?: "basic" | "advanced";
|
||||||
|
chunksPerSource?: number;
|
||||||
|
includeImages?: boolean;
|
||||||
|
includeFavicon?: boolean;
|
||||||
|
format?: string;
|
||||||
|
}
|
||||||
|
|
||||||
export interface NormalizedSearchRequest {
|
export interface NormalizedSearchRequest {
|
||||||
query: string;
|
query: string;
|
||||||
limit?: number;
|
limit?: number;
|
||||||
@@ -7,6 +27,7 @@ export interface NormalizedSearchRequest {
|
|||||||
endPublishedDate?: string;
|
endPublishedDate?: string;
|
||||||
category?: string;
|
category?: string;
|
||||||
provider?: string;
|
provider?: string;
|
||||||
|
tavily?: TavilySearchOptions;
|
||||||
}
|
}
|
||||||
|
|
||||||
export interface NormalizedSearchResult {
|
export interface NormalizedSearchResult {
|
||||||
@@ -16,12 +37,16 @@ export interface NormalizedSearchResult {
|
|||||||
publishedDate?: string;
|
publishedDate?: string;
|
||||||
author?: string;
|
author?: string;
|
||||||
score?: number;
|
score?: number;
|
||||||
|
content?: string;
|
||||||
|
rawContent?: string;
|
||||||
|
images?: string[];
|
||||||
}
|
}
|
||||||
|
|
||||||
export interface NormalizedSearchResponse {
|
export interface NormalizedSearchResponse {
|
||||||
providerName: string;
|
providerName: string;
|
||||||
requestId?: string;
|
requestId?: string;
|
||||||
searchTime?: number;
|
searchTime?: number;
|
||||||
|
answer?: string;
|
||||||
results: NormalizedSearchResult[];
|
results: NormalizedSearchResult[];
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -32,6 +57,7 @@ export interface NormalizedFetchRequest {
|
|||||||
summary?: boolean;
|
summary?: boolean;
|
||||||
textMaxCharacters?: number;
|
textMaxCharacters?: number;
|
||||||
provider?: string;
|
provider?: string;
|
||||||
|
tavily?: TavilyFetchOptions;
|
||||||
}
|
}
|
||||||
|
|
||||||
export interface NormalizedFetchResult {
|
export interface NormalizedFetchResult {
|
||||||
@@ -40,6 +66,8 @@ export interface NormalizedFetchResult {
|
|||||||
text?: string;
|
text?: string;
|
||||||
highlights?: string[];
|
highlights?: string[];
|
||||||
summary?: string;
|
summary?: string;
|
||||||
|
images?: string[];
|
||||||
|
favicon?: string;
|
||||||
error?: string;
|
error?: string;
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
85
.pi/agent/extensions/web-search/src/runtime.test.ts
Normal file
85
.pi/agent/extensions/web-search/src/runtime.test.ts
Normal file
@@ -0,0 +1,85 @@
|
|||||||
|
import test from "node:test";
|
||||||
|
import assert from "node:assert/strict";
|
||||||
|
import { createWebSearchRuntime } from "./runtime.ts";
|
||||||
|
|
||||||
|
function createProvider(name: string, type: string, handlers: Partial<any>) {
|
||||||
|
return {
|
||||||
|
name,
|
||||||
|
type,
|
||||||
|
async search(request: any) {
|
||||||
|
return handlers.search?.(request);
|
||||||
|
},
|
||||||
|
async fetch(request: any) {
|
||||||
|
return handlers.fetch?.(request);
|
||||||
|
},
|
||||||
|
};
|
||||||
|
}
|
||||||
|
|
||||||
|
test("search retries Tavily failures once with Exa", async () => {
|
||||||
|
const runtime = createWebSearchRuntime({
|
||||||
|
loadConfig: async () => ({
|
||||||
|
path: "test.json",
|
||||||
|
defaultProviderName: "tavily-main",
|
||||||
|
defaultProvider: { name: "tavily-main", type: "tavily", apiKey: "tvly" },
|
||||||
|
providers: [
|
||||||
|
{ name: "tavily-main", type: "tavily", apiKey: "tvly" },
|
||||||
|
{ name: "exa-fallback", type: "exa", apiKey: "exa" },
|
||||||
|
],
|
||||||
|
providersByName: new Map([
|
||||||
|
["tavily-main", { name: "tavily-main", type: "tavily", apiKey: "tvly" }],
|
||||||
|
["exa-fallback", { name: "exa-fallback", type: "exa", apiKey: "exa" }],
|
||||||
|
]),
|
||||||
|
}),
|
||||||
|
createProvider(providerConfig) {
|
||||||
|
if (providerConfig.type === "tavily") {
|
||||||
|
return createProvider(providerConfig.name, providerConfig.type, {
|
||||||
|
search: async () => {
|
||||||
|
throw new Error("503 upstream unavailable");
|
||||||
|
},
|
||||||
|
});
|
||||||
|
}
|
||||||
|
return createProvider(providerConfig.name, providerConfig.type, {
|
||||||
|
search: async () => ({
|
||||||
|
providerName: providerConfig.name,
|
||||||
|
results: [{ title: "Exa hit", url: "https://exa.ai" }],
|
||||||
|
}),
|
||||||
|
});
|
||||||
|
},
|
||||||
|
});
|
||||||
|
|
||||||
|
const result = await runtime.search({ query: "pi docs" });
|
||||||
|
|
||||||
|
assert.equal(result.execution.actualProviderName, "exa-fallback");
|
||||||
|
assert.equal(result.execution.failoverFromProviderName, "tavily-main");
|
||||||
|
assert.match(result.execution.failoverReason ?? "", /503/);
|
||||||
|
});
|
||||||
|
|
||||||
|
test("search does not retry when Exa was explicitly selected", async () => {
|
||||||
|
const runtime = createWebSearchRuntime({
|
||||||
|
loadConfig: async () => ({
|
||||||
|
path: "test.json",
|
||||||
|
defaultProviderName: "tavily-main",
|
||||||
|
defaultProvider: { name: "tavily-main", type: "tavily", apiKey: "tvly" },
|
||||||
|
providers: [
|
||||||
|
{ name: "tavily-main", type: "tavily", apiKey: "tvly" },
|
||||||
|
{ name: "exa-fallback", type: "exa", apiKey: "exa" },
|
||||||
|
],
|
||||||
|
providersByName: new Map([
|
||||||
|
["tavily-main", { name: "tavily-main", type: "tavily", apiKey: "tvly" }],
|
||||||
|
["exa-fallback", { name: "exa-fallback", type: "exa", apiKey: "exa" }],
|
||||||
|
]),
|
||||||
|
}),
|
||||||
|
createProvider(providerConfig) {
|
||||||
|
return createProvider(providerConfig.name, providerConfig.type, {
|
||||||
|
search: async () => {
|
||||||
|
throw new Error(`boom:${providerConfig.name}`);
|
||||||
|
},
|
||||||
|
});
|
||||||
|
},
|
||||||
|
});
|
||||||
|
|
||||||
|
await assert.rejects(
|
||||||
|
() => runtime.search({ query: "pi docs", provider: "exa-fallback" }),
|
||||||
|
/boom:exa-fallback/,
|
||||||
|
);
|
||||||
|
});
|
||||||
139
.pi/agent/extensions/web-search/src/runtime.ts
Normal file
139
.pi/agent/extensions/web-search/src/runtime.ts
Normal file
@@ -0,0 +1,139 @@
|
|||||||
|
import { loadWebSearchConfig, type ResolvedWebSearchConfig } from "./config.ts";
|
||||||
|
import { createExaProvider } from "./providers/exa.ts";
|
||||||
|
import { createTavilyProvider } from "./providers/tavily.ts";
|
||||||
|
import type {
|
||||||
|
NormalizedFetchRequest,
|
||||||
|
NormalizedFetchResponse,
|
||||||
|
NormalizedSearchRequest,
|
||||||
|
NormalizedSearchResponse,
|
||||||
|
WebProvider,
|
||||||
|
} from "./providers/types.ts";
|
||||||
|
import type { WebSearchProviderConfig } from "./schema.ts";
|
||||||
|
|
||||||
|
export interface ProviderExecutionMeta {
|
||||||
|
requestedProviderName?: string;
|
||||||
|
actualProviderName: string;
|
||||||
|
failoverFromProviderName?: string;
|
||||||
|
failoverReason?: string;
|
||||||
|
}
|
||||||
|
|
||||||
|
export interface RuntimeSearchResponse extends NormalizedSearchResponse {
|
||||||
|
execution: ProviderExecutionMeta;
|
||||||
|
}
|
||||||
|
|
||||||
|
export interface RuntimeFetchResponse extends NormalizedFetchResponse {
|
||||||
|
execution: ProviderExecutionMeta;
|
||||||
|
}
|
||||||
|
|
||||||
|
export function createWebSearchRuntime(
|
||||||
|
deps: {
|
||||||
|
loadConfig?: () => Promise<ResolvedWebSearchConfig>;
|
||||||
|
createProvider?: (providerConfig: WebSearchProviderConfig) => WebProvider;
|
||||||
|
} = {},
|
||||||
|
) {
|
||||||
|
const loadConfig = deps.loadConfig ?? loadWebSearchConfig;
|
||||||
|
const createProvider = deps.createProvider ?? ((providerConfig: WebSearchProviderConfig) => {
|
||||||
|
switch (providerConfig.type) {
|
||||||
|
case "tavily":
|
||||||
|
return createTavilyProvider(providerConfig);
|
||||||
|
case "exa":
|
||||||
|
return createExaProvider(providerConfig);
|
||||||
|
}
|
||||||
|
});
|
||||||
|
|
||||||
|
async function resolveConfigAndProvider(providerName?: string) {
|
||||||
|
const config = await loadConfig();
|
||||||
|
const selectedName = providerName ?? config.defaultProviderName;
|
||||||
|
const selectedConfig = config.providersByName.get(selectedName);
|
||||||
|
|
||||||
|
if (!selectedConfig) {
|
||||||
|
throw new Error(
|
||||||
|
`Unknown web-search provider \"${selectedName}\". Configured providers: ${[...config.providersByName.keys()].join(", ")}`,
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
return {
|
||||||
|
config,
|
||||||
|
selectedName,
|
||||||
|
selectedConfig,
|
||||||
|
selectedProvider: createProvider(selectedConfig),
|
||||||
|
};
|
||||||
|
}
|
||||||
|
|
||||||
|
async function search(request: NormalizedSearchRequest): Promise<RuntimeSearchResponse> {
|
||||||
|
const { config, selectedName, selectedConfig, selectedProvider } = await resolveConfigAndProvider(request.provider);
|
||||||
|
|
||||||
|
try {
|
||||||
|
const response = await selectedProvider.search(request);
|
||||||
|
return {
|
||||||
|
...response,
|
||||||
|
execution: {
|
||||||
|
requestedProviderName: request.provider,
|
||||||
|
actualProviderName: selectedName,
|
||||||
|
},
|
||||||
|
};
|
||||||
|
} catch (error) {
|
||||||
|
if (selectedConfig.type !== "tavily") {
|
||||||
|
throw error;
|
||||||
|
}
|
||||||
|
|
||||||
|
const fallbackConfig = [...config.providersByName.values()].find((provider) => provider.type === "exa");
|
||||||
|
if (!fallbackConfig) {
|
||||||
|
throw error;
|
||||||
|
}
|
||||||
|
|
||||||
|
const fallbackProvider = createProvider(fallbackConfig);
|
||||||
|
const fallbackResponse = await fallbackProvider.search({ ...request, provider: fallbackConfig.name });
|
||||||
|
return {
|
||||||
|
...fallbackResponse,
|
||||||
|
execution: {
|
||||||
|
requestedProviderName: request.provider,
|
||||||
|
actualProviderName: fallbackConfig.name,
|
||||||
|
failoverFromProviderName: selectedName,
|
||||||
|
failoverReason: (error as Error).message,
|
||||||
|
},
|
||||||
|
};
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
async function fetch(request: NormalizedFetchRequest): Promise<RuntimeFetchResponse> {
|
||||||
|
const { config, selectedName, selectedConfig, selectedProvider } = await resolveConfigAndProvider(request.provider);
|
||||||
|
|
||||||
|
try {
|
||||||
|
const response = await selectedProvider.fetch(request);
|
||||||
|
return {
|
||||||
|
...response,
|
||||||
|
execution: {
|
||||||
|
requestedProviderName: request.provider,
|
||||||
|
actualProviderName: selectedName,
|
||||||
|
},
|
||||||
|
};
|
||||||
|
} catch (error) {
|
||||||
|
if (selectedConfig.type !== "tavily") {
|
||||||
|
throw error;
|
||||||
|
}
|
||||||
|
|
||||||
|
const fallbackConfig = [...config.providersByName.values()].find((provider) => provider.type === "exa");
|
||||||
|
if (!fallbackConfig) {
|
||||||
|
throw error;
|
||||||
|
}
|
||||||
|
|
||||||
|
const fallbackProvider = createProvider(fallbackConfig);
|
||||||
|
const fallbackResponse = await fallbackProvider.fetch({ ...request, provider: fallbackConfig.name });
|
||||||
|
return {
|
||||||
|
...fallbackResponse,
|
||||||
|
execution: {
|
||||||
|
requestedProviderName: request.provider,
|
||||||
|
actualProviderName: fallbackConfig.name,
|
||||||
|
failoverFromProviderName: selectedName,
|
||||||
|
failoverReason: (error as Error).message,
|
||||||
|
},
|
||||||
|
};
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
return {
|
||||||
|
search,
|
||||||
|
fetch,
|
||||||
|
};
|
||||||
|
}
|
||||||
@@ -13,9 +13,43 @@ export const ExaProviderConfigSchema = Type.Object({
|
|||||||
options: Type.Optional(ProviderOptionsSchema),
|
options: Type.Optional(ProviderOptionsSchema),
|
||||||
});
|
});
|
||||||
|
|
||||||
|
export const TavilyProviderOptionsSchema = Type.Object({
|
||||||
|
defaultSearchLimit: Type.Optional(Type.Integer({ minimum: 1, maximum: 20 })),
|
||||||
|
defaultFetchTextMaxCharacters: Type.Optional(Type.Integer({ minimum: 1 })),
|
||||||
|
});
|
||||||
|
|
||||||
|
export const TavilyProviderConfigSchema = Type.Object({
|
||||||
|
name: Type.String({ minLength: 1 }),
|
||||||
|
type: Type.Literal("tavily"),
|
||||||
|
apiKey: Type.String({ minLength: 1 }),
|
||||||
|
options: Type.Optional(TavilyProviderOptionsSchema),
|
||||||
|
});
|
||||||
|
|
||||||
|
export const WebSearchProviderConfigSchema = Type.Union([ExaProviderConfigSchema, TavilyProviderConfigSchema]);
|
||||||
|
|
||||||
export const WebSearchConfigSchema = Type.Object({
|
export const WebSearchConfigSchema = Type.Object({
|
||||||
defaultProvider: Type.String({ minLength: 1 }),
|
defaultProvider: Type.String({ minLength: 1 }),
|
||||||
providers: Type.Array(ExaProviderConfigSchema, { minItems: 1 }),
|
providers: Type.Array(WebSearchProviderConfigSchema, { minItems: 1 }),
|
||||||
|
});
|
||||||
|
|
||||||
|
export const TavilySearchToolOptionsSchema = Type.Object({
|
||||||
|
searchDepth: Type.Optional(Type.String()),
|
||||||
|
topic: Type.Optional(Type.String()),
|
||||||
|
timeRange: Type.Optional(Type.String()),
|
||||||
|
days: Type.Optional(Type.Integer({ minimum: 1 })),
|
||||||
|
chunksPerSource: Type.Optional(Type.Integer({ minimum: 1 })),
|
||||||
|
includeAnswer: Type.Optional(Type.Boolean()),
|
||||||
|
includeRawContent: Type.Optional(Type.Boolean()),
|
||||||
|
includeImages: Type.Optional(Type.Boolean()),
|
||||||
|
});
|
||||||
|
|
||||||
|
export const TavilyFetchToolOptionsSchema = Type.Object({
|
||||||
|
query: Type.Optional(Type.String()),
|
||||||
|
extractDepth: Type.Optional(Type.String()),
|
||||||
|
chunksPerSource: Type.Optional(Type.Integer({ minimum: 1 })),
|
||||||
|
includeImages: Type.Optional(Type.Boolean()),
|
||||||
|
includeFavicon: Type.Optional(Type.Boolean()),
|
||||||
|
format: Type.Optional(Type.String()),
|
||||||
});
|
});
|
||||||
|
|
||||||
export const WebSearchParamsSchema = Type.Object({
|
export const WebSearchParamsSchema = Type.Object({
|
||||||
@@ -27,6 +61,7 @@ export const WebSearchParamsSchema = Type.Object({
|
|||||||
endPublishedDate: Type.Optional(Type.String()),
|
endPublishedDate: Type.Optional(Type.String()),
|
||||||
category: Type.Optional(Type.String()),
|
category: Type.Optional(Type.String()),
|
||||||
provider: Type.Optional(Type.String()),
|
provider: Type.Optional(Type.String()),
|
||||||
|
tavily: Type.Optional(TavilySearchToolOptionsSchema),
|
||||||
});
|
});
|
||||||
|
|
||||||
export const WebFetchParamsSchema = Type.Object({
|
export const WebFetchParamsSchema = Type.Object({
|
||||||
@@ -36,10 +71,16 @@ export const WebFetchParamsSchema = Type.Object({
|
|||||||
summary: Type.Optional(Type.Boolean()),
|
summary: Type.Optional(Type.Boolean()),
|
||||||
textMaxCharacters: Type.Optional(Type.Integer({ minimum: 1 })),
|
textMaxCharacters: Type.Optional(Type.Integer({ minimum: 1 })),
|
||||||
provider: Type.Optional(Type.String()),
|
provider: Type.Optional(Type.String()),
|
||||||
|
tavily: Type.Optional(TavilyFetchToolOptionsSchema),
|
||||||
});
|
});
|
||||||
|
|
||||||
export type ProviderOptions = Static<typeof ProviderOptionsSchema>;
|
export type ProviderOptions = Static<typeof ProviderOptionsSchema>;
|
||||||
|
export type TavilyProviderOptions = Static<typeof TavilyProviderOptionsSchema>;
|
||||||
export type ExaProviderConfig = Static<typeof ExaProviderConfigSchema>;
|
export type ExaProviderConfig = Static<typeof ExaProviderConfigSchema>;
|
||||||
|
export type TavilyProviderConfig = Static<typeof TavilyProviderConfigSchema>;
|
||||||
|
export type WebSearchProviderConfig = Static<typeof WebSearchProviderConfigSchema>;
|
||||||
export type WebSearchConfig = Static<typeof WebSearchConfigSchema>;
|
export type WebSearchConfig = Static<typeof WebSearchConfigSchema>;
|
||||||
|
export type TavilySearchToolOptions = Static<typeof TavilySearchToolOptionsSchema>;
|
||||||
|
export type TavilyFetchToolOptions = Static<typeof TavilyFetchToolOptionsSchema>;
|
||||||
export type WebSearchParams = Static<typeof WebSearchParamsSchema>;
|
export type WebSearchParams = Static<typeof WebSearchParamsSchema>;
|
||||||
export type WebFetchParams = Static<typeof WebFetchParamsSchema>;
|
export type WebFetchParams = Static<typeof WebFetchParamsSchema>;
|
||||||
|
|||||||
@@ -4,7 +4,7 @@ import { createWebFetchTool } from "./web-fetch.ts";
|
|||||||
|
|
||||||
test("web_fetch prepareArguments folds a single url into urls", () => {
|
test("web_fetch prepareArguments folds a single url into urls", () => {
|
||||||
const tool = createWebFetchTool({
|
const tool = createWebFetchTool({
|
||||||
resolveProvider: async () => {
|
executeFetch: async () => {
|
||||||
throw new Error("not used");
|
throw new Error("not used");
|
||||||
},
|
},
|
||||||
});
|
});
|
||||||
@@ -15,43 +15,51 @@ test("web_fetch prepareArguments folds a single url into urls", () => {
|
|||||||
});
|
});
|
||||||
});
|
});
|
||||||
|
|
||||||
test("web_fetch defaults to text and returns formatted fetch results", async () => {
|
test("web_fetch forwards nested Tavily extract options to the runtime", async () => {
|
||||||
let capturedRequest: Record<string, unknown> | undefined;
|
let capturedRequest: any;
|
||||||
|
|
||||||
const tool = createWebFetchTool({
|
const tool = createWebFetchTool({
|
||||||
resolveProvider: async () => ({
|
executeFetch: async (request) => {
|
||||||
name: "exa-main",
|
capturedRequest = request;
|
||||||
type: "exa",
|
return {
|
||||||
async search() {
|
providerName: "tavily-main",
|
||||||
throw new Error("not used");
|
results: [
|
||||||
},
|
{
|
||||||
async fetch(request) {
|
url: "https://pi.dev",
|
||||||
capturedRequest = request as unknown as Record<string, unknown>;
|
title: "Docs",
|
||||||
return {
|
text: "Body",
|
||||||
providerName: "exa-main",
|
},
|
||||||
results: [
|
],
|
||||||
{
|
execution: { actualProviderName: "tavily-main" },
|
||||||
url: "https://exa.ai/docs",
|
};
|
||||||
title: "Docs",
|
},
|
||||||
text: "Body",
|
|
||||||
},
|
|
||||||
],
|
|
||||||
};
|
|
||||||
},
|
|
||||||
}),
|
|
||||||
});
|
});
|
||||||
|
|
||||||
const result = await tool.execute("tool-1", { urls: ["https://exa.ai/docs"] }, undefined, undefined, undefined);
|
const result = await tool.execute(
|
||||||
|
"tool-1",
|
||||||
|
{
|
||||||
|
urls: ["https://pi.dev"],
|
||||||
|
tavily: {
|
||||||
|
query: "installation",
|
||||||
|
extractDepth: "advanced",
|
||||||
|
includeImages: true,
|
||||||
|
},
|
||||||
|
},
|
||||||
|
undefined,
|
||||||
|
undefined,
|
||||||
|
undefined,
|
||||||
|
);
|
||||||
|
|
||||||
assert.equal(capturedRequest?.text, true);
|
assert.equal(capturedRequest.tavily.query, "installation");
|
||||||
|
assert.equal(capturedRequest.tavily.extractDepth, "advanced");
|
||||||
|
assert.equal(capturedRequest.text, true);
|
||||||
assert.match((result.content[0] as { text: string }).text, /Body/);
|
assert.match((result.content[0] as { text: string }).text, /Body/);
|
||||||
assert.equal((result.details as { results: Array<{ title: string }> }).results[0]?.title, "Docs");
|
|
||||||
});
|
});
|
||||||
|
|
||||||
test("web_fetch rejects malformed URLs", async () => {
|
test("web_fetch rejects malformed URLs", async () => {
|
||||||
const tool = createWebFetchTool({
|
const tool = createWebFetchTool({
|
||||||
resolveProvider: async () => {
|
executeFetch: async () => {
|
||||||
throw new Error("should not resolve provider for invalid URLs");
|
throw new Error("should not execute fetch for invalid URLs");
|
||||||
},
|
},
|
||||||
});
|
});
|
||||||
|
|
||||||
|
|||||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user