fix(coding-agent): guard session context replay tail

2026-07-23 16:38:22 +08:00
parent 7504d4c24d
commit 6fae89aaea
5 changed files with 397 additions and 26 deletions
@@ -4,6 +4,7 @@
 ### Fixed

 - Fixed the setup wizard hiding the selected row on short terminals (e.g. 24x80): the provider sign-in, theme, and web-search lists now fit their windows to the visible height, and decorative chrome (sign-in hint, theme mock preview) yields to the list when space is tight.
+- Fixed restored sessions replaying terminal aborted or errored assistant turns, which could repeatedly fail continuation from an assistant role; `/retry` now consults the persisted transcript so the failed turn remains retryable without re-entering provider context.

 ## [17.0.8] - 2026-07-22

@@ -1881,6 +1881,22 @@ function titleConversationTurnFromMessage(message: AgentMessage): TitleConversat
 	return { role: message.role, ...(text ? { text } : {}), ...(thinking ? { thinking } : {}) };
 }

+function syntheticToolResultTailStart(messages: readonly AgentMessage[]): number {
+	let index = messages.length;
+	while (index > 0 && isSyntheticToolResultMessage(messages[index - 1])) {
+		index--;
+	}
+	return index;
+}
+
+function retryableAssistantTurnEnd(messages: readonly AgentMessage[]): number | undefined {
+	const turnEnd = syntheticToolResultTailStart(messages);
+	const message = messages[turnEnd - 1];
+	if (message?.role !== "assistant") return undefined;
+	if (message.stopReason !== "error" && message.stopReason !== "aborted") return undefined;
+	return turnEnd;
+}
+
 export class AgentSession {
 	readonly agent: Agent;
 	readonly sessionManager: SessionManager;
@@ -8951,12 +8967,16 @@ export class AgentSession {
 				);
 			}

-			// Check if we need to compact before sending (catches aborted responses). Run
-			// inline (allowDefer=false) so the handoff/maintenance fully settles before this
-			// prompt's agent loop starts — otherwise a deferred handoff would fire on the
-			// next microtask alongside the new turn.
+			// Recover a previously failed/incomplete assistant turn before sending.
+			// Successful historical turns take the cheaper pre-prompt threshold path
+			// below; re-running the full post-turn check on resume can synchronously
+			// rewrite/re-render old context before the new prompt starts.
 			const lastAssistant = this.#findLastAssistantMessage();
