feat(coding-agent): added Gemini reasoning runaway detection and interruption

- Implemented a monitoring system to identify Gemini model runaway behavior during thinking steps using consecutive header detection. - Introduced a configurable tool-call reminder prompt to inject corrective context when reasoning stalls. - Added session-level logic to automatically interrupt and prune stalled assistant turns from the conversation history. - Provided comprehensive test coverage for the detection logic and stream interruption scenarios.
2026-06-27 01:01:02 +02:00
parent eb39fc4a50
commit 71e044cfe1
8 changed files with 597 additions and 25 deletions
@@ -0,0 +1,252 @@
+import { afterEach, beforeEach, describe, expect, it, vi } from "bun:test";
+import * as path from "node:path";
+import { Agent } from "@oh-my-pi/pi-agent-core";
+import type {
+	Api,
+	AssistantMessage,
+	Context,
+	Message,
+	Model,
+	SimpleStreamOptions,
+	ThinkingContent,
+} from "@oh-my-pi/pi-ai";
+import { createMockModel } from "@oh-my-pi/pi-ai/providers/mock";
+import { AssistantMessageEventStream } from "@oh-my-pi/pi-ai/utils/event-stream";
+import { GEMINI_HEADER_RUNAWAY_THRESHOLD } from "@oh-my-pi/pi-ai/utils/thinking-loop";
+import { ModelRegistry } from "@oh-my-pi/pi-coding-agent/config/model-registry";
+import { Settings } from "@oh-my-pi/pi-coding-agent/config/settings";
+import { AgentSession, type AgentSessionEvent } from "@oh-my-pi/pi-coding-agent/session/agent-session";
+import { AuthStorage } from "@oh-my-pi/pi-coding-agent/session/auth-storage";
+import { convertToLlm } from "@oh-my-pi/pi-coding-agent/session/messages";
+import { SessionManager } from "@oh-my-pi/pi-coding-agent/session/session-manager";
+import { TempDir } from "@oh-my-pi/pi-utils";
+
+function emptyUsage(): AssistantMessage["usage"] {
+	return {
+		input: 0,
+		output: 0,
+		cacheRead: 0,
+		cacheWrite: 0,
+		totalTokens: 0,
+		cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
+	};
+}
+
+/** Concatenate the text of a developer/user/custom LLM message. */
+function messageText(message: Message): string {
+	if (typeof message.content === "string") return message.content;
+	let text = "";
+	for (const block of message.content) {
+		if (block.type === "text") text += block.text;
+	}
+	return text;
+}
+
+/**
+ * First-call stream: a genuinely-distinct planning runaway — each thought summary
+ * has a fresh title + a paragraph naming new code anchors, so the similarity loop
+ * guard never fires; only the header-count guard catches it. Mirrors
+ * `streaming-edit-abort`: an abort listener pushes the terminal `aborted` event,
+ * and deltas are spaced with `Bun.sleep(0)` so the interceptor's `agent.abort()`
+ * lands before the turn would otherwise finish `stop`.
+ *
+ * When `finalText` is provided, a non-interrupted run ends with visible prose so
+ * the empty-stop handler does not auto-retry; used to prove the reminder setting
+ * alone controls this feature.
+ */
+function headerRunawayStream(
+	model: Model<Api>,
+	options?: SimpleStreamOptions,
+	finalText?: string,
+): AssistantMessageEventStream {
+	const stream = new AssistantMessageEventStream();
+	const thinking: ThinkingContent = { type: "thinking", thinking: "" };
+	const timestamp = Date.now();
+	const partial: AssistantMessage = {
+		role: "assistant",
+		content: [thinking],
+		api: model.api,
+		provider: model.provider,
+		model: model.id,
+		usage: emptyUsage(),
+		stopReason: "stop",
+		timestamp,
+	};
+	let aborted = false;
+	options?.signal?.addEventListener(
+		"abort",
+		() => {
+			if (aborted) return;
+			aborted = true;
+			stream.push({
+				type: "error",
+				reason: "aborted",
+				error: { ...partial, content: [{ ...thinking }], stopReason: "aborted" },
+			});
+		},
+		{ once: true },
+	);
+
+	void (async () => {
+		stream.push({ type: "start", partial });
+		stream.push({ type: "thinking_start", contentIndex: 0, partial });
+		for (let i = 0; i < GEMINI_HEADER_RUNAWAY_THRESHOLD + 2; i++) {
+			if (aborted) return;
+			const delta = `**Refining Stage ${i}**\n\nReworking module_${i} so handler_${i} routes Stage${i}Result through render_${i}.\n\n`;
+			thinking.thinking += delta;
+			stream.push({ type: "thinking_delta", contentIndex: 0, delta, partial });
+			await Bun.sleep(0);
+		}
+		if (aborted) return;
+		stream.push({ type: "thinking_end", contentIndex: 0, content: thinking.thinking, partial });
+		if (!finalText) {
+			stream.push({ type: "done", reason: "stop", message: partial });
+			return;
+		}
+		const finalMessage: AssistantMessage = {
+			...partial,
+			content: [{ ...thinking }, { type: "text", text: finalText }],
+		};
+		stream.push({ type: "text_start", contentIndex: 1, partial: finalMessage });
+		stream.push({ type: "text_delta", contentIndex: 1, delta: finalText, partial: finalMessage });
+		stream.push({ type: "text_end", contentIndex: 1, content: finalText, partial: finalMessage });
+		stream.push({ type: "done", reason: "stop", message: finalMessage });
+	})();
+	return stream;
+}
+
+function successStream(model: Model<Api>, text: string): AssistantMessageEventStream {
+	const stream = new AssistantMessageEventStream();
+	queueMicrotask(() => {
+		const message: AssistantMessage = {
+			role: "assistant",
+			content: [{ type: "text", text }],
+			api: model.api,
+			provider: model.provider,
+			model: model.id,
+			usage: emptyUsage(),
+			stopReason: "stop",
+			timestamp: Date.now(),
+		};
+		stream.push({ type: "start", partial: message });
+		stream.push({ type: "text_start", contentIndex: 0, partial: message });
+		stream.push({ type: "text_delta", contentIndex: 0, delta: text, partial: message });
+		stream.push({ type: "text_end", contentIndex: 0, content: text, partial: message });
+		stream.push({ type: "done", reason: "stop", message });
+	});
+	return stream;
+}
+
+describe("AgentSession Gemini header-runaway interrupt", () => {
+	let tempDir: TempDir;
+	let authStorage: AuthStorage;
+	let session: AgentSession | undefined;
+
+	beforeEach(async () => {
+		tempDir = TempDir.createSync("@pi-gemini-header-interrupt-");
+		authStorage = await AuthStorage.create(path.join(tempDir.path(), "auth.db"));
+		authStorage.setRuntimeApiKey("openrouter", "openrouter-test-key");
+	});
+
+	afterEach(async () => {
+		if (session) {
+			await session.dispose();
+			session = undefined;
+		}
+		authStorage.close();
+		tempDir.removeSync();
+		vi.restoreAllMocks();
+	});
+
+	function buildSession(streamFn: Agent["streamFn"], overrides?: Record<string, unknown>): void {
+		const model = createMockModel({ provider: "openrouter", id: "google/gemini-3.5-flash" }).model;
+		const modelRegistry = new ModelRegistry(authStorage);
+		const agent = new Agent({
+			getApiKey: requestedModel => `${requestedModel.provider}-test-key`,
+			initialState: { model, systemPrompt: ["Test"], tools: [], messages: [] },
+			streamFn,
+			convertToLlm,
+		});
+		const settings = Settings.isolated({
+			"compaction.enabled": false,
+			"retry.enabled": false,
+			"todo.enabled": false,
+			"advisor.enabled": false,
+			"model.loopGuard.enabled": true,
+			"model.loopGuard.toolCallReminder": true,
+			...overrides,
+		});
+		settings.setModelRole("default", `${model.provider}/${model.id}`);
+		session = new AgentSession({ agent, sessionManager: SessionManager.inMemory(), settings, modelRegistry });
+	}
+
+	it("interrupts the reasoning runaway, injects a tool-call reminder, and continues", async () => {
+		const contexts: Context[] = [];
+		let call = 0;
+		buildSession((model, context, options) => {
+			contexts.push(context);
+			call++;
+			return call === 1 ? headerRunawayStream(model, options) : successStream(model, "Acted: called a tool.");
+		});
+		const notices: Array<Extract<AgentSessionEvent, { type: "notice" }>> = [];
+		session?.subscribe(event => {
+			if (event.type === "notice") notices.push(event);
+		});
+
+		await session?.prompt("Do the task");
+		await session?.waitForIdle();
+
+		// The runaway was interrupted and the turn was re-driven.
+		expect(call).toBe(2);
+
+		// The user saw a transparency notice from the loop guard.
+		const guardNotice = notices.find(n => n.source === "loop-guard");
+		expect(guardNotice).toBeDefined();
+		// The continuation carried the hidden tool-call reminder (custom -> developer).
+		const reminderInContext = contexts[1].messages.some(
+			m =>
+				m.role === "developer" &&
+				/consecutive planning headers/.test(messageText(m)) &&
+				/tool call/.test(messageText(m)),
+		);
+		expect(reminderInContext).toBe(true);
+		// It names the header count that tripped the guard.
+		const reminderText = contexts[1].messages.map(messageText).join("\n");
+		expect(reminderText).toContain(String(GEMINI_HEADER_RUNAWAY_THRESHOLD));
+
+		// The stalled reasoning-only turn was discarded (not replayed as loop fuel).
+		const messages = session?.agent.state.messages ?? [];
+		const assistants = messages.filter((m): m is AssistantMessage => m.role === "assistant");
+		expect(assistants).toHaveLength(1);
+		expect(assistants[0].content).toEqual([{ type: "text", text: "Acted: called a tool." }]);
+		const replaysHeaders = messages.some(m => m.role === "assistant" && /Refining Stage/.test(messageText(m)));
+		expect(replaysHeaders).toBe(false);
+	});
+
+	it("does not interrupt when the tool-call reminder setting is off", async () => {
+		let call = 0;
+		buildSession(
+			(model, _context, options) => {
+				call++;
+				return headerRunawayStream(model, options, "Visible final answer.");
+			},
+			{ "model.loopGuard.toolCallReminder": false },
+		);
+		const notices: Array<Extract<AgentSessionEvent, { type: "notice" }>> = [];
+		session?.subscribe(event => {
+			if (event.type === "notice") notices.push(event);
+		});
+
+		await session?.prompt("Do the task");
+		await session?.waitForIdle();
+
+		expect(call).toBe(1);
+		expect(notices.some(n => n.source === "loop-guard")).toBe(false);
+		const messages = session?.agent.state.messages ?? [];
+		const reminderInjected = messages.some(m => m.role === "custom" && m.customType === "gemini-tool-call-reminder");
+		expect(reminderInjected).toBe(false);
+		const assistants = messages.filter((m): m is AssistantMessage => m.role === "assistant");
+		expect(assistants).toHaveLength(1);
+		expect(assistants[0].content.at(-1)).toEqual({ type: "text", text: "Visible final answer." });
+	});
+});