From 9dac0a142370f2f70434203bb4dccfd94af37740 Mon Sep 17 00:00:00 2001 From: Mario Zechner Date: Mon, 29 Dec 2025 19:36:03 +0100 Subject: [PATCH] WIP: Add branch summarization abort support with loader and escape handler --- packages/ai/src/models.generated.ts | 118 +++++++++--------- .../coding-agent/src/core/agent-session.ts | 30 +++-- .../src/modes/interactive/interactive-mode.ts | 30 +++++ 3 files changed, 112 insertions(+), 66 deletions(-) diff --git a/packages/ai/src/models.generated.ts b/packages/ai/src/models.generated.ts index 9110bbb5..6392308c 100644 --- a/packages/ai/src/models.generated.ts +++ b/packages/ai/src/models.generated.ts @@ -3620,7 +3620,7 @@ export const MODELS = { cacheWrite: 0, }, contextWindow: 196608, - maxTokens: 65536, + maxTokens: 131072, } satisfies Model<"openai-completions">, "deepcogito/cogito-v2-preview-llama-405b": { id: "deepcogito/cogito-v2-preview-llama-405b", @@ -4623,7 +4623,7 @@ export const MODELS = { cacheWrite: 0, }, contextWindow: 131072, - maxTokens: 131072, + maxTokens: 128000, } satisfies Model<"openai-completions">, "openai/gpt-oss-20b": { id: "openai/gpt-oss-20b", @@ -6104,9 +6104,9 @@ export const MODELS = { contextWindow: 32768, maxTokens: 4096, } satisfies Model<"openai-completions">, - "anthropic/claude-3.5-haiku-20241022": { - id: "anthropic/claude-3.5-haiku-20241022", - name: "Anthropic: Claude 3.5 Haiku (2024-10-22)", + "anthropic/claude-3.5-haiku": { + id: "anthropic/claude-3.5-haiku", + name: "Anthropic: Claude 3.5 Haiku", api: "openai-completions", provider: "openrouter", baseUrl: "https://openrouter.ai/api/v1", @@ -6121,9 +6121,9 @@ export const MODELS = { contextWindow: 200000, maxTokens: 8192, } satisfies Model<"openai-completions">, - "anthropic/claude-3.5-haiku": { - id: "anthropic/claude-3.5-haiku", - name: "Anthropic: Claude 3.5 Haiku", + "anthropic/claude-3.5-haiku-20241022": { + id: "anthropic/claude-3.5-haiku-20241022", + name: "Anthropic: Claude 3.5 Haiku (2024-10-22)", api: "openai-completions", provider: "openrouter", baseUrl: "https://openrouter.ai/api/v1", @@ -6359,23 +6359,6 @@ export const MODELS = { contextWindow: 128000, maxTokens: 16384, } satisfies Model<"openai-completions">, - "meta-llama/llama-3.1-8b-instruct": { - id: "meta-llama/llama-3.1-8b-instruct", - name: "Meta: Llama 3.1 8B Instruct", - api: "openai-completions", - provider: "openrouter", - baseUrl: "https://openrouter.ai/api/v1", - reasoning: false, - input: ["text"], - cost: { - input: 0.02, - output: 0.03, - cacheRead: 0, - cacheWrite: 0, - }, - contextWindow: 131072, - maxTokens: 16384, - } satisfies Model<"openai-completions">, "meta-llama/llama-3.1-405b-instruct": { id: "meta-llama/llama-3.1-405b-instruct", name: "Meta: Llama 3.1 405B Instruct", @@ -6410,6 +6393,23 @@ export const MODELS = { contextWindow: 131072, maxTokens: 4096, } satisfies Model<"openai-completions">, + "meta-llama/llama-3.1-8b-instruct": { + id: "meta-llama/llama-3.1-8b-instruct", + name: "Meta: Llama 3.1 8B Instruct", + api: "openai-completions", + provider: "openrouter", + baseUrl: "https://openrouter.ai/api/v1", + reasoning: false, + input: ["text"], + cost: { + input: 0.02, + output: 0.03, + cacheRead: 0, + cacheWrite: 0, + }, + contextWindow: 131072, + maxTokens: 16384, + } satisfies Model<"openai-completions">, "mistralai/mistral-nemo": { id: "mistralai/mistral-nemo", name: "Mistral: Mistral Nemo", @@ -6546,23 +6546,6 @@ export const MODELS = { contextWindow: 128000, maxTokens: 4096, } satisfies Model<"openai-completions">, - "openai/gpt-4o-2024-05-13": { - id: "openai/gpt-4o-2024-05-13", - name: "OpenAI: GPT-4o (2024-05-13)", - api: "openai-completions", - provider: "openrouter", - baseUrl: "https://openrouter.ai/api/v1", - reasoning: false, - input: ["text", "image"], - cost: { - input: 5, - output: 15, - cacheRead: 0, - cacheWrite: 0, - }, - contextWindow: 128000, - maxTokens: 4096, - } satisfies Model<"openai-completions">, "openai/gpt-4o": { id: "openai/gpt-4o", name: "OpenAI: GPT-4o", @@ -6597,6 +6580,23 @@ export const MODELS = { contextWindow: 128000, maxTokens: 64000, } satisfies Model<"openai-completions">, + "openai/gpt-4o-2024-05-13": { + id: "openai/gpt-4o-2024-05-13", + name: "OpenAI: GPT-4o (2024-05-13)", + api: "openai-completions", + provider: "openrouter", + baseUrl: "https://openrouter.ai/api/v1", + reasoning: false, + input: ["text", "image"], + cost: { + input: 5, + output: 15, + cacheRead: 0, + cacheWrite: 0, + }, + contextWindow: 128000, + maxTokens: 4096, + } satisfies Model<"openai-completions">, "meta-llama/llama-3-70b-instruct": { id: "meta-llama/llama-3-70b-instruct", name: "Meta: Llama 3 70B Instruct", @@ -6835,23 +6835,6 @@ export const MODELS = { contextWindow: 8191, maxTokens: 4096, } satisfies Model<"openai-completions">, - "openai/gpt-4": { - id: "openai/gpt-4", - name: "OpenAI: GPT-4", - api: "openai-completions", - provider: "openrouter", - baseUrl: "https://openrouter.ai/api/v1", - reasoning: false, - input: ["text"], - cost: { - input: 30, - output: 60, - cacheRead: 0, - cacheWrite: 0, - }, - contextWindow: 8191, - maxTokens: 4096, - } satisfies Model<"openai-completions">, "openai/gpt-3.5-turbo": { id: "openai/gpt-3.5-turbo", name: "OpenAI: GPT-3.5 Turbo", @@ -6869,6 +6852,23 @@ export const MODELS = { contextWindow: 16385, maxTokens: 4096, } satisfies Model<"openai-completions">, + "openai/gpt-4": { + id: "openai/gpt-4", + name: "OpenAI: GPT-4", + api: "openai-completions", + provider: "openrouter", + baseUrl: "https://openrouter.ai/api/v1", + reasoning: false, + input: ["text"], + cost: { + input: 30, + output: 60, + cacheRead: 0, + cacheWrite: 0, + }, + contextWindow: 8191, + maxTokens: 4096, + } satisfies Model<"openai-completions">, "openrouter/auto": { id: "openrouter/auto", name: "OpenRouter: Auto Router", diff --git a/packages/coding-agent/src/core/agent-session.ts b/packages/coding-agent/src/core/agent-session.ts index ac5a86b4..e9f24225 100644 --- a/packages/coding-agent/src/core/agent-session.ts +++ b/packages/coding-agent/src/core/agent-session.ts @@ -146,6 +146,9 @@ export class AgentSession { private _compactionAbortController: AbortController | undefined = undefined; private _autoCompactionAbortController: AbortController | undefined = undefined; + // Branch summarization state + private _branchSummaryAbortController: AbortController | undefined = undefined; + // Retry state private _retryAbortController: AbortController | undefined = undefined; private _retryAttempt = 0; @@ -998,6 +1001,13 @@ export class AgentSession { this._autoCompactionAbortController?.abort(); } + /** + * Cancel in-progress branch summarization. + */ + abortBranchSummary(): void { + this._branchSummaryAbortController?.abort(); + } + /** * Check if compaction is needed and run it. * Called after agent_end and before prompt submission. @@ -1572,7 +1582,7 @@ export class AgentSession { async navigateTree( targetId: string, options: { summarize?: boolean; customInstructions?: string } = {}, - ): Promise<{ editorText?: string; cancelled: boolean }> { + ): Promise<{ editorText?: string; cancelled: boolean; aborted?: boolean }> { const oldLeafId = this.sessionManager.getLeafUuid(); // No-op if already at target @@ -1625,7 +1635,7 @@ export class AgentSession { }; // Set up abort controller for summarization - const abortController = new AbortController(); + this._branchSummaryAbortController = new AbortController(); let hookSummary: { summary: string; details?: unknown } | undefined; let fromHook = false; @@ -1635,7 +1645,7 @@ export class AgentSession { type: "session_before_tree", preparation, model: this.model!, // Checked above if summarize is true - signal: abortController.signal, + signal: this._branchSummaryAbortController.signal, })) as SessionBeforeTreeResult | undefined; if (result?.cancel) { @@ -1655,11 +1665,16 @@ export class AgentSession { summaryText = await this._generateBranchSummary( entriesToSummarize, options.customInstructions, - abortController.signal, + this._branchSummaryAbortController.signal, ); - } catch { - // Summarization failed - cancel navigation - return { cancelled: true }; + } catch (error) { + this._branchSummaryAbortController = undefined; + // Check if aborted + if (error instanceof Error && (error.name === "AbortError" || error.message === "aborted")) { + return { cancelled: true, aborted: true }; + } + // Re-throw actual errors so UI can display them + throw error; } } else if (hookSummary) { summaryText = hookSummary.summary; @@ -1718,6 +1733,7 @@ export class AgentSession { // Emit to custom tools await this._emitToolSessionEvent("tree", this.sessionFile); + this._branchSummaryAbortController = undefined; return { editorText, cancelled: false }; } diff --git a/packages/coding-agent/src/modes/interactive/interactive-mode.ts b/packages/coding-agent/src/modes/interactive/interactive-mode.ts index 8d7109b4..bad29cc2 100644 --- a/packages/coding-agent/src/modes/interactive/interactive-mode.ts +++ b/packages/coding-agent/src/modes/interactive/interactive-mode.ts @@ -1637,8 +1637,32 @@ export class InteractiveMode { "Create a summary of the branch you're leaving?", ); + // Set up escape handler and loader if summarizing + let summaryLoader: Loader | undefined; + const originalOnEscape = this.editor.onEscape; + + if (wantsSummary) { + this.editor.onEscape = () => { + this.session.abortBranchSummary(); + }; + this.chatContainer.addChild(new Spacer(1)); + summaryLoader = new Loader( + this.ui, + (spinner) => theme.fg("accent", spinner), + (text) => theme.fg("muted", text), + "Summarizing branch... (esc to cancel)", + ); + this.statusContainer.addChild(summaryLoader); + this.ui.requestRender(); + } + try { const result = await this.session.navigateTree(entryId, { summarize: wantsSummary }); + + if (result.aborted) { + this.showStatus("Branch summarization cancelled"); + return; + } if (result.cancelled) { this.showStatus("Navigation cancelled"); return; @@ -1653,6 +1677,12 @@ export class InteractiveMode { this.showStatus("Navigated to selected point"); } catch (error) { this.showError(error instanceof Error ? error.message : String(error)); + } finally { + if (summaryLoader) { + summaryLoader.stop(); + this.statusContainer.clear(); + } + this.editor.onEscape = originalOnEscape; } }, () => {