-			if (lastAssistant && !options?.skipCompactionCheck) {
+			if (
+				lastAssistant &&
+				!options?.skipCompactionCheck &&
+				(lastAssistant.stopReason === "error" || lastAssistant.stopReason === "length")
+			) {
 				await this.#checkCompaction(lastAssistant, false, false, false);
 			}

@@ -15780,7 +15800,8 @@ export class AgentSession {
 	}
 	/**
 	 * Manually retry the last failed assistant turn.
-	 * Removes the error message from agent state and re-attempts with a fresh retry budget.
+	 * Removes the error message from active agent state when present and
+	 * re-attempts with a fresh retry budget.
 	 *
 	 * A stream that stalls or aborts mid-tool-call ends the turn with
 	 * `stopReason: "error" | "aborted"` and then appends one synthetic
@@ -15791,31 +15812,29 @@ export class AgentSession {
 	 * checking the assistant message; it strips both the placeholders and the
 	 * failed turn before re-attempting.
 	 *
+	 * A restored session deliberately omits failed assistant turns from provider
+	 * context. In that case, the persisted display transcript remains the source
+	 * of truth for whether the current branch has a retryable failed tail.
+	 *
 	 * @returns true if retry was initiated, false if no failed turn to retry or agent is busy
 	 */
 	async retry(): Promise<boolean> {
 		if (this.isStreaming || this.isCompacting || this.isRetrying) return false;

 		const messages = this.agent.state.messages;
-
-		// Walk back past trailing synthetic tool_result placeholders emitted for
-		// tool calls that never ran because the turn stalled/aborted mid-tool-call.
-		// They shadow the failed assistant turn from the single-message lookback.
-		let turnEnd = messages.length;
-		while (turnEnd > 0 && isSyntheticToolResultMessage(messages[turnEnd - 1])) {
-			turnEnd--;
+		const activeTurnEnd = retryableAssistantTurnEnd(messages);
+		if (activeTurnEnd !== undefined) {
+			// Remove the failed/aborted assistant message plus its synthetic tool
+			// results (same as auto-retry does before re-attempting).
+			this.agent.replaceMessages(messages.slice(0, activeTurnEnd - 1));
+		} else {
+			// A restored session already dropped the failed assistant turn (and its
+			// paired synthetic tool results) from provider context, so the persisted
+			// display transcript is the source of truth for a retryable failed tail.
+			const transcriptMessages = this.sessionManager.buildSessionContext({ transcript: true }).messages;
+			if (retryableAssistantTurnEnd(transcriptMessages) === undefined) return false;
 		}

-		const lastMsg = messages[turnEnd - 1];
-		if (lastMsg?.role !== "assistant") return false;
-
-		const assistantMsg = lastMsg as AssistantMessage;
-		if (assistantMsg.stopReason !== "error" && assistantMsg.stopReason !== "aborted") return false;
-
-		// Remove the failed/aborted assistant message plus its synthetic tool
-		// results (same as auto-retry does before re-attempting).
-		this.agent.replaceMessages(messages.slice(0, turnEnd - 1));
-
 		// Reset retry budget for a fresh attempt
 		this.#retryAttempt = 0;

@@ -1,7 +1,8 @@
 import { describe, expect, it } from "bun:test";
 import type { AgentMessage } from "@oh-my-pi/pi-agent-core";
+import type { AssistantMessage } from "@oh-my-pi/pi-ai";
 import * as snapcompact from "@oh-my-pi/snapcompact";
-import type { CompactionSummaryMessage } from "./messages";
+import { type CompactionSummaryMessage, INTERRUPTED_THINKING_MESSAGE_TYPE } from "./messages";
 import { buildSessionContext, type StrippedToolCallsMarker } from "./session-context";
 import type { SessionEntry } from "./session-entries";

@@ -159,3 +160,225 @@ describe("buildSessionContext dangling toolCalls", () => {
 		expect(context.messages.some(message => message.role === "assistant")).toBe(false);
 	});
 });
+
+const assistantUsage: AssistantMessage["usage"] = {
+	input: 0,
+	output: 0,
+	cacheRead: 0,
+	cacheWrite: 0,
+	totalTokens: 0,
+	cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
+};
+
+function userEntry(id: string, parentId: string | null, content: string, messageTimestamp: number): SessionEntry {
+	return {
+		type: "message",
+		id,
+		parentId,
+		timestamp,
+		message: { role: "user", content, timestamp: messageTimestamp } as AgentMessage,
+	};
+}
+
+function assistantEntry(
+	id: string,
+	parentId: string | null,
+	stopReason: AssistantMessage["stopReason"],
+	text: string,
+	messageTimestamp: number,
+): SessionEntry {
+	return {
+		type: "message",
+		id,
+		parentId,
+		timestamp,
+		message: {
+			role: "assistant",
+			content: [{ type: "text", text }],
+			api: "anthropic-messages",
+			provider: "anthropic",
+			model: "claude-sonnet-4-5",
+			usage: assistantUsage,
+			stopReason,
+			timestamp: messageTimestamp,
+		} satisfies AssistantMessage,
+	};
+}
+
+function toolCallAssistantEntry(
+	id: string,
+	parentId: string | null,
+	stopReason: AssistantMessage["stopReason"],
+	toolCallId: string,
+	messageTimestamp: number,
+): SessionEntry {
+	return {
+		type: "message",
+		id,
+		parentId,
+		timestamp,
+		message: {
+			role: "assistant",
+			content: [{ type: "toolCall", id: toolCallId, name: "write", arguments: { path: "plan.md", content: "x" } }],
+			api: "anthropic-messages",
+			provider: "anthropic",
+			model: "claude-sonnet-4-5",
+			usage: assistantUsage,
+			stopReason,
+			timestamp: messageTimestamp,
+		} satisfies AssistantMessage,
+	};
+}
+
+function syntheticToolResultEntry(
+	id: string,
+	parentId: string | null,
+	toolCallId: string,
+	messageTimestamp: number,
+): SessionEntry {
+	return {
+		type: "message",
+		id,
+		parentId,
+		timestamp,
+		message: {
+			role: "toolResult",
+			toolCallId,
+			toolName: "write",
+			content: [
+				{ type: "text", text: "Tool call was not executed because the provider stream ended with an error." },
+			],
+			details: { __synthetic: true, source: "assistant_stop_error", executed: false },
+			isError: true,
+			timestamp: messageTimestamp,
+		} as AgentMessage,
+	};
+}
+
+function hiddenContinuityEntry(id: string, parentId: string | null): SessionEntry {
+	return {
+		type: "custom_message",
+		id,
+		parentId,
+		timestamp,
+		customType: INTERRUPTED_THINKING_MESSAGE_TYPE,
+		content: "preserved interrupted thinking",
+		display: false,
+		attribution: "agent",
+	};
+}
+
+function expectUserTail(messages: AgentMessage[], content: string): void {
+	const tail = messages.at(-1);
+	expect(tail?.role).toBe("user");
+	if (tail?.role !== "user") {
+		throw new Error(`Expected user tail, received ${tail?.role ?? "none"}`);
+	}
+	expect(tail.content).toBe(content);
+}
+
+describe("buildSessionContext failed replay tails", () => {
+	it("terminates on cyclic parent links and includes each reachable message once", () => {
+		const entries = [userEntry("A", "B", "from A", 1), userEntry("B", "A", "from B", 2)];
+
+		const context = buildSessionContext(entries, "A");
+
+		expect(context.messages.map(message => (message.role === "user" ? message.content : message.role))).toEqual([
+			"from B",
+			"from A",
+		]);
+	});
+
+	it("omits a terminal aborted assistant from normal context", () => {
+		const context = buildSessionContext([
+			userEntry("user", null, "continue", 1),
+			assistantEntry("assistant", "user", "aborted", "partial unsafe replay", 2),
+		]);
+
+		expect(context.messages.some(message => message.role === "assistant")).toBe(false);
+		expectUserTail(context.messages, "continue");
+	});
+
+	it("omits an earlier aborted assistant before a later user from normal context", () => {
+		const context = buildSessionContext([
+			userEntry("user-1", null, "first prompt", 1),
+			assistantEntry("assistant", "user-1", "aborted", "partial unsafe replay", 2),
+			userEntry("user-2", "assistant", "retry", 3),
+		]);
+
+		expect(context.messages.some(message => message.role === "assistant")).toBe(false);
+		expectUserTail(context.messages, "retry");
+	});
+
+	it("preserves a terminal aborted assistant in transcript mode", () => {
+		const context = buildSessionContext(
+			[
+				userEntry("user", null, "continue", 1),
+				assistantEntry("assistant", "user", "aborted", "visible transcript error", 2),
+			],
+			undefined,
+			undefined,
+			{ transcript: true },
+		);
+
+		const assistant = context.messages.find(message => message.role === "assistant");
+		expect(assistant?.role).toBe("assistant");
+		if (assistant?.role !== "assistant") {
+			throw new Error(`Expected transcript assistant, received ${assistant?.role ?? "none"}`);
+		}
+		expect(assistant.stopReason).toBe("aborted");
+		expect(assistant.content).toEqual([{ type: "text", text: "visible transcript error" }]);
+	});
+
+	it("omits a terminal error assistant from normal context", () => {
+		const context = buildSessionContext([
+			userEntry("user", null, "retry with smaller input", 1),
+			assistantEntry("assistant", "user", "error", "provider rejected the request", 2),
+		]);
+
+		expect(context.messages.some(message => message.role === "assistant")).toBe(false);
+		expectUserTail(context.messages, "retry with smaller input");
+	});
+
+	it("keeps an aborted assistant when hidden interrupted-thinking continuity follows it", () => {
+		const context = buildSessionContext([
+			userEntry("user", null, "keep reasoning continuity", 1),
+			assistantEntry("assistant", "user", "aborted", "partial answer before interrupt", 2),
+			hiddenContinuityEntry("continuity", "assistant"),
+		]);
+
+		const assistant = context.messages.find(message => message.role === "assistant");
+		expect(assistant?.role).toBe("assistant");
+		if (assistant?.role !== "assistant") {
+			throw new Error(`Expected assistant before continuity, received ${assistant?.role ?? "none"}`);
+		}
+		expect(assistant.stopReason).toBe("aborted");
+		expect(context.messages.at(-1)?.role).toBe("custom");
+	});
+
+	it("drops synthetic tool results paired with a dropped failed tool-call turn", () => {
+		const context = buildSessionContext([
+			userEntry("user", null, "write the plan", 1),
+			toolCallAssistantEntry("assistant", "user", "error", "call-1", 2),
+			syntheticToolResultEntry("result", "assistant", "call-1", 3),
+		]);
+
+		expect(context.messages.map(message => message.role)).toEqual(["user"]);
+		expectUserTail(context.messages, "write the plan");
+	});
+
+	it("keeps the failed tool-call turn and its result in transcript mode", () => {
+		const context = buildSessionContext(
+			[
+				userEntry("user", null, "write the plan", 1),
+				toolCallAssistantEntry("assistant", "user", "error", "call-1", 2),
+				syntheticToolResultEntry("result", "assistant", "call-1", 3),
+			],
+			undefined,
+			undefined,
+			{ transcript: true, keepDanglingToolCalls: true },
+		);
+
+		expect(context.messages.map(message => message.role)).toEqual(["user", "assistant", "toolResult"]);
+	});
+});
@@ -5,6 +5,7 @@ import {
 	createBranchSummaryMessage,
 	createCompactionSummaryMessage,
 	createCustomMessage,
+	INTERRUPTED_THINKING_MESSAGE_TYPE,
 	isCustomMessageContent,
 	normalizeCustomMessagePayload,
 } from "./messages";
@@ -197,10 +198,13 @@ export function buildSessionContext(
 		};
 	}

