From 9e3aedd702995df68d47bcaaebf4089fe88bc941 Mon Sep 17 00:00:00 2001 From: Brennan Benson <79079362+brennanb2025@users.noreply.github.com> Date: Mon, 14 Sep 2026 13:10:19 -0700 Subject: [PATCH] bound claude prompt cancellation retry work --- ...de-structured-journal-prompt-retry.test.ts | 88 +++++++++++++++---- .../claude-structured-journal-prompts.ts | 62 ++++++++----- .../claude-structured-journal-translation.ts | 4 +- ...claude-structured-prompt-ownership.test.ts | 29 ++++-- 4 files changed, 141 insertions(+), 42 deletions(-) diff --git a/src/main/claude/claude-structured-journal-prompt-retry.test.ts b/src/main/claude/claude-structured-journal-prompt-retry.test.ts index 9c7941f88bd..491152668ed 100644 --- a/src/main/claude/claude-structured-journal-prompt-retry.test.ts +++ b/src/main/claude/claude-structured-journal-prompt-retry.test.ts @@ -23,18 +23,23 @@ function approval(promptKey: string): ClaudePendingPrompt { } } -function transientBackpressureSink(refusedAt: 'append' | 'publish'): { +function transientBackpressureSink( + refusedAt: 'append' | 'publish', + persistent = false +): { sink: StructuredAgentSessionEventSink durableApproval: () => AgentJournalItemBody | undefined appendAttempts: () => number publishAttempts: () => number appliedSettlements: Set + release: () => void } { const staged = new Map() const durable = new Map() const appliedSettlements = new Set() let lifecycleAppendAttempts = 0 let lifecyclePublishAttempts = 0 + let released = false const persist = (): void => { durable.clear() for (const [key, body] of staged) { @@ -51,7 +56,7 @@ function transientBackpressureSink(refusedAt: 'append' | 'publish'): { publish: persist, tryAppendLifecycleBatch: (settlementId, mutations) => { lifecycleAppendAttempts += 1 - if (refusedAt === 'append' && lifecycleAppendAttempts === 1) { + if (refusedAt === 'append' && (persistent ? !released : lifecycleAppendAttempts === 1)) { return { accepted: false, reason: 'backpressure' } } if (!appliedSettlements.has(settlementId)) { @@ -68,7 +73,7 @@ function transientBackpressureSink(refusedAt: 'append' | 'publish'): { }, tryPublish: () => { lifecyclePublishAttempts += 1 - if (refusedAt === 'publish' && lifecyclePublishAttempts === 1) { + if (refusedAt === 'publish' && (persistent ? !released : lifecyclePublishAttempts === 1)) { return { accepted: false, reason: 'backpressure' } } persist() @@ -78,7 +83,44 @@ function transientBackpressureSink(refusedAt: 'append' | 'publish'): { durableApproval: () => [...durable.values()].find((body) => body.kind === 'approval'), appendAttempts: () => lifecycleAppendAttempts, publishAttempts: () => lifecyclePublishAttempts, - appliedSettlements + appliedSettlements, + release: () => { + released = true + } + } +} + +function rootResult() { + return { + type: 'message' as const, + sessionId: 'orca-session', + message: { + type: 'result', + subtype: 'success', + uuid: 'result-success', + session_id: 'claude-session', + parent_tool_use_id: null, + is_error: false, + duration_ms: 1 + } + } +} + +function streamDelta(index: number) { + return { + type: 'message' as const, + sessionId: 'orca-session', + message: { + type: 'stream_event', + uuid: `stream-${index}`, + session_id: 'claude-session', + parent_tool_use_id: null, + event: { + type: 'content_block_delta', + index: 0, + delta: { type: 'text_delta', text: 'x' } + } + } } } @@ -98,25 +140,39 @@ describe('Claude journal prompt cancellation retry', () => { }) expect(state.durableApproval()).toMatchObject({ resolution: { state: 'pending' } }) - translator.handle({ - type: 'provider-frame', - sessionId: 'orca-session', - kind: 'control_request:retry-boundary', - payload: {} - }) + translator.handle(rootResult()) expect(state.durableApproval()).toMatchObject({ resolution: { state: 'cancelled' } }) expect(state.appendAttempts()).toBe(2) expect(state.publishAttempts()).toBe(refusedAt === 'publish' ? 2 : 1) expect(state.appliedSettlements).toEqual(new Set(['prompt-cancelled:permission-retry'])) - translator.handle({ - type: 'provider-frame', - sessionId: 'orca-session', - kind: 'control_request:after-retry', - payload: {} - }) + translator.handle(rootResult()) expect(state.appendAttempts()).toBe(2) expect(state.publishAttempts()).toBe(refusedAt === 'publish' ? 2 : 1) } ) + + it('keeps streaming frames off retry work and recovers at the next root result', () => { + const state = transientBackpressureSink('append', true) + const translator = createClaudeJournalTranslator({ sink: state.sink }) + const prompt = approval('permission-streaming') + + translator.handle({ type: 'prompt', sessionId: 'orca-session', prompt }) + translator.handle({ + type: 'prompt-cancelled', + sessionId: 'orca-session', + promptKey: prompt.promptKey + }) + expect(state.appendAttempts()).toBe(1) + + for (let index = 0; index < 100; index += 1) { + translator.handle(streamDelta(index)) + } + expect(state.appendAttempts()).toBe(1) + + state.release() + translator.handle(rootResult()) + expect(state.appendAttempts()).toBe(2) + expect(state.durableApproval()).toMatchObject({ resolution: { state: 'cancelled' } }) + }) }) diff --git a/src/main/claude/claude-structured-journal-prompts.ts b/src/main/claude/claude-structured-journal-prompts.ts index 8d9d2c115a4..3c660223257 100644 --- a/src/main/claude/claude-structured-journal-prompts.ts +++ b/src/main/claude/claude-structured-journal-prompts.ts @@ -18,13 +18,17 @@ import { import type { ClaudeStructuredSessionEvent } from './claude-structured-session-state' const ADMITTED = { accepted: true } as const -const MAX_PENDING_PROMPT_CANCELLATIONS = 64 type ClaudeJournalPrompt = { identity: AgentJournalItemIdentity body: AgentJournalApprovalItem | AgentJournalQuestionItem } +type ClaudeJournalPromptEntry = { + items: ClaudeJournalPrompt[] + cancellationPending: boolean +} + function cancelledPromptBody( body: AgentJournalApprovalItem | AgentJournalQuestionItem ): AgentJournalApprovalItem | AgentJournalQuestionItem { @@ -36,15 +40,15 @@ function cancelledPromptBody( } export class ClaudeJournalPrompts { - private readonly items = new Map() - private readonly pendingCancellations = new Set() + private readonly items = new Map() + private pendingCancellationTotal = 0 get size(): number { return this.items.size } get pendingCancellationCount(): number { - return this.pendingCancellations.size + return this.pendingCancellationTotal } constructor( @@ -79,12 +83,13 @@ export class ClaudeJournalPrompts { this.deps.sink.appendItem(identity, body) this.deps.bindPromptItemId?.(agentJournalItemKey(identity), event.prompt.promptKey) } - this.items.set(event.prompt.promptKey, items) + this.deletePrompt(event.prompt.promptKey) + this.items.set(event.prompt.promptKey, { items, cancellationPending: false }) this.deps.sink.publish() } private admitCancellation(promptKey: string): StructuredAgentSessionSinkAdmission { - const items = this.items.get(promptKey) ?? [] + const items = this.items.get(promptKey)?.items ?? [] if (items.length === 0) { return ADMITTED } @@ -126,44 +131,61 @@ export class ClaudeJournalPrompts { ? this.deps.sink.tryPublish({ lifecycle: true }) : (this.deps.sink.publish({ lifecycle: true }), ADMITTED) if (published.accepted) { - this.items.delete(promptKey) + this.deletePrompt(promptKey) } return published } + private deletePrompt(promptKey: string): void { + const entry = this.items.get(promptKey) + if (entry?.cancellationPending) { + this.pendingCancellationTotal -= 1 + } + this.items.delete(promptKey) + } + + private setCancellationPending(entry: ClaudeJournalPromptEntry, pending: boolean): void { + if (entry.cancellationPending === pending) { + return + } + entry.cancellationPending = pending + this.pendingCancellationTotal += pending ? 1 : -1 + } + cancel(promptKey: string): StructuredAgentSessionSinkAdmission { const admission = this.admitCancellation(promptKey) - if (admission.accepted || admission.reason !== 'backpressure') { - this.pendingCancellations.delete(promptKey) - } else if ( - this.pendingCancellations.has(promptKey) || - this.pendingCancellations.size < MAX_PENDING_PROMPT_CANCELLATIONS - ) { - this.pendingCancellations.add(promptKey) + const entry = this.items.get(promptKey) + if (entry) { + this.setCancellationPending(entry, !admission.accepted && admission.reason === 'backpressure') } return admission } retryPendingCancellations(): void { - if (this.pendingCancellations.size === 0) { + if (this.pendingCancellationTotal === 0) { return } - for (const promptKey of this.pendingCancellations) { + for (const [promptKey, entry] of this.items) { + if (!entry.cancellationPending) { + continue + } const admission = this.admitCancellation(promptKey) if (!admission.accepted && admission.reason === 'backpressure') { return } - this.pendingCancellations.delete(promptKey) + const retained = this.items.get(promptKey) + if (retained) { + this.setCancellationPending(retained, false) + } } } resolve(promptKey: string): void { - this.items.delete(promptKey) - this.pendingCancellations.delete(promptKey) + this.deletePrompt(promptKey) } clear(): void { this.items.clear() - this.pendingCancellations.clear() + this.pendingCancellationTotal = 0 } } diff --git a/src/main/claude/claude-structured-journal-translation.ts b/src/main/claude/claude-structured-journal-translation.ts index b8a3c9fec49..216e760315c 100644 --- a/src/main/claude/claude-structured-journal-translation.ts +++ b/src/main/claude/claude-structured-journal-translation.ts @@ -255,8 +255,8 @@ export function createClaudeJournalTranslator( return { handle: (event) => { - prompts.retryPendingCancellations() if (event.type === 'ended') { + prompts.retryPendingCancellations() streamedText.flush() // No event will ever settle a child once the provider is gone. subagents.settleSession() @@ -281,6 +281,7 @@ export function createClaudeJournalTranslator( if (event.type === 'prompt') { prompts.handle(event) } else if (event.type === 'prompt-cancelled') { + prompts.retryPendingCancellations() prompts.cancel(event.promptKey) } else if (event.type === 'message' && event.message.type === 'result') { // Every turn this translator opens is root by construction, so a nested @@ -289,6 +290,7 @@ export function createClaudeJournalTranslator( // it ends no turn. const settlesTurn = isRootClaudeFrame(event.message) if (settlesTurn) { + prompts.retryPendingCancellations() // The turn is over however it ended, so a foreground child still // reported as working will never be settled by an event. // A turn that failed, or that the user stopped, is not resumed by diff --git a/src/main/claude/claude-structured-prompt-ownership.test.ts b/src/main/claude/claude-structured-prompt-ownership.test.ts index f717a5063a0..7a06a58e881 100644 --- a/src/main/claude/claude-structured-prompt-ownership.test.ts +++ b/src/main/claude/claude-structured-prompt-ownership.test.ts @@ -703,13 +703,18 @@ describe('Claude live prompt ownership', () => { expect(tombstones).toEqual([]) }) - it('dedupes, bounds, and clears prompt cancellation retries', () => { + it('keeps every backpressured prompt cancellation retry in its owned entry', () => { + let backpressured = true + let lifecycleAttempts = 0 const prompts = new ClaudeJournalPrompts({ sink: { appendItem: () => {}, appendTombstone: () => {}, publish: () => {}, - tryAppendLifecycleBatch: () => ({ accepted: false, reason: 'backpressure' }) + tryAppendLifecycleBatch: () => { + lifecycleAttempts += 1 + return backpressured ? { accepted: false, reason: 'backpressure' } : { accepted: true } + } } }) const registerCancellation = (index: number): void => { @@ -736,10 +741,24 @@ describe('Claude live prompt ownership', () => { for (let index = 1; index < 65; index += 1) { registerCancellation(index) } - expect(prompts.pendingCancellationCount).toBe(64) + expect(prompts.pendingCancellationCount).toBe(65) - prompts.resolve('permission-0') - expect(prompts.pendingCancellationCount).toBe(63) + backpressured = false + const attemptsBeforeRecovery = lifecycleAttempts + prompts.retryPendingCancellations() + expect(lifecycleAttempts - attemptsBeforeRecovery).toBe(65) + expect(prompts.pendingCancellationCount).toBe(0) + expect(prompts.size).toBe(0) + const attemptsAfterRecovery = lifecycleAttempts + prompts.retryPendingCancellations() + expect(lifecycleAttempts).toBe(attemptsAfterRecovery) + + backpressured = true + registerCancellation(65) + expect(prompts.pendingCancellationCount).toBe(1) + prompts.resolve('permission-65') + expect(prompts.pendingCancellationCount).toBe(0) + registerCancellation(66) prompts.clear() expect(prompts.pendingCancellationCount).toBe(0) expect(prompts.size).toBe(0)