94 lines
2.6 KiB
TypeScript
94 lines
2.6 KiB
TypeScript
/**
|
|
* pi-loop-guard — Infinite loop prevention extension
|
|
*
|
|
* Counts tool calls across turns and aborts the agent when a configurable
|
|
* threshold is exceeded. This prevents the agent from looping indefinitely
|
|
* when errors repeat (e.g., database connection failures, missing API keys).
|
|
*
|
|
* Usage:
|
|
* pi install git:unraid.local:3003/pi-dev/pi-loop-guard
|
|
*
|
|
* Configuration (in .pi/settings.json or ~/.pi/agent/settings.json):
|
|
* {
|
|
* "loopGuard": {
|
|
* "maxToolCalls": 50,
|
|
* "maxTurns": 20
|
|
* }
|
|
* }
|
|
*/
|
|
|
|
import type { ExtensionAPI, ExtensionContext } from "@earendil-works/pi-coding-agent";
|
|
|
|
interface LoopGuardSettings {
|
|
maxToolCalls: number;
|
|
maxTurns: number;
|
|
}
|
|
|
|
const DEFAULT_SETTINGS: LoopGuardSettings = {
|
|
maxToolCalls: 50,
|
|
maxTurns: 20,
|
|
};
|
|
|
|
function readSettings(ctx: ExtensionContext): LoopGuardSettings {
|
|
try {
|
|
const raw = ctx.sessionManager.getEntries();
|
|
for (const entry of raw) {
|
|
if (entry.type === "custom" && entry.customType === "loopguard-settings") {
|
|
const s = entry.data as Partial<LoopGuardSettings>;
|
|
return {
|
|
maxToolCalls: s.maxToolCalls ?? DEFAULT_SETTINGS.maxToolCalls,
|
|
maxTurns: s.maxTurns ?? DEFAULT_SETTINGS.maxTurns,
|
|
};
|
|
}
|
|
}
|
|
} catch {
|
|
// Ignore errors reading settings
|
|
}
|
|
return { ...DEFAULT_SETTINGS };
|
|
}
|
|
|
|
export default function (pi: ExtensionAPI) {
|
|
let toolCallCount = 0;
|
|
let turnCount = 0;
|
|
let settings: LoopGuardSettings = { ...DEFAULT_SETTINGS };
|
|
|
|
// Load settings on session start
|
|
pi.on("session_start", async (_event, ctx) => {
|
|
settings = readSettings(ctx);
|
|
});
|
|
|
|
// Count tool calls and turns
|
|
pi.on("turn_end", async (_event, ctx) => {
|
|
const entries = ctx.sessionManager.getEntries();
|
|
const lastEntry = entries[entries.length - 1];
|
|
|
|
if (lastEntry?.type === "message" && lastEntry.message?.role === "assistant") {
|
|
const content = lastEntry.message.content as Array<{ type: string }>;
|
|
const toolCallsInTurn = content.filter((c) => c.type === "toolCall").length;
|
|
toolCallCount += toolCallsInTurn;
|
|
}
|
|
|
|
turnCount++;
|
|
|
|
if (toolCallCount >= settings.maxToolCalls) {
|
|
ctx.abort();
|
|
ctx.ui.notify(
|
|
`Loop guard: stopped after ${toolCallCount} tool calls (threshold: ${settings.maxToolCalls}). Possible infinite loop detected.`,
|
|
"error",
|
|
);
|
|
} else if (turnCount >= settings.maxTurns) {
|
|
ctx.abort();
|
|
ctx.ui.notify(
|
|
`Loop guard: stopped after ${turnCount} turns (threshold: ${settings.maxTurns}). Possible infinite loop detected.`,
|
|
"error",
|
|
);
|
|
}
|
|
});
|
|
|
|
// Reset counters on session shutdown
|
|
pi.on("session_shutdown", async () => {
|
|
toolCallCount = 0;
|
|
turnCount = 0;
|
|
});
|
|
}
|