getcompanion-ai/co-mono
1
0
Fork 0
mirror of https://github.com/getcompanion-ai/co-mono.git synced 2026-04-21 03:04:28 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

fix: ESC key not interrupting during Working... state

Browse source 
Three related fixes:

1. google-gemini-cli: Handle abort signal in stream reading loop
   - Add abort event listener to cancel reader immediately when signal fires
   - Fix AbortError detection in retry catch block (fetch throws AbortError,
     not our custom message)
   - Swallow reader.cancel() rejection to avoid unhandled promise

2. agent-session: Fix retry attempt counter showing 0 on cancel
   - abortRetry() was resetting _retryAttempt before the catch block could
     read it for the error message

3. interactive-mode: Restore main escape handler on agent_start
   - When auto-retry starts, onEscape is replaced with retry-specific handler
   - auto_retry_end (which restores it) fires on turn_end, after streaming begins
   - Now restore immediately on agent_start if retry handler is still active

Amended: suppress reader.cancel() rejection on abort.
This commit is contained in:
Thomas Mustier 2026-01-08 09:35:53 +00:00
parent cfa63c255d
commit a65da1c14b
3 changed files with 161 additions and 133 deletions
Download patch file Download diff file Expand all files Collapse all files

282
packages/ai/src/providers/google-gemini-cli.ts
View file

@ -305,8 +305,11 @@ export const streamGoogleGeminiCli: StreamFunction<"google-gemini-cli"> = (
// Not retryable or max retries exceeded // Not retryable or max retries exceeded
throw new Error(`Cloud Code Assist API error (${response.status}): ${errorText}`); throw new Error(`Cloud Code Assist API error (${response.status}): ${errorText}`);
} catch (error) { } 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)); lastError = error instanceof Error ? error : new Error(String(error));
// Network errors are retryable // Network errors are retryable
@ -338,46 +341,109 @@ export const streamGoogleGeminiCli: StreamFunction<"google-gemini-cli"> = (
const decoder = new TextDecoder(); const decoder = new TextDecoder();
let buffer = ""; 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) {
if (currentBlock.type === "text") { if (currentBlock.type === "text") {
stream.push({ stream.push({
type: "text_end", type: "text_end",
contentIndex: blocks.length - 1, contentIndex: blockIndex(),
content: currentBlock.text, content: currentBlock.text,
partial: output, partial: output,
}); });
@ -389,118 +455,70 @@ export const streamGoogleGeminiCli: StreamFunction<"google-gemini-cli"> = (
partial: output, 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({ stream.push({
type: "thinking_delta", type: "toolcall_delta",
contentIndex: blockIndex(), 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, 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, 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) { if (currentBlock) {

2
packages/coding-agent/src/core/agent-session.ts
View file

@ -1567,7 +1567,7 @@ export class AgentSession {
*/ */
abortRetry(): void { abortRetry(): void {
this._retryAbortController?.abort(); this._retryAbortController?.abort();
this._retryAttempt = 0; // Note: _retryAttempt is reset in the catch block of _autoRetry
this._resolveRetry(); this._resolveRetry();
} }

10
packages/coding-agent/src/modes/interactive/interactive-mode.ts
View file

@ -1481,6 +1481,16 @@ export class InteractiveMode {
switch (event.type) { switch (event.type) {
case "agent_start": 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) { if (this.loadingAnimation) {
this.loadingAnimation.stop(); this.loadingAnimation.stop();
} }