invest-copilot/.pi/extensions/pi-loop-guard/index.ts

/**
 * 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;
	});
}