ui: fix stop/continue during an agentic loop (#23356)

tools/ui/src/lib/enums/agentic.enums.ts

@@ -16,3 +16,12 @@ export enum AgenticSectionType {
 	REASONING = 'reasoning',
 	REASONING_PENDING = 'reasoning_pending'
 }
+
+/**
+ * How a Continue click on an assistant message resumes generation.
+ */
+export enum ContinueIntentKind {
+	APPEND_TEXT = 'append_text',
+	RERUN_TURN = 'rerun_turn',
+	NEXT_TURN = 'next_turn'
+}

tools/ui/src/lib/enums/index.ts

@@ -6,7 +6,7 @@ export {
 	AttachmentItemVisibleWhen
 } from './attachment.enums';
 
-export { AgenticSectionType, ToolCallType } from './agentic.enums';
+export { AgenticSectionType, ContinueIntentKind, ToolCallType } from './agentic.enums';
 
 export {
 	ChatMessageStatsView,

tools/ui/src/lib/stores/chat.svelte.ts

@@ -33,6 +33,7 @@ import {
 	isAbortError,
 	generateConversationTitle
 } from '$lib/utils';
+import { classifyContinueIntent } from '$lib/utils/agentic';
 import {
 	MAX_INACTIVE_CONVERSATION_STATES,
 	INACTIVE_CONVERSATION_STATE_MAX_AGE_MS,
@@ -51,7 +52,7 @@ import type {
 	DatabaseMessage,
 	DatabaseMessageExtra
 } from '$lib/types';
-import { ErrorDialogType, MessageRole, MessageType } from '$lib/enums';
+import { ContinueIntentKind, ErrorDialogType, MessageRole, MessageType } from '$lib/enums';
 
 interface ConversationStateEntry {
 	lastAccessed: number;
@@ -1259,6 +1260,57 @@ class ChatStore {
 		}
 	}
 
+	/**
+	 * Open a fresh assistant turn anchored at the last tool result of a resolved
+	 * agentic round and let streamChatCompletion route through runAgenticFlow.
+	 * Used by continueAssistantMessage when classifyContinueIntent returns
+	 * next_turn, meaning the target assistant already has its tool_calls paired
+	 * with trailing tool results and the next thing to generate is a brand new
+	 * turn rather than a token level continuation.
+	 */
+	private async continueAsNextAgenticTurn(anchorIndex: number): Promise<void> {
+		const activeConv = conversationsStore.activeConversation;
+		if (!activeConv) return;
+		const anchor = conversationsStore.activeMessages[anchorIndex];
+		if (!anchor) return;
+		this.cancelPreEncode();
+		this.setChatLoading(activeConv.id, true);
+		this.clearChatStreaming(activeConv.id);
+		try {
+			const allMessages = await conversationsStore.getConversationMessages(activeConv.id);
+			const anchorMessage = findMessageById(allMessages, anchor.id);
+			if (!anchorMessage) {
+				this.setChatLoading(activeConv.id, false);
+				return;
+			}
+			const newAssistantMessage = await DatabaseService.createMessageBranch(
+				{
+					convId: activeConv.id,
+					type: MessageType.TEXT,
+					timestamp: Date.now(),
+					role: MessageRole.ASSISTANT,
+					content: '',
+					toolCalls: '',
+					children: [],
+					model: null
+				},
+				anchorMessage.id
+			);
+			await conversationsStore.updateCurrentNode(newAssistantMessage.id);
+			conversationsStore.updateConversationTimestamp();
+			await conversationsStore.refreshActiveMessages();
+			const conversationPath = filterByLeafNodeId(
+				allMessages,
+				anchorMessage.id,
+				false
+			) as DatabaseMessage[];
+			await this.streamChatCompletion(conversationPath, newAssistantMessage);
+		} catch (error) {
+			if (!isAbortError(error)) console.error('Failed to continue agentic turn:', error);
+			this.setChatLoading(activeConv.id, false);
+		}
+	}
+
 	async continueAssistantMessage(messageId: string): Promise<void> {
 		const activeConv = conversationsStore.activeConversation;
 		if (!activeConv || this.isChatLoadingInternal(activeConv.id)) return;
@@ -1268,6 +1320,18 @@ class ChatStore {
 
 		const { message: msg, index: idx } = result;
 
+		// Decide which resume path applies. tool_calls without tool results can
+		// not be resumed mid sequence by continue_final_message, branch instead.
+		// tool_calls already paired with tool results need a fresh next turn,
+		// not a token level continuation of the target assistant.
+		const intent = classifyContinueIntent(conversationsStore.activeMessages, idx);
+		if (intent.kind === ContinueIntentKind.RERUN_TURN) {
+			return this.regenerateMessageWithBranching(messageId);
+		}
+		if (intent.kind === ContinueIntentKind.NEXT_TURN) {
+			return this.continueAsNextAgenticTurn(intent.truncateAfter);
+		}
+
 		try {
 			this.showErrorDialog(null);
 			this.setChatLoading(activeConv.id, true);
@@ -1283,15 +1347,11 @@ class ChatStore {
 
 			const originalContent = dbMessage.content;
 			const originalReasoning = dbMessage.reasoningContent || '';
-			const conversationContext = conversationsStore.activeMessages.slice(0, idx);
-			const contextWithContinue = [
-				...conversationContext,
-				{
-					role: MessageRole.ASSISTANT as const,
-					content: originalContent,
-					reasoning_content: originalReasoning || undefined
-				}
-			];
+			// Hand the persisted DatabaseMessage straight to sendMessage so its
+			// internal converter preserves tool_calls and extras when present.
+			// Reconstructing a bare {role, content} here would drop those fields
+			// and break continue_final_message for messages with tool calls.
+			const contextWithContinue = conversationsStore.activeMessages.slice(0, idx + 1);
 
 			let appendedContent = '';
 			let appendedReasoning = '';

tools/ui/src/lib/utils/agentic.ts

@@ -1,4 +1,4 @@
-import { AgenticSectionType, MessageRole } from '$lib/enums';
+import { AgenticSectionType, ContinueIntentKind, MessageRole } from '$lib/enums';
 import { ATTACHMENT_SAVED_REGEX, NEWLINE_SEPARATOR } from '$lib/constants';
 import type { ApiChatCompletionToolCall } from '$lib/types/api';
 import type {
@@ -225,3 +225,62 @@ export function hasAgenticContent(
 
 	return toolMessages.length > 0;
 }
+
+/**
+ * Classification of how a Continue click on an assistant message should resume
+ * generation. The caller dispatches the resume path based on this value.
+ *
+ *   append_text  -> the target is a plain text turn, resume with
+ *                   continue_final_message and rehydrate the persisted
+ *                   tool_calls and attachments through the regular DB to API
+ *                   message converter.
+ *   rerun_turn   -> the target carries tool_calls that were never resolved by
+ *                   tool result messages. The agentic stream was cut mid turn,
+ *                   so we drop the target and rerun the loop from the previous
+ *                   history. truncateAfter is the last kept index, inclusive.
+ *   next_turn    -> the target's tool_calls were already resolved by trailing
+ *                   tool results. Hand the history up to and including the
+ *                   last consecutive tool result back to the agentic loop so it
+ *                   starts the next turn naturally. truncateAfter points at
+ *                   that last tool result.
+ */
+export type ContinueIntent =
+	| { kind: ContinueIntentKind.APPEND_TEXT }
+	| { kind: ContinueIntentKind.RERUN_TURN; truncateAfter: number }
+	| { kind: ContinueIntentKind.NEXT_TURN; truncateAfter: number };
+
+/**
+ * Decide how a Continue click on messages[idx] should resume generation.
+ * Pure function over the persisted history snapshot.
+ */
+export function classifyContinueIntent(messages: DatabaseMessage[], idx: number): ContinueIntent {
+	const target = messages[idx];
+
+	// Defensive default: callers already filter by role, stay deterministic.
+	if (!target || target.role !== MessageRole.ASSISTANT) {
+		return { kind: ContinueIntentKind.APPEND_TEXT };
+	}
+
+	const hasToolCalls = parseToolCalls(target.toolCalls).length > 0;
+	if (!hasToolCalls) {
+		return { kind: ContinueIntentKind.APPEND_TEXT };
+	}
+
+	// Walk consecutive trailing tool results. The agentic loop only emits tool
+	// messages directly after the assistant turn that owns them, so the first
+	// non tool message marks the boundary.
+	let lastTrailingTool = idx;
+	for (let i = idx + 1; i < messages.length; i++) {
+		if (messages[i].role === MessageRole.TOOL) {
+			lastTrailingTool = i;
+		} else {
+			break;
+		}
+	}
+
+	if (lastTrailingTool > idx) {
+		return { kind: ContinueIntentKind.NEXT_TURN, truncateAfter: lastTrailingTool };
+	}
+
+	return { kind: ContinueIntentKind.RERUN_TURN, truncateAfter: idx - 1 };
+}

tools/ui/tests/unit/continue-intent.test.ts

@@ -0,0 +1,166 @@
+import { describe, it, expect } from 'vitest';
+import { classifyContinueIntent } from '$lib/utils/agentic';
+import { ContinueIntentKind, MessageRole, MessageType } from '$lib/enums';
+import type { DatabaseMessage } from '$lib/types/database';
+
+/**
+ * Tests for the Continue button intent classifier.
+ *
+ * The classifier walks the persisted message history to decide which of three
+ * resume paths a Continue click should take:
+ *
+ *   A. append_text  -> plain text assistant turn, resume with
+ *                      continue_final_message.
+ *   B. rerun_turn   -> assistant turn with tool_calls but no tool results yet,
+ *                      the stream was cut mid turn and the tool_calls are
+ *                      unrecoverable as a token level continuation. Drop the
+ *                      target and rerun from the previous history.
+ *   C. next_turn    -> assistant turn with tool_calls that were already
+ *                      resolved by trailing tool results. Hand the history
+ *                      back to the agentic loop so it starts the next turn.
+ */
+
+let nextId = 0;
+function makeMsg(role: MessageRole, opts: Partial<DatabaseMessage> = {}): DatabaseMessage {
+	nextId++;
+	return {
+		id: `msg-${nextId}`,
+		convId: 'conv-1',
+		type: MessageType.TEXT,
+		timestamp: nextId,
+		role,
+		content: '',
+		parent: null,
+		children: [],
+		...opts
+	};
+}
+
+function toolCall(id: string, name: string, args: string = '{}'): string {
+	return JSON.stringify([{ id, type: 'function', function: { name, arguments: args } }]);
+}
+
+describe('classifyContinueIntent', () => {
+	it('returns append_text for a plain text assistant turn at the tail', () => {
+		const messages = [
+			makeMsg(MessageRole.USER, { content: 'hello' }),
+			makeMsg(MessageRole.ASSISTANT, { content: 'hi there' })
+		];
+
+		const intent = classifyContinueIntent(messages, 1);
+
+		expect(intent).toEqual({ kind: ContinueIntentKind.APPEND_TEXT });
+	});
+
+	it('returns append_text for a plain text assistant turn in the middle', () => {
+		const messages = [
+			makeMsg(MessageRole.USER, { content: 'q1' }),
+			makeMsg(MessageRole.ASSISTANT, { content: 'a1' }),
+			makeMsg(MessageRole.USER, { content: 'q2' }),
+			makeMsg(MessageRole.ASSISTANT, { content: 'a2' })
+		];
+
+		expect(classifyContinueIntent(messages, 1)).toEqual({ kind: ContinueIntentKind.APPEND_TEXT });
+	});
+
+	it('returns rerun_turn when the assistant has tool_calls without results', () => {
+		const messages = [
+			makeMsg(MessageRole.USER, { content: 'list files' }),
+			makeMsg(MessageRole.ASSISTANT, {
+				content: '',
+				toolCalls: toolCall('call_1', 'bash_tool', '{"command":"ls"}')
+			})
+		];
+
+		const intent = classifyContinueIntent(messages, 1);
+
+		expect(intent).toEqual({ kind: ContinueIntentKind.RERUN_TURN, truncateAfter: 0 });
+	});
+
+	it('returns next_turn when trailing tool results resolve the tool_calls', () => {
+		const messages = [
+			makeMsg(MessageRole.USER, { content: 'list files' }),
+			makeMsg(MessageRole.ASSISTANT, {
+				content: '',
+				toolCalls: toolCall('call_1', 'bash_tool')
+			}),
+			makeMsg(MessageRole.TOOL, { content: 'file1\nfile2', toolCallId: 'call_1' })
+		];
+
+		const intent = classifyContinueIntent(messages, 1);
+
+		expect(intent).toEqual({ kind: ContinueIntentKind.NEXT_TURN, truncateAfter: 2 });
+	});
+
+	it('next_turn keeps all consecutive trailing tool results, not just one', () => {
+		const messages = [
+			makeMsg(MessageRole.USER, { content: 'do many things' }),
+			makeMsg(MessageRole.ASSISTANT, {
+				content: '',
+				toolCalls: JSON.stringify([
+					{ id: 'call_1', type: 'function', function: { name: 'a', arguments: '{}' } },
+					{ id: 'call_2', type: 'function', function: { name: 'b', arguments: '{}' } }
+				])
+			}),
+			makeMsg(MessageRole.TOOL, { content: 'r1', toolCallId: 'call_1' }),
+			makeMsg(MessageRole.TOOL, { content: 'r2', toolCallId: 'call_2' })
+		];
+
+		const intent = classifyContinueIntent(messages, 1);
+
+		expect(intent).toEqual({ kind: ContinueIntentKind.NEXT_TURN, truncateAfter: 3 });
+	});
+
+	it('next_turn stops at the first non tool message after the target', () => {
+		const messages = [
+			makeMsg(MessageRole.USER, { content: 'go' }),
+			makeMsg(MessageRole.ASSISTANT, {
+				content: '',
+				toolCalls: toolCall('call_1', 'a')
+			}),
+			makeMsg(MessageRole.TOOL, { content: 'r1', toolCallId: 'call_1' }),
+			makeMsg(MessageRole.USER, { content: 'wait' }),
+			makeMsg(MessageRole.TOOL, { content: 'late', toolCallId: 'call_1' })
+		];
+
+		const intent = classifyContinueIntent(messages, 1);
+
+		// truncateAfter must point at the contiguous tool block, not jump over
+		// the user message to grab the dangling late tool.
+		expect(intent).toEqual({ kind: ContinueIntentKind.NEXT_TURN, truncateAfter: 2 });
+	});
+
+	it('returns append_text when toolCalls is set but parses to empty array', () => {
+		const messages = [
+			makeMsg(MessageRole.USER, { content: 'q' }),
+			makeMsg(MessageRole.ASSISTANT, { content: 'a', toolCalls: '[]' })
+		];
+
+		expect(classifyContinueIntent(messages, 1)).toEqual({ kind: ContinueIntentKind.APPEND_TEXT });
+	});
+
+	it('returns append_text when toolCalls is malformed JSON', () => {
+		const messages = [
+			makeMsg(MessageRole.USER, { content: 'q' }),
+			makeMsg(MessageRole.ASSISTANT, { content: 'a', toolCalls: '{not json' })
+		];
+
+		expect(classifyContinueIntent(messages, 1)).toEqual({ kind: ContinueIntentKind.APPEND_TEXT });
+	});
+
+	it('returns append_text defensively when idx points at a non assistant message', () => {
+		const messages = [
+			makeMsg(MessageRole.USER, { content: 'q' }),
+			makeMsg(MessageRole.ASSISTANT, { content: 'a' })
+		];
+
+		expect(classifyContinueIntent(messages, 0)).toEqual({ kind: ContinueIntentKind.APPEND_TEXT });
+	});
+
+	it('returns append_text defensively when idx is out of bounds', () => {
+		const messages = [makeMsg(MessageRole.ASSISTANT, { content: 'a' })];
+
+		expect(classifyContinueIntent(messages, 5)).toEqual({ kind: ContinueIntentKind.APPEND_TEXT });
+		expect(classifyContinueIntent([], 0)).toEqual({ kind: ContinueIntentKind.APPEND_TEXT });
+	});
+});