pi repair toolcall

toolcall-repair.ts

// for dum dum models ie qwen3.5 9b
import type { ExtensionAPI } from "@mariozechner/pi-coding-agent";
import type { AssistantMessage } from "@mariozechner/pi-ai";

type RepairMode = "repair" | "continue";

interface RepairState {
	enabled: boolean;
	mode: RepairMode;
	consecutiveRepairs: number;
	maxConsecutiveRepairs: number;
	lastHandledFingerprint?: string;
}

function isAssistantMessage(message: unknown): message is AssistantMessage {
	return !!message && typeof message === "object" && (message as AssistantMessage).role === "assistant";
}

function detectMalformedXmlToolCall(message: AssistantMessage): { snippet: string; fingerprint: string } | null {
	if (!Array.isArray(message.content)) return null;

	const hasStructuredToolCall = message.content.some((block) => block?.type === "toolCall");
	if (hasStructuredToolCall) return null;

	const text = message.content
		.filter((block) => block?.type === "thinking" || block?.type === "text")
		.map((block) => {
			if (block.type === "thinking") return block.thinking || "";
			if (block.type === "text") return block.text || "";
			return "";
		})
		.join("\n");

	if (!text) return null;

	const hasXmlToolCall = /<tool_call>/i.test(text) && /<\/tool_call>/i.test(text) && /<function=/i.test(text);
	if (!hasXmlToolCall) return null;

	const match = text.match(/<tool_call>[\s\S]*?<\/tool_call>/i);
	const snippet = (match?.[0] || text).slice(0, 1000);
	const fingerprint = `${message.timestamp}:${snippet}`;

	return { snippet, fingerprint };
}

function getFollowUp(mode: RepairMode): string {
	if (mode === "continue") return "continue";

	return [
		"Your previous response emitted a literal <tool_call> block in text/thinking.",
		"Re-emit the exact intended action as a proper structured tool call now.",
		"Output only the tool call.",
	].join(" ");
}

export default function (pi: ExtensionAPI) {
	const state: RepairState = {
		enabled: true,
		mode: "repair",
		consecutiveRepairs: 0,
		maxConsecutiveRepairs: 3,
	};

	pi.on("session_start", async (_event, _ctx) => {
		// Always start enabled for each session/runtime load.
		state.enabled = true;
		state.consecutiveRepairs = 0;
		state.lastHandledFingerprint = undefined;
	});

	pi.registerCommand("toolcall-repair", {
		description: "Configure malformed <tool_call> auto-recovery (status|on|off|mode repair|mode continue)",
		handler: async (args, ctx) => {
			const input = args.trim().toLowerCase();

			if (!input || input === "status") {
				ctx.ui.notify(
					`toolcall-repair: ${state.enabled ? "on" : "off"}, mode=${state.mode}, maxConsecutive=${state.maxConsecutiveRepairs}`,
					"info",
				);
				return;
			}

			if (input === "on") {
				state.enabled = true;
				ctx.ui.notify("toolcall-repair enabled", "info");
				return;
			}

			if (input === "off") {
				state.enabled = false;
				ctx.ui.notify("toolcall-repair disabled", "info");
				return;
			}

			if (input === "mode repair") {
				state.mode = "repair";
				ctx.ui.notify("toolcall-repair mode set to repair", "info");
				return;
			}

			if (input === "mode continue") {
				state.mode = "continue";
				ctx.ui.notify("toolcall-repair mode set to continue", "info");
				return;
			}

			ctx.ui.notify("Usage: /toolcall-repair [status|on|off|mode repair|mode continue]", "warning");
		},
	});

	pi.on("message_end", async (event, ctx) => {
		const msg = event.message;
		if (!isAssistantMessage(msg)) return;

		if (!state.enabled) return;

		const malformed = detectMalformedXmlToolCall(msg);

		if (!malformed) {
			// Reset backoff after a normal assistant response.
			state.consecutiveRepairs = 0;
			state.lastHandledFingerprint = undefined;
			return;
		}

		// Only repair completed stop responses; avoid interfering with toolUse/error/aborted turns.
		if (msg.stopReason !== "stop") return;

		if (state.lastHandledFingerprint === malformed.fingerprint) return;

		if (state.consecutiveRepairs >= state.maxConsecutiveRepairs) {
			ctx.ui.notify(
				`toolcall-repair: hit max consecutive repairs (${state.maxConsecutiveRepairs}); waiting for manual input`,
				"warning",
			);
			return;
		}

		state.consecutiveRepairs += 1;
		state.lastHandledFingerprint = malformed.fingerprint;

		const followUp = getFollowUp(state.mode);
		if (ctx.isIdle()) {
			pi.sendUserMessage(followUp);
		} else {
			pi.sendUserMessage(followUp, { deliverAs: "followUp" });
		}

		ctx.ui.notify(`toolcall-repair: malformed <tool_call> detected, sent auto-${state.mode}`, "warning");
	});
}