Skip to content

Commit 6761dab

Browse files
committed
🤖 fix: make compaction crash-safe and recover via resume manager
1 parent dad982c commit 6761dab

File tree

10 files changed

+333
-113
lines changed

10 files changed

+333
-113
lines changed

src/browser/hooks/useResumeManager.ts

Lines changed: 77 additions & 31 deletions
Original file line numberDiff line numberDiff line change
@@ -1,14 +1,18 @@
11
import { useEffect, useRef } from "react";
22
import { useWorkspaceStoreRaw, type WorkspaceState } from "@/browser/stores/WorkspaceStore";
33
import { CUSTOM_EVENTS, type CustomEventType } from "@/common/constants/events";
4-
import { getAutoRetryKey, getRetryStateKey } from "@/common/constants/storage";
4+
import {
5+
getAutoRetryKey,
6+
getRetryStateKey,
7+
getCancelledCompactionKey,
8+
} from "@/common/constants/storage";
59
import { getSendOptionsFromStorage } from "@/browser/utils/messages/sendOptions";
610
import { readPersistedState, updatePersistedState } from "./usePersistedState";
711
import {
812
isEligibleForAutoRetry,
913
isNonRetryableSendError,
1014
} from "@/browser/utils/messages/retryEligibility";
11-
import { applyCompactionOverrides } from "@/browser/utils/messages/compactionOptions";
15+
import { executeCompaction } from "@/browser/utils/chatCommands";
1216
import type { SendMessageError } from "@/common/types/errors";
1317
import {
1418
createFailedRetryState,
@@ -23,6 +27,15 @@ export interface RetryState {
2327
lastError?: SendMessageError;
2428
}
2529

30+
/**
31+
* Persisted marker for user-cancelled compaction.
32+
* Used to distinguish intentional cancellation (Ctrl+C) from crash/force-exit.
33+
*/
34+
export interface CancelledCompactionMarker {
35+
messageId: string;
36+
timestamp: number;
37+
}
38+
2639
/**
2740
* Centralized auto-resume manager for interrupted streams
2841
*
@@ -163,44 +176,77 @@ export function useResumeManager() {
163176
);
164177

165178
try {
179+
if (!api) {
180+
retryingRef.current.delete(workspaceId);
181+
return;
182+
}
183+
166184
// Start with workspace defaults
167-
let options = getSendOptionsFromStorage(workspaceId);
185+
const options = getSendOptionsFromStorage(workspaceId);
168186

169187
// Check if last user message was a compaction request
170188
const state = workspaceStatesRef.current.get(workspaceId);
171-
if (state) {
172-
const lastUserMsg = [...state.messages].reverse().find((msg) => msg.type === "user");
173-
if (lastUserMsg?.compactionRequest) {
174-
// Apply compaction overrides using shared function (same as ChatInput)
175-
// This ensures custom model/tokens are preserved across resume
176-
options = applyCompactionOverrides(options, {
177-
model: lastUserMsg.compactionRequest.parsed.model,
178-
maxOutputTokens: lastUserMsg.compactionRequest.parsed.maxOutputTokens,
179-
continueMessage: {
180-
text: lastUserMsg.compactionRequest.parsed.continueMessage?.text ?? "",
181-
imageParts: lastUserMsg.compactionRequest.parsed.continueMessage?.imageParts,
182-
model: lastUserMsg.compactionRequest.parsed.continueMessage?.model ?? options.model,
183-
},
184-
});
189+
const lastUserMsg = state?.messages
190+
? [...state.messages].reverse().find((msg) => msg.type === "user")
191+
: undefined;
192+
193+
if (lastUserMsg?.compactionRequest) {
194+
// Check if this compaction was user-cancelled (Ctrl+C)
195+
const cancelledMarker = readPersistedState<CancelledCompactionMarker | null>(
196+
getCancelledCompactionKey(workspaceId),
197+
null
198+
);
199+
200+
if (cancelledMarker && cancelledMarker.messageId === lastUserMsg.id) {
201+
// User explicitly cancelled this compaction - don't auto-retry
202+
// Clear the marker (one-shot) so manual retry works
203+
updatePersistedState(getCancelledCompactionKey(workspaceId), () => null);
204+
console.debug(
205+
`[retry] ${workspaceId} skipping cancelled compaction (messageId=${lastUserMsg.id})`
206+
);
207+
return;
185208
}
186-
}
187209

188-
if (!api) {
189-
retryingRef.current.delete(workspaceId);
190-
return;
191-
}
192-
const result = await api.workspace.resumeStream({ workspaceId, options });
210+
// Retry compaction via executeCompaction (re-sends the compaction request)
211+
// This properly rebuilds the compaction-specific behavior including continueMessage queuing
212+
console.debug(`[retry] ${workspaceId} retrying interrupted compaction`);
213+
const { parsed } = lastUserMsg.compactionRequest;
214+
const result = await executeCompaction({
215+
api,
216+
workspaceId,
217+
sendMessageOptions: options,
218+
model: parsed.model,
219+
maxOutputTokens: parsed.maxOutputTokens,
220+
continueMessage: parsed.continueMessage,
221+
editMessageId: lastUserMsg.id, // Edit the existing compaction request message
222+
});
193223

194-
if (!result.success) {
195-
// Store error in retry state so RetryBarrier can display it
196-
const newState = createFailedRetryState(attempt, result.error);
197-
console.debug(
198-
`[retry] ${workspaceId} resumeStream failed: attempt ${attempt}${newState.attempt}`
199-
);
200-
updatePersistedState(getRetryStateKey(workspaceId), newState);
224+
if (!result.success) {
225+
const errorData: SendMessageError = {
226+
type: "unknown",
227+
raw: result.error ?? "Failed to retry compaction",
228+
};
229+
const newState = createFailedRetryState(attempt, errorData);
230+
console.debug(
231+
`[retry] ${workspaceId} compaction failed: attempt ${attempt}${newState.attempt}`
232+
);
233+
updatePersistedState(getRetryStateKey(workspaceId), newState);
234+
}
235+
} else {
236+
// Normal stream resume (non-compaction)
237+
const result = await api.workspace.resumeStream({ workspaceId, options });
238+
239+
if (!result.success) {
240+
// Store error in retry state so RetryBarrier can display it
241+
const newState = createFailedRetryState(attempt, result.error);
242+
console.debug(
243+
`[retry] ${workspaceId} resumeStream failed: attempt ${attempt}${newState.attempt}`
244+
);
245+
updatePersistedState(getRetryStateKey(workspaceId), newState);
246+
}
201247
}
202248
// Note: Don't clear retry state on success - stream-end event will handle that
203-
// resumeStream success just means "stream initiated", not "stream completed"
249+
// resumeStream/executeCompaction success just means "stream initiated", not "stream completed"
204250
// Clearing here causes backoff reset bug when stream starts then immediately fails
205251
} catch (error) {
206252
// Store error in retry state for display

src/browser/utils/chatCommands.ts

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,8 @@ import {
2626
resolveCompactionModel,
2727
isValidModelFormat,
2828
} from "@/browser/utils/messages/compactionModelPreference";
29+
import { getCancelledCompactionKey } from "@/common/constants/storage";
30+
import { updatePersistedState } from "@/browser/hooks/usePersistedState";
2931
import type { ImageAttachment } from "../components/ImageAttachments";
3032
import { dispatchWorkspaceSwitch } from "./workspaceEvents";
3133
import { getRuntimeKey, copyWorkspaceStorage } from "@/common/constants/storage";
@@ -691,6 +693,10 @@ export function prepareCompactionMessage(options: CompactionOptions): {
691693
export async function executeCompaction(
692694
options: CompactionOptions & { api: RouterClient<AppRouter> }
693695
): Promise<CompactionResult> {
696+
// Clear any cancelled-compaction marker since we're (re-)starting compaction
697+
// This allows auto-retry to work if this attempt is interrupted by crash/force-exit
698+
updatePersistedState(getCancelledCompactionKey(options.workspaceId), () => null);
699+
694700
const { messageText, metadata, sendOptions } = prepareCompactionMessage(options);
695701

696702
const result = await options.api.workspace.sendMessage({

src/browser/utils/compaction/handler.ts

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,9 @@
77

88
import type { StreamingMessageAggregator } from "@/browser/utils/messages/StreamingMessageAggregator";
99
import type { APIClient } from "@/browser/contexts/API";
10+
import { getCancelledCompactionKey } from "@/common/constants/storage";
11+
import { updatePersistedState } from "@/browser/hooks/usePersistedState";
12+
import type { CancelledCompactionMarker } from "@/browser/hooks/useResumeManager";
1013

1114
/**
1215
* Check if the workspace is currently in a compaction stream
@@ -76,6 +79,14 @@ export async function cancelCompaction(
7679
return false;
7780
}
7881

82+
// Mark this compaction as user-cancelled so auto-retry doesn't pick it up
83+
// This distinguishes intentional Ctrl+C from crash/force-exit
84+
const marker: CancelledCompactionMarker = {
85+
messageId: compactionRequestMsg.id,
86+
timestamp: Date.now(),
87+
};
88+
updatePersistedState(getCancelledCompactionKey(workspaceId), () => marker);
89+
7990
// Interrupt stream with abandonPartial flag
8091
// Backend detects this and skips compaction (Ctrl+C flow)
8192
await client.workspace.interruptStream({ workspaceId, options: { abandonPartial: true } });

src/common/orpc/schemas/stream.ts

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -322,7 +322,7 @@ export const WorkspaceChatMessageSchema = z.discriminatedUnion("type", [
322322
UsageDeltaEventSchema,
323323
QueuedMessageChangedEventSchema,
324324
RestoreToInputEventSchema,
325-
// Idle compaction notification
325+
// Compaction notifications
326326
IdleCompactionNeededEventSchema,
327327
// Init events
328328
...WorkspaceInitEventSchema.def.options,

src/node/services/agentSession.ts

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -39,7 +39,6 @@ import type { PostCompactionAttachment, PostCompactionExclusions } from "@/commo
3939
import { TURNS_BETWEEN_ATTACHMENTS } from "@/common/constants/attachments";
4040
import { extractEditedFileDiffs } from "@/common/utils/messages/extractEditedFiles";
4141
import { isValidModelFormat } from "@/common/utils/ai/models";
42-
4342
/**
4443
* Tracked file state for detecting external edits.
4544
* Uses timestamp-based polling with diff injection.
@@ -298,6 +297,8 @@ export class AgentSession {
298297
workspaceId: this.workspaceId,
299298
message: { type: "caught-up" },
300299
});
300+
// Note: Aborted compaction recovery is handled by useResumeManager on the frontend,
301+
// which detects interrupted compaction-request messages and retries via executeCompaction.
301302
}
302303
}
303304

0 commit comments

Comments
 (0)