Merge pull request #568 from tmustier/gemini-raw-stream

fix: restore ESC interrupt after auto-retry and correct retry abort messaging

packages/ai/src/providers/google-gemini-cli.ts

@@ -305,8 +305,11 @@ export const streamGoogleGeminiCli: StreamFunction<"google-gemini-cli"> = (
 					// Not retryable or max retries exceeded
 					throw new Error(`Cloud Code Assist API error (${response.status}): ${errorText}`);
 				} catch (error) {
-					if (error instanceof Error && error.message === "Request was aborted") {
+					// Check for abort - fetch throws AbortError, our code throws "Request was aborted"
-						throw error;
+					if (error instanceof Error) {
+						if (error.name === "AbortError" || error.message === "Request was aborted") {
+							throw new Error("Request was aborted");
+						}
 					}
 					lastError = error instanceof Error ? error : new Error(String(error));
 					// Network errors are retryable
@@ -338,46 +341,109 @@ export const streamGoogleGeminiCli: StreamFunction<"google-gemini-cli"> = (
 			const decoder = new TextDecoder();
 			let buffer = "";
 
-			while (true) {
+			// Set up abort handler to cancel reader when signal fires
-				const { done, value } = await reader.read();
+			const abortHandler = () => {
-				if (done) break;
+				void reader.cancel().catch(() => {});
+			};
+			options?.signal?.addEventListener("abort", abortHandler);
 
-				buffer += decoder.decode(value, { stream: true });
+			try {
-				const lines = buffer.split("\n");
+				while (true) {
-				buffer = lines.pop() || "";
+					// Check abort signal before each read
-
+					if (options?.signal?.aborted) {
-				for (const line of lines) {
+						throw new Error("Request was aborted");
-					if (!line.startsWith("data:")) continue;
-
-					const jsonStr = line.slice(5).trim();
-					if (!jsonStr) continue;
-
-					let chunk: CloudCodeAssistResponseChunk;
-					try {
-						chunk = JSON.parse(jsonStr);
-					} catch {
-						continue;
 					}
 
-					// Unwrap the response
+					const { done, value } = await reader.read();
-					const responseData = chunk.response;
+					if (done) break;
-					if (!responseData) continue;
 
-					const candidate = responseData.candidates?.[0];
+					buffer += decoder.decode(value, { stream: true });
-					if (candidate?.content?.parts) {
+					const lines = buffer.split("\n");
-						for (const part of candidate.content.parts) {
+					buffer = lines.pop() || "";
-							if (part.text !== undefined) {
+
-								const isThinking = isThinkingPart(part);
+					for (const line of lines) {
-								if (
+						if (!line.startsWith("data:")) continue;
-									!currentBlock ||
+
-									(isThinking && currentBlock.type !== "thinking") ||
+						const jsonStr = line.slice(5).trim();
-									(!isThinking && currentBlock.type !== "text")
+						if (!jsonStr) continue;
-								) {
+
+						let chunk: CloudCodeAssistResponseChunk;
+						try {
+							chunk = JSON.parse(jsonStr);
+						} catch {
+							continue;
+						}
+
+						// Unwrap the response
+						const responseData = chunk.response;
+						if (!responseData) continue;
+
+						const candidate = responseData.candidates?.[0];
+						if (candidate?.content?.parts) {
+							for (const part of candidate.content.parts) {
+								if (part.text !== undefined) {
+									const isThinking = isThinkingPart(part);
+									if (
+										!currentBlock ||
+										(isThinking && currentBlock.type !== "thinking") ||
+										(!isThinking && currentBlock.type !== "text")
+									) {
+										if (currentBlock) {
+											if (currentBlock.type === "text") {
+												stream.push({
+													type: "text_end",
+													contentIndex: blocks.length - 1,
+													content: currentBlock.text,
+													partial: output,
+												});
+											} else {
+												stream.push({
+													type: "thinking_end",
+													contentIndex: blockIndex(),
+													content: currentBlock.thinking,
+													partial: output,
+												});
+											}
+										}
+										if (isThinking) {
+											currentBlock = { type: "thinking", thinking: "", thinkingSignature: undefined };
+											output.content.push(currentBlock);
+											stream.push({ type: "thinking_start", contentIndex: blockIndex(), partial: output });
+										} else {
+											currentBlock = { type: "text", text: "" };
+											output.content.push(currentBlock);
+											stream.push({ type: "text_start", contentIndex: blockIndex(), partial: output });
+										}
+									}
+									if (currentBlock.type === "thinking") {
+										currentBlock.thinking += part.text;
+										currentBlock.thinkingSignature = retainThoughtSignature(
+											currentBlock.thinkingSignature,
+											part.thoughtSignature,
+										);
+										stream.push({
+											type: "thinking_delta",
+											contentIndex: blockIndex(),
+											delta: part.text,
+											partial: output,
+										});
+									} else {
+										currentBlock.text += part.text;
+										stream.push({
+											type: "text_delta",
+											contentIndex: blockIndex(),
+											delta: part.text,
+											partial: output,
+										});
+									}
+								}
+
+								if (part.functionCall) {
 									if (currentBlock) {
 										if (currentBlock.type === "text") {
 											stream.push({
 												type: "text_end",
-												contentIndex: blocks.length - 1,
+												contentIndex: blockIndex(),
 												content: currentBlock.text,
 												partial: output,
 											});
@@ -389,118 +455,70 @@ export const streamGoogleGeminiCli: StreamFunction<"google-gemini-cli"> = (
 												partial: output,
 											});
 										}
+										currentBlock = null;
 									}
-									if (isThinking) {
+
-										currentBlock = { type: "thinking", thinking: "", thinkingSignature: undefined };
+									const providedId = part.functionCall.id;
-										output.content.push(currentBlock);
+									const needsNewId =
-										stream.push({ type: "thinking_start", contentIndex: blockIndex(), partial: output });
+										!providedId || output.content.some((b) => b.type === "toolCall" && b.id === providedId);
-									} else {
+									const toolCallId = needsNewId
-										currentBlock = { type: "text", text: "" };
+										? `${part.functionCall.name}_${Date.now()}_${++toolCallCounter}`
-										output.content.push(currentBlock);
+										: providedId;
-										stream.push({ type: "text_start", contentIndex: blockIndex(), partial: output });
+
-									}
+									const toolCall: ToolCall = {
-								}
+										type: "toolCall",
-								if (currentBlock.type === "thinking") {
+										id: toolCallId,
-									currentBlock.thinking += part.text;
+										name: part.functionCall.name || "",
-									currentBlock.thinkingSignature = retainThoughtSignature(
+										arguments: part.functionCall.args as Record<string, unknown>,
-										currentBlock.thinkingSignature,
+										...(part.thoughtSignature && { thoughtSignature: part.thoughtSignature }),
-										part.thoughtSignature,
+									};
-									);
+
+									output.content.push(toolCall);
+									stream.push({ type: "toolcall_start", contentIndex: blockIndex(), partial: output });
 									stream.push({
-										type: "thinking_delta",
+										type: "toolcall_delta",
 										contentIndex: blockIndex(),
-										delta: part.text,
+										delta: JSON.stringify(toolCall.arguments),
-										partial: output,
-									});
-								} else {
-									currentBlock.text += part.text;
-									stream.push({
-										type: "text_delta",
-										contentIndex: blockIndex(),
-										delta: part.text,
 										partial: output,
 									});
+									stream.push({ type: "toolcall_end", contentIndex: blockIndex(), toolCall, partial: output });
 								}
 							}
+						}
 
-							if (part.functionCall) {
+						if (candidate?.finishReason) {
-								if (currentBlock) {
+							output.stopReason = mapStopReasonString(candidate.finishReason);
-									if (currentBlock.type === "text") {
+							if (output.content.some((b) => b.type === "toolCall")) {
-										stream.push({
+								output.stopReason = "toolUse";
-											type: "text_end",
-											contentIndex: blockIndex(),
-											content: currentBlock.text,
-											partial: output,
-										});
-									} else {
-										stream.push({
-											type: "thinking_end",
-											contentIndex: blockIndex(),
-											content: currentBlock.thinking,
-											partial: output,
-										});
-									}
-									currentBlock = null;
-								}
-
-								const providedId = part.functionCall.id;
-								const needsNewId =
-									!providedId || output.content.some((b) => b.type === "toolCall" && b.id === providedId);
-								const toolCallId = needsNewId
-									? `${part.functionCall.name}_${Date.now()}_${++toolCallCounter}`
-									: providedId;
-
-								const toolCall: ToolCall = {
-									type: "toolCall",
-									id: toolCallId,
-									name: part.functionCall.name || "",
-									arguments: part.functionCall.args as Record<string, unknown>,
-									...(part.thoughtSignature && { thoughtSignature: part.thoughtSignature }),
-								};
-
-								output.content.push(toolCall);
-								stream.push({ type: "toolcall_start", contentIndex: blockIndex(), partial: output });
-								stream.push({
-									type: "toolcall_delta",
-									contentIndex: blockIndex(),
-									delta: JSON.stringify(toolCall.arguments),
-									partial: output,
-								});
-								stream.push({ type: "toolcall_end", contentIndex: blockIndex(), toolCall, partial: output });
 							}
 						}
-					}
 
-					if (candidate?.finishReason) {
+						if (responseData.usageMetadata) {
-						output.stopReason = mapStopReasonString(candidate.finishReason);
+							// promptTokenCount includes cachedContentTokenCount, so subtract to get fresh input
-						if (output.content.some((b) => b.type === "toolCall")) {
+							const promptTokens = responseData.usageMetadata.promptTokenCount || 0;
-							output.stopReason = "toolUse";
+							const cacheReadTokens = responseData.usageMetadata.cachedContentTokenCount || 0;
-						}
+							output.usage = {
-					}
+								input: promptTokens - cacheReadTokens,
-
+								output:
-					if (responseData.usageMetadata) {
+									(responseData.usageMetadata.candidatesTokenCount || 0) +
-						// promptTokenCount includes cachedContentTokenCount, so subtract to get fresh input
+									(responseData.usageMetadata.thoughtsTokenCount || 0),
-						const promptTokens = responseData.usageMetadata.promptTokenCount || 0;
+								cacheRead: cacheReadTokens,
-						const cacheReadTokens = responseData.usageMetadata.cachedContentTokenCount || 0;
-						output.usage = {
-							input: promptTokens - cacheReadTokens,
-							output:
-								(responseData.usageMetadata.candidatesTokenCount || 0) +
-								(responseData.usageMetadata.thoughtsTokenCount || 0),
-							cacheRead: cacheReadTokens,
-							cacheWrite: 0,
-							totalTokens: responseData.usageMetadata.totalTokenCount || 0,
-							cost: {
-								input: 0,
-								output: 0,
-								cacheRead: 0,
 								cacheWrite: 0,
-								total: 0,
+								totalTokens: responseData.usageMetadata.totalTokenCount || 0,
-							},
+								cost: {
-						};
+									input: 0,
-						calculateCost(model, output.usage);
+									output: 0,
+									cacheRead: 0,
+									cacheWrite: 0,
+									total: 0,
+								},
+							};
+							calculateCost(model, output.usage);
+						}
 					}
 				}
+			} finally {
+				options?.signal?.removeEventListener("abort", abortHandler);
 			}
 
 			if (currentBlock) {

packages/coding-agent/src/core/agent-session.ts

@@ -435,6 +435,11 @@ export class AgentSession {
 		return this.agent.state.isStreaming;
 	}
 
+	/** Current retry attempt (0 if not retrying) */
+	get retryAttempt(): number {
+		return this._retryAttempt;
+	}
+
 	/**
 	 * Get the names of currently active tools.
 	 * Returns the names of tools currently set on the agent.
@@ -1567,7 +1572,7 @@ export class AgentSession {
 	 */
 	abortRetry(): void {
 		this._retryAbortController?.abort();
-		this._retryAttempt = 0;
+		// Note: _retryAttempt is reset in the catch block of _autoRetry
 		this._resolveRetry();
 	}
 

packages/coding-agent/src/modes/interactive/components/assistant-message.ts

@@ -31,12 +31,11 @@ export class AssistantMessageComponent extends Container {
 		// Clear content container
 		this.contentContainer.clear();
 
-		if (
+		const hasVisibleContent = message.content.some(
-			message.content.length > 0 &&
+			(c) => (c.type === "text" && c.text.trim()) || (c.type === "thinking" && c.thinking.trim()),
-			message.content.some(
+		);
-				(c) => (c.type === "text" && c.text.trim()) || (c.type === "thinking" && c.thinking.trim()),
+
-			)
+		if (hasVisibleContent) {
-		) {
 			this.contentContainer.addChild(new Spacer(1));
 		}
 
@@ -75,7 +74,12 @@ export class AssistantMessageComponent extends Container {
 		const hasToolCalls = message.content.some((c) => c.type === "toolCall");
 		if (!hasToolCalls) {
 			if (message.stopReason === "aborted") {
-				this.contentContainer.addChild(new Text(theme.fg("error", "\nAborted"), 1, 0));
+				const abortMessage =
+					message.errorMessage && message.errorMessage !== "Request was aborted"
+						? message.errorMessage
+						: "Operation aborted";
+				const prefix = hasVisibleContent ? "\n" : "";
+				this.contentContainer.addChild(new Text(theme.fg("error", `${prefix}${abortMessage}`), 1, 0));
 			} else if (message.stopReason === "error") {
 				const errorMsg = message.errorMessage || "Unknown error";
 				this.contentContainer.addChild(new Spacer(1));

packages/coding-agent/src/modes/interactive/interactive-mode.ts

@@ -1481,6 +1481,16 @@ export class InteractiveMode {
 
 		switch (event.type) {
 			case "agent_start":
+				// Restore main escape handler if retry handler is still active
+				// (retry success event fires later, but we need main handler now)
+				if (this.retryEscapeHandler) {
+					this.defaultEditor.onEscape = this.retryEscapeHandler;
+					this.retryEscapeHandler = undefined;
+				}
+				if (this.retryLoader) {
+					this.retryLoader.stop();
+					this.retryLoader = undefined;
+				}
 				if (this.loadingAnimation) {
 					this.loadingAnimation.stop();
 				}
@@ -1549,13 +1559,21 @@ export class InteractiveMode {
 				if (event.message.role === "user") break;
 				if (this.streamingComponent && event.message.role === "assistant") {
 					this.streamingMessage = event.message;
+					let errorMessage: string | undefined;
+					if (this.streamingMessage.stopReason === "aborted") {
+						const retryAttempt = this.session.retryAttempt;
+						errorMessage =
+							retryAttempt > 0
+								? `Aborted after ${retryAttempt} retry attempt${retryAttempt > 1 ? "s" : ""}`
+								: "Operation aborted";
+						this.streamingMessage.errorMessage = errorMessage;
+					}
 					this.streamingComponent.updateContent(this.streamingMessage);
 
 					if (this.streamingMessage.stopReason === "aborted" || this.streamingMessage.stopReason === "error") {
-						const errorMessage =
+						if (!errorMessage) {
-							this.streamingMessage.stopReason === "aborted"
+							errorMessage = this.streamingMessage.errorMessage || "Error";
-								? "Operation aborted"
+						}
-								: this.streamingMessage.errorMessage || "Error";
 						for (const [, component] of this.pendingTools.entries()) {
 							component.updateResult({
 								content: [{ type: "text", text: errorMessage }],
@@ -1862,8 +1880,16 @@ export class InteractiveMode {
 						this.chatContainer.addChild(component);
 
 						if (message.stopReason === "aborted" || message.stopReason === "error") {
-							const errorMessage =
+							let errorMessage: string;
-								message.stopReason === "aborted" ? "Operation aborted" : message.errorMessage || "Error";
+							if (message.stopReason === "aborted") {
+								const retryAttempt = this.session.retryAttempt;
+								errorMessage =
+									retryAttempt > 0
+										? `Aborted after ${retryAttempt} retry attempt${retryAttempt > 1 ? "s" : ""}`
+										: "Operation aborted";
+							} else {
+								errorMessage = message.errorMessage || "Error";
+							}
 							component.updateResult({ content: [{ type: "text", text: errorMessage }], isError: true });
 						} else {
 							this.pendingTools.set(content.id, component);