-	// Walk from leaf to root, collecting path
+	// Walk from leaf to root, collecting path. Corrupt/pre-fix files can contain
+	// parent cycles; stop at the first repeat so session load is bounded.
 	const path: SessionEntry[] = [];
+	const seenPathIds = new Set<string>();
 	let current: SessionEntry | undefined = leaf;
-	while (current) {
+	while (current && !seenPathIds.has(current.id)) {
+		seenPathIds.add(current.id);
 		path.push(current);
 		current = current.parentId ? byId.get(current.parentId) : undefined;
 	}
@@ -498,6 +502,41 @@ export function buildSessionContext(
 		}
 	}

+	// Error/abort assistant turns are transcript events, not safe assistant
+	// turns to replay into the next provider request. Drop them even when a
+	// later user message follows through non-context entries (`session_exit`,
+	// labels, etc.); otherwise a resumed session replays a dead partial turn
+	// and can spend minutes reprocessing old context before the new prompt.
+	// Keep the interrupted-thinking continuity pair: convertToLlm strips the
+	// unsafe trailing thinking from that assistant and sends the hidden
+	// continuity note instead.
+	if (!options?.transcript) {
+		for (let i = messages.length - 1; i >= 0; i--) {
+			const message = messages[i];
+			if (message?.role !== "assistant") continue;
+			if (message.stopReason !== "aborted" && message.stopReason !== "error") continue;
+			const next = messages[i + 1];
+			if (next?.role === "custom" && next.customType === INTERRUPTED_THINKING_MESSAGE_TYPE) continue;
+			// A failed turn that emitted tool calls persists paired synthetic
+			// tool_result placeholders after it. Dropping only the assistant would
+			// strand those results with no preceding tool_use — a shape providers
+			// reject — so remove the paired results alongside the turn.
+			const droppedToolCallIds = new Set<string>();
+			for (const block of message.content) {
+				if (block.type === "toolCall") droppedToolCallIds.add(block.id);
+			}
+			messages.splice(i, 1);
+			if (droppedToolCallIds.size > 0) {
+				for (let j = messages.length - 1; j >= i; j--) {
+					const candidate = messages[j];
+					if (candidate?.role === "toolResult" && droppedToolCallIds.has(candidate.toolCallId)) {
+						messages.splice(j, 1);
+					}
+				}
+			}
+		}
+	}
+
 	return {
 		messages,
 		cacheMissExplainedAt: options?.transcript ? cacheMissExplainedAt : undefined,
@@ -173,4 +173,93 @@ describe("AgentSession manual retry", () => {
 			text: "recovered after stalled tool call",
 		});
 	});
