feat: add retry for failed message (conwnet#757)

Author: conwnet

extensions/github1s-ai/assets/chat.css

@@ -1144,6 +1144,16 @@ summary:focus-visible,
 	font: inherit;
 }
 
+.message-retry {
+	gap: 4px;
+	margin-top: 6px;
+}
+
+.message-retry .icon {
+	width: 12px;
+	height: 12px;
+}
+
 .message-status {
 	margin: 5px 0 0;
 	color: var(--vscode-descriptionForeground);

extensions/github1s-ai/src/common/conversation.ts

@@ -34,6 +34,22 @@ export interface Conversation extends ConversationSummary {
 	messages: ConversationMessage[];
 }
 
+export const isInterruptedMessage = (message: ConversationMessage): boolean =>
+	message.metadata.status === 'failed' ||
+	message.metadata.status === 'aborted' ||
+	message.metadata.status === 'unknown';
+
+export const getRetryableMessage = (messages: readonly ConversationMessage[]): ConversationMessage | undefined => {
+	const assistant = messages.at(-1);
+	const user = messages.at(-2);
+	return assistant?.role === 'assistant' &&
+		isInterruptedMessage(assistant) &&
+		user?.role === 'user' &&
+		user.metadata.turnId === assistant.metadata.turnId
+		? assistant
+		: undefined;
+};
+
 export const createUserMessage = (
 	turnId: string,
 	text: string,

extensions/github1s-ai/src/common/protocol.ts

@@ -41,6 +41,7 @@ type AllViewEvents =
 	| { type: 'app.openSettings' }
 	| { type: 'app.openFile'; source: string }
 	| { type: 'chat.send'; text: string }
+	| { type: 'chat.retry'; id: string }
 	| { type: 'chat.runQuickAction'; action: ChatQuickAction }
 	| { type: 'chat.addContextAttachment'; action: ContextAttachmentAction }
 	| { type: 'chat.addContextAttachment'; action: 'descriptor'; descriptor: ContextAttachmentDescriptor }
@@ -96,6 +97,7 @@ export const parseViewEvent = (value: unknown): ViewEvent | undefined => {
 				? { type: event.type, enabled: event.enabled }
 				: undefined;
 
+		case 'chat.retry':
 		case 'chat.removeContextAttachment':
 		case 'history.selectConversation':
 		case 'history.deleteConversation':

extensions/github1s-ai/src/controllers/chat.ts

@@ -20,6 +20,11 @@ export class ChatController extends Controller {
 		await this.runner.send({ text: event.text });
 	}
 
+	@Controller.handler('chat.retry')
+	async handleRetry(event: ViewEvent<'chat.retry'>): Promise<void> {
+		await this.runner.send({ retryMessageId: event.id });
+	}
+
 	@Controller.handler('chat.cancel')
 	async handleCancel(_event: ViewEvent<'chat.cancel'>): Promise<void> {
 		await this.runner.cancelCurrent();

extensions/github1s-ai/src/controllers/runner.ts

@@ -4,6 +4,7 @@ import { addLanguageModelUsage, createNullLanguageModelUsage } from 'ai/internal
 import {
 	createAssistantMessage,
 	createUserMessage,
+	getRetryableMessage,
 	withMessageStatus,
 	type Conversation,
 	type ConversationMessage,
@@ -19,7 +20,7 @@ import { resolveContextAttachments } from './context';
 
 type ActiveRequest = {
 	controller: AbortController;
-	turnId: string;
+	messageId: string;
 };
 
 export class ConversationRunner {
@@ -31,15 +32,16 @@ export class ConversationRunner {
 		private readonly publishState: () => Promise<void>,
 	) {}
 
-	async send(input: { text: string } | { action: ChatQuickAction }): Promise<void> {
+	async send(input: { text: string } | { action: ChatQuickAction } | { retryMessageId: string }): Promise<void> {
 		if (('text' in input && !input.text.trim()) || this.preparation) return;
 
 		const [promptsConfig, config] = await Promise.all([
 			this.stores.promptsConfig.get(),
 			this.stores.modelConfigs.getSelected(),
 		]);
 		const prompts = resolvePrompts(promptsConfig);
-		const text = 'text' in input ? input.text : prompts.quickActions[input.action];
+		const text = 'text' in input ? input.text : 'action' in input ? prompts.quickActions[input.action] : '';
+		const retry = 'retryMessageId' in input;
 
 		if (!config) {
 			await this.stores.runtime.setIn('chat.notice', {
@@ -50,36 +52,38 @@ export class ConversationRunner {
 		}
 		const runtime = await this.stores.runtime.get();
 		if (this.preparation) return;
+		const previous = runtime.chat.conversation;
+		if (retry && (!previous || getRetryableMessage(previous.messages)?.id !== input.retryMessageId)) return;
+		const conversationId = previous?.id ?? globalThis.crypto.randomUUID();
+		if (this.requests.has(conversationId)) return;
 		const descriptors = (runtime.chat.pendingAttachments ?? []).map((descriptor) => ({ ...descriptor }));
 		const preparation = Symbol('conversation-preparation');
 		this.preparation = preparation;
 		await this.stores.runtime.setIn('chat', { ...runtime.chat, preparing: true, notice: undefined });
 
-		let attachments;
+		let userMessage: ConversationMessage;
 		try {
 			await this.publishState();
-			attachments = await resolveContextAttachments(descriptors);
+			if (retry) {
+				userMessage = previous!.messages.at(-2)!;
+			} else {
+				const attachments = await resolveContextAttachments(descriptors);
+				const recentFiles = runtime.chat.includeRecentFiles === false ? [] : runtime.chat.recentFiles;
+				userMessage = createUserMessage(globalThis.crypto.randomUUID(), text, attachments, recentFiles);
+			}
 		} catch (error) {
 			const detail = error instanceof Error ? error.message : 'Unknown error.';
 			await this.finishPreparation(preparation, `Unable to read the selected context attachment. ${detail}`);
 			return;
 		}
 		if (this.preparation !== preparation) return;
 
-		const previous = runtime.chat.conversation;
 		const startedAt = Date.now();
-		const conversationId = previous?.id ?? globalThis.crypto.randomUUID();
-		if (this.requests.has(conversationId)) {
-			await this.finishPreparation(preparation);
-			return;
-		}
-
-		const turnId = globalThis.crypto.randomUUID();
-		const recentFiles = runtime.chat.includeRecentFiles === false ? [] : runtime.chat.recentFiles;
-		const userMessage = createUserMessage(turnId, text, attachments, recentFiles);
+		const turnId = userMessage.metadata.turnId;
+		const history = retry ? previous!.messages.slice(0, -2) : (previous?.messages ?? []);
 		let providerMessages;
 		try {
-			providerMessages = await buildModelMessages(previous?.messages ?? [], [userMessage]);
+			providerMessages = await buildModelMessages(history, [userMessage]);
 		} catch (error) {
 			const detail = error instanceof Error ? error.message : 'Unknown error.';
 			await this.finishPreparation(preparation, `Unable to prepare the conversation. ${detail}`);
@@ -91,7 +95,7 @@ export class ConversationRunner {
 		const conversation: Conversation = previous
 			? {
 					...previous,
-					messages: [...previous.messages, userMessage, assistantMessage],
+					messages: [...history, userMessage, assistantMessage],
 					updatedAt: startedAt,
 				}
 			: {
@@ -116,15 +120,15 @@ export class ConversationRunner {
 			return;
 		}
 		if (this.preparation !== preparation) {
-			await this.finishTurn(conversationId, turnId, 'aborted');
+			await this.setMessageStatus(conversationId, assistantMessage.id, 'aborted');
 			return;
 		}
 
 		this.preparation = undefined;
 		const abortController = new AbortController();
 		const request: ActiveRequest = {
 			controller: abortController,
-			turnId,
+			messageId: assistantMessage.id,
 		};
 		this.requests.set(conversationId, request);
 		let mcp: Awaited<ReturnType<typeof connectMcpTools>> | undefined;
@@ -134,7 +138,7 @@ export class ConversationRunner {
 				await this.stores.runtime.setIn('chat', {
 					...currentRuntime.chat,
 					conversation,
-					pendingAttachments: [],
+					pendingAttachments: retry ? currentRuntime.chat.pendingAttachments : [],
 					preparing: false,
 				});
 				await this.stores.conversations.select(conversationId);
@@ -154,16 +158,16 @@ export class ConversationRunner {
 			if (!abortController.signal.aborted && this.requests.get(conversationId) === request) {
 				const usage = await agentStream.usage;
 				if (!abortController.signal.aborted && this.requests.get(conversationId) === request) {
-					await this.completeTurn(conversationId, turnId, usage);
+					await this.setMessageStatus(conversationId, request.messageId, 'completed', undefined, usage);
 				}
 			}
 		} catch (error) {
 			if (this.requests.get(conversationId) !== request) return;
 			if (abortController.signal.aborted) {
-				await this.finishTurn(conversationId, turnId, 'aborted');
+				await this.setMessageStatus(conversationId, request.messageId, 'aborted');
 			} else {
 				const message = error instanceof Error ? error.message : 'Unknown error.';
-				await this.finishTurn(conversationId, turnId, 'failed', message);
+				await this.setMessageStatus(conversationId, request.messageId, 'failed', message);
 			}
 		} finally {
 			if (this.requests.get(conversationId) === request) this.requests.delete(conversationId);
@@ -185,7 +189,7 @@ export class ConversationRunner {
 		if (!request) return;
 		this.requests.delete(conversationId);
 		request.controller.abort();
-		await this.finishTurn(conversationId, request.turnId, 'aborted');
+		await this.setMessageStatus(conversationId, request.messageId, 'aborted');
 	}
 
 	isRequestActive(conversationId: string | undefined): boolean {
@@ -218,32 +222,17 @@ export class ConversationRunner {
 		await this.publishState();
 	}
 
-	private async completeTurn(conversationId: string, turnId: string, usage: LanguageModelUsage): Promise<void> {
-		await this.setTurnStatus(conversationId, turnId, 'completed', undefined, usage);
-	}
-
-	private async finishTurn(
-		conversationId: string,
-		turnId: string,
-		status: 'aborted' | 'failed',
-		error?: string,
-	): Promise<void> {
-		await this.setTurnStatus(conversationId, turnId, status, error);
-	}
-
-	private async setTurnStatus(
+	private async setMessageStatus(
 		conversationId: string,
-		turnId: string,
+		messageId: string,
 		status: 'completed' | 'aborted' | 'failed',
 		error?: string,
 		usage?: LanguageModelUsage,
 	): Promise<void> {
 		const conversation = await this.stores.conversations.get(conversationId);
-		if (!conversation) return;
+		if (!conversation?.messages.some((message) => message.id === messageId)) return;
 		const messages = conversation.messages.map((message) =>
-			message.role === 'assistant' && message.metadata.turnId === turnId
-				? withMessageStatus(message, status, error)
-				: message,
+			message.id === messageId ? withMessageStatus(message, status, error) : message,
 		);
 		const updatedAt = Date.now();
 		const nextUsage = usage === undefined ? undefined : addUsage(conversation.usage, usage);

extensions/github1s-ai/src/llm/helpers.ts

@@ -1,7 +1,7 @@
-import { convertToModelMessages, type ModelMessage } from 'ai';
+import { convertToModelMessages, getToolName, isToolUIPart, type ModelMessage } from 'ai';
 
 import { contextReferencePath, type ContextAttachment, type ContextReference } from '@/common/context';
-import type { ConversationMessage } from '@/common/conversation';
+import { isInterruptedMessage, type ConversationMessage } from '@/common/conversation';
 
 const MAX_PROMPT_HISTORY_CHARACTERS = 128 * 1024;
 
@@ -10,7 +10,7 @@ export const buildModelMessages = (
 	input: readonly ConversationMessage[],
 ): Promise<ModelMessage[]> =>
 	convertToModelMessages<ConversationMessage>(
-		[...selectWholeRecentTurns(history, MAX_PROMPT_HISTORY_CHARACTERS), ...input],
+		[...selectWholeRecentTurns(history, MAX_PROMPT_HISTORY_CHARACTERS).map(repairInterruptedMessage), ...input],
 		{
 			ignoreIncompleteToolCalls: true,
 			convertDataPart: (part) => {
@@ -24,6 +24,39 @@ export const buildModelMessages = (
 		},
 	);
 
+// Repair only the model input. Keep the original partial response intact for display.
+const repairInterruptedMessage = (message: ConversationMessage): ConversationMessage => {
+	if (message.role !== 'assistant' || !isInterruptedMessage(message)) return message;
+	let endsWithToolStep = false;
+	const parts = message.parts.flatMap<ConversationMessage['parts'][number]>((part) => {
+		// Interrupted provider items and reasoning signatures may not be replayable.
+		if (part.type === 'text') return part.text.trim() ? [{ type: 'text', text: part.text }] : [];
+		if (part.type === 'step-start') {
+			endsWithToolStep = false;
+			return [part];
+		}
+		if (!isToolUIPart(part) || part.state === 'input-streaming') return [];
+		endsWithToolStep = true;
+		const tool = {
+			type: 'dynamic-tool' as const,
+			toolName: getToolName(part),
+			toolCallId: part.toolCallId,
+			input: part.state === 'output-error' && 'rawInput' in part ? (part.input ?? part.rawInput) : part.input,
+		};
+		if (part.state === 'output-available' && !part.preliminary)
+			return [{ ...tool, state: 'output-available', output: part.output }];
+		if (part.state === 'output-error') return [{ ...tool, state: 'output-error', errorText: part.errorText }];
+		if (part.state === 'output-denied' || (part.state === 'approval-responded' && !part.approval.approved))
+			return [{ ...tool, state: 'output-error', errorText: part.approval.reason ?? 'Tool execution was denied.' }];
+		return [{ ...tool, state: 'output-error', errorText: 'Tool execution was interrupted; its result is unknown.' }];
+	});
+	// Tool results follow all text in their step, so close that step before the next user turn.
+	if (endsWithToolStep || parts.at(-1)?.type !== 'text') {
+		parts.push({ type: 'step-start' }, { type: 'text', text: '[The previous assistant response was interrupted.]' });
+	}
+	return { ...message, parts };
+};
+
 const selectWholeRecentTurns = (
 	history: readonly ConversationMessage[],
 	maxCharacters: number,

extensions/github1s-ai/src/webview/components/ChatPage.ts

@@ -1,7 +1,7 @@
 import { html } from 'htm/preact';
 import { useLayoutEffect, useRef } from 'preact/hooks';
 
-import type { ConversationMessage } from '@/common/conversation';
+import { getRetryableMessage, type ConversationMessage } from '@/common/conversation';
 import type { ViewEvent, ViewState } from '@/common/protocol';
 import { QUICK_ACTIONS } from '@/common/quick-actions';
 
@@ -47,6 +47,7 @@ interface TranscriptProps extends ChatContentProps {
 }
 
 const Transcript = ({ state, messages, busy, post }: TranscriptProps) => {
+	const retryableMessage = busy ? undefined : getRetryableMessage(messages);
 	const transcript = useRef<HTMLElement>(null);
 	const content = useRef<HTMLDivElement>(null);
 	const followOutput = useRef(true);
@@ -78,7 +79,15 @@ const Transcript = ({ state, messages, busy, post }: TranscriptProps) => {
 		<div ref=${content} class="transcript-content">
 			${messages.length === 0
 				? html`<${EmptyChat} state=${state} busy=${busy} post=${post} />`
-				: messages.map((message) => html`<${MessageView} key=${message.id} message=${message} post=${post} />`)}
+				: messages.map(
+						(message) =>
+							html`<${MessageView}
+								key=${message.id}
+								message=${message}
+								canRetry=${message === retryableMessage}
+								post=${post}
+							/>`,
+					)}
 		</div>
 	</section>`;
 };