+
+	it("retries a persisted failed turn after rebuilding provider context", async () => {
+		const model = getBundledModel("anthropic", "claude-sonnet-4-5");
+		if (!model) {
+			throw new Error("Expected bundled Anthropic test model to exist");
+		}
+
+		const mock = createMockModel({
+			responses: [
+				{
+					content: [{ type: "toolCall", name: "write", arguments: { path: "plan.md", content: "x" } }],
+					stopReason: "error",
+					errorMessage: "stream stalled before the tool ran",
+				},
+				{ content: ["recovered after session reopen"], stopReason: "stop" },
+			],
+		});
+		const sessionManager = SessionManager.inMemory();
+		const agent = new Agent({
+			getApiKey: model => `${model.provider}-test-key`,
+			initialState: {
+				model,
+				systemPrompt: ["Test"],
+				tools: [],
+				messages: [],
+			},
+			streamFn: mock.stream,
+		});
+		session = new AgentSession({
+			agent,
+			sessionManager,
+			settings: Settings.isolated({ "compaction.enabled": false, "retry.enabled": false }),
+			modelRegistry: new ModelRegistry(authStorage),
+		});
+		session.subscribe(() => {});
+
+		await session.prompt("write before reopen");
+		await session.waitForIdle();
+		const failedAssistant = session.agent.state.messages.findLast(
+			(message): message is AssistantMessage => message.role === "assistant",
+		);
+		expect(failedAssistant?.stopReason).toBe("error");
+
+		const reopenedManager = SessionManager.inMemory();
+		reopenedManager.restoreState(sessionManager.captureState());
+		await session.dispose();
+		session = undefined;
+
+		const restoredMessages = reopenedManager.buildSessionContext().messages;
+		// The failed tool-call turn AND its paired synthetic tool result are both
+		// dropped from provider context — leaving a stranded tool result with no
+		// preceding tool_use would be rejected by provider converters.
+		expect(restoredMessages.map(message => message.role)).toEqual(["user"]);
+		const transcriptMessages = reopenedManager.buildSessionContext({ transcript: true }).messages;
+		expect(transcriptMessages.at(-1)?.role).toBe("toolResult");
+		const transcriptAssistant = transcriptMessages.findLast(
+			(message): message is AssistantMessage => message.role === "assistant",
+		);
+		expect(transcriptAssistant?.stopReason).toBe("error");
+
+		const reopenedAgent = new Agent({
+			getApiKey: model => `${model.provider}-test-key`,
+			initialState: {
+				model,
+				systemPrompt: ["Test"],
+				tools: [],
+				messages: restoredMessages,
+			},
+			streamFn: mock.stream,
+		});
+		session = new AgentSession({
+			agent: reopenedAgent,
+			sessionManager: reopenedManager,
+			settings: Settings.isolated({ "compaction.enabled": false, "retry.enabled": false }),
+			modelRegistry: new ModelRegistry(authStorage),
+		});
+		session.subscribe(() => {});
+
+		await expect(session.retry()).resolves.toBe(true);
+		await session.waitForIdle();
+		expect(session.agent.state.messages.map(message => message.role)).toEqual(["user", "assistant"]);
+
+		expect(mock.calls.length).toBe(2);
+		expect(lastAgentMessage(session).stopReason).toBe("stop");
+		expect(lastAgentMessage(session).content).toContainEqual({
+			type: "text",
+			text: "recovered after session reopen",
+		});
+	});
 });