Skip to content

Commit 087e85f

Browse files
ethanndicksonethanndickson
committed
feat: add auto-compaction with progressive warnings
1 parent 9e1f0aa commit 087e85f

File tree

11 files changed

+294
-56
lines changed

11 files changed

+294
-56
lines changed

‎src/browser/components/AIView.tsx‎

Lines changed: 24 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -23,14 +23,21 @@ import { formatKeybind, KEYBINDS } from "@/browser/utils/ui/keybinds";
2323
import { useAutoScroll } from "@/browser/hooks/useAutoScroll";
2424
import { usePersistedState } from "@/browser/hooks/usePersistedState";
2525
import { useThinking } from "@/browser/contexts/ThinkingContext";
26-
import { useWorkspaceState, useWorkspaceAggregator } from "@/browser/stores/WorkspaceStore";
26+
import {
27+
useWorkspaceState,
28+
useWorkspaceAggregator,
29+
useWorkspaceUsage,
30+
} from "@/browser/stores/WorkspaceStore";
2731
import { WorkspaceHeader } from "./WorkspaceHeader";
2832
import { getModelName } from "@/common/utils/ai/models";
2933
import type { DisplayedMessage } from "@/common/types/message";
3034
import type { RuntimeConfig } from "@/common/types/runtime";
3135
import { useAIViewKeybinds } from "@/browser/hooks/useAIViewKeybinds";
3236
import { evictModelFromLRU } from "@/browser/hooks/useModelLRU";
3337
import { QueuedMessage } from "./Messages/QueuedMessage";
38+
import { CompactionWarning } from "./CompactionWarning";
39+
import { shouldAutoCompact } from "@/browser/utils/compaction/autoCompactionCheck";
40+
import { use1MContext } from "@/browser/hooks/use1MContext";
3441

3542
interface AIViewProps {
3643
workspaceId: string;
@@ -74,6 +81,8 @@ const AIViewInner: React.FC<AIViewProps> = ({
7481

7582
const workspaceState = useWorkspaceState(workspaceId);
7683
const aggregator = useWorkspaceAggregator(workspaceId);
84+
const workspaceUsage = useWorkspaceUsage(workspaceId);
85+
const [use1M] = use1MContext();
7786
const handledModelErrorsRef = useRef<Set<string>>(new Set());
7887

7988
useEffect(() => {
@@ -318,6 +327,13 @@ const AIViewInner: React.FC<AIViewProps> = ({
318327
// Get active stream message ID for token counting
319328
const activeStreamMessageId = aggregator.getActiveStreamMessageId();
320329

330+
const autoCompactionCheck = currentModel
331+
? shouldAutoCompact(workspaceUsage, currentModel, use1M)
332+
: { shouldShowWarning: false, usagePercentage: 0, thresholdPercentage: 70 };
333+
334+
// Show warning when: shouldShowWarning flag is true AND not currently compacting
335+
const shouldShowCompactionWarning = !isCompacting && autoCompactionCheck.shouldShowWarning;
336+
321337
// Note: We intentionally do NOT reset autoRetry when streams start.
322338
// If user pressed the interrupt key, autoRetry stays false until they manually retry.
323339
// This makes state transitions explicit and predictable.
@@ -503,6 +519,12 @@ const AIViewInner: React.FC<AIViewProps> = ({
503519
</button>
504520
)}
505521
</div>
522+
{shouldShowCompactionWarning && (
523+
<CompactionWarning
524+
usagePercentage={autoCompactionCheck.usagePercentage}
525+
thresholdPercentage={autoCompactionCheck.thresholdPercentage}
526+
/>
527+
)}
506528
<ChatInput
507529
variant="workspace"
508530
workspaceId={workspaceId}
@@ -516,6 +538,7 @@ const AIViewInner: React.FC<AIViewProps> = ({
516538
onEditLastUserMessage={() => void handleEditLastUserMessage()}
517539
canInterrupt={canInterrupt}
518540
onReady={handleChatInputReady}
541+
autoCompactionCheck={autoCompactionCheck}
519542
/>
520543
</div>
521544

‎src/browser/components/ChatInput/index.tsx‎

Lines changed: 91 additions & 32 deletions
Original file line numberDiff line numberDiff line change
@@ -30,6 +30,7 @@ import {
3030
handleCompactCommand,
3131
forkWorkspace,
3232
prepareCompactionMessage,
33+
executeCompaction,
3334
type CommandHandlerContext,
3435
} from "@/browser/utils/chatCommands";
3536
import { CUSTOM_EVENTS } from "@/common/constants/events";
@@ -472,6 +473,32 @@ export const ChatInput: React.FC<ChatInputProps> = (props) => {
472473
// Workspace variant: full command handling + message send
473474
if (variant !== "workspace") return; // Type guard
474475

476+
// Prepare image parts if any
477+
const imageParts = imageAttachments.map((img, index) => {
478+
// Validate before sending to help with debugging
479+
if (!img.url || typeof img.url !== "string") {
480+
console.error(
481+
`Image attachment [${index}] has invalid url:`,
482+
typeof img.url,
483+
img.url?.slice(0, 50)
484+
);
485+
}
486+
if (!img.url?.startsWith("data:")) {
487+
console.error(`Image attachment [${index}] url is not a data URL:`, img.url?.slice(0, 100));
488+
}
489+
if (!img.mediaType || typeof img.mediaType !== "string") {
490+
console.error(
491+
`Image attachment [${index}] has invalid mediaType:`,
492+
typeof img.mediaType,
493+
img.mediaType
494+
);
495+
}
496+
return {
497+
url: img.url,
498+
mediaType: img.mediaType,
499+
};
500+
});
501+
475502
try {
476503
// Parse command
477504
const parsed = parseCommand(messageText);
@@ -571,8 +598,10 @@ export const ChatInput: React.FC<ChatInputProps> = (props) => {
571598
const context: CommandHandlerContext = {
572599
workspaceId: props.workspaceId,
573600
sendMessageOptions,
601+
imageParts,
574602
editMessageId: editingMessage?.id,
575603
setInput,
604+
setImageAttachments,
576605
setIsSending,
577606
setToast,
578607
onCancelEdit: props.onCancelEdit,
@@ -636,7 +665,9 @@ export const ChatInput: React.FC<ChatInputProps> = (props) => {
636665
const context: CommandHandlerContext = {
637666
workspaceId: props.workspaceId,
638667
sendMessageOptions,
668+
imageParts: undefined, // /new doesn't use images
639669
setInput,
670+
setImageAttachments,
640671
setIsSending,
641672
setToast,
642673
};
@@ -656,42 +687,70 @@ export const ChatInput: React.FC<ChatInputProps> = (props) => {
656687
}
657688
}
658689

659-
// Regular message - send directly via API
660-
setIsSending(true);
661-
662690
// Save current state for restoration on error
663691
const previousImageAttachments = [...imageAttachments];
664692

665-
try {
666-
// Prepare image parts if any
667-
const imageParts = imageAttachments.map((img, index) => {
668-
// Validate before sending to help with debugging
669-
if (!img.url || typeof img.url !== "string") {
670-
console.error(
671-
`Image attachment [${index}] has invalid url:`,
672-
typeof img.url,
673-
img.url?.slice(0, 50)
674-
);
675-
}
676-
if (!img.url?.startsWith("data:")) {
677-
console.error(
678-
`Image attachment [${index}] url is not a data URL:`,
679-
img.url?.slice(0, 100)
680-
);
681-
}
682-
if (!img.mediaType || typeof img.mediaType !== "string") {
683-
console.error(
684-
`Image attachment [${index}] has invalid mediaType:`,
685-
typeof img.mediaType,
686-
img.mediaType
687-
);
693+
// Auto-compaction check (workspace variant only)
694+
// Check if we should auto-compact before sending this message
695+
// Result is computed in parent (AIView) and passed down to avoid duplicate calculation
696+
const shouldAutoCompact =
697+
props.autoCompactionCheck &&
698+
props.autoCompactionCheck.usagePercentage >= props.autoCompactionCheck.thresholdPercentage;
699+
if (variant === "workspace" && !editingMessage && shouldAutoCompact) {
700+
// Clear input immediately for responsive UX
701+
setInput("");
702+
setImageAttachments([]);
703+
setIsSending(true);
704+
705+
try {
706+
const result = await executeCompaction({
707+
workspaceId: props.workspaceId,
708+
continueMessage: {
709+
text: messageText,
710+
imageParts,
711+
},
712+
sendMessageOptions,
713+
});
714+
715+
if (!result.success) {
716+
// Restore on error
717+
setInput(messageText);
718+
setImageAttachments(previousImageAttachments);
719+
setToast({
720+
id: Date.now().toString(),
721+
type: "error",
722+
title: "Auto-Compaction Failed",
723+
message: result.error ?? "Failed to start auto-compaction",
724+
});
725+
} else {
726+
setToast({
727+
id: Date.now().toString(),
728+
type: "success",
729+
message: `Context threshold reached - auto-compacting...`,
730+
});
688731
}
689-
return {
690-
url: img.url,
691-
mediaType: img.mediaType,
692-
};
693-
});
732+
} catch (error) {
733+
// Restore on unexpected error
734+
setInput(messageText);
735+
setImageAttachments(previousImageAttachments);
736+
setToast({
737+
id: Date.now().toString(),
738+
type: "error",
739+
title: "Auto-Compaction Failed",
740+
message:
741+
error instanceof Error ? error.message : "Unexpected error during auto-compaction",
742+
});
743+
} finally {
744+
setIsSending(false);
745+
}
694746

747+
return; // Skip normal send
748+
}
749+
750+
// Regular message - send directly via API
751+
setIsSending(true);
752+
753+
try {
695754
// When editing a /compact command, regenerate the actual summarization request
696755
let actualMessageText = messageText;
697756
let muxMetadata: MuxFrontendMetadata | undefined;
@@ -707,7 +766,7 @@ export const ChatInput: React.FC<ChatInputProps> = (props) => {
707766
} = prepareCompactionMessage({
708767
workspaceId: props.workspaceId,
709768
maxOutputTokens: parsed.maxOutputTokens,
710-
continueMessage: parsed.continueMessage,
769+
continueMessage: { text: parsed.continueMessage ?? "", imageParts },
711770
model: parsed.model,
712771
sendMessageOptions,
713772
});

‎src/browser/components/ChatInput/types.ts‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
import type { ImagePart } from "@/common/types/ipc";
22
import type { FrontendWorkspaceMetadata } from "@/common/types/workspace";
3+
import type { AutoCompactionCheckResult } from "@/browser/utils/compaction/autoCompactionCheck";
34

45
export interface ChatInputAPI {
56
focus: () => void;
@@ -23,6 +24,7 @@ export interface ChatInputWorkspaceVariant {
2324
canInterrupt?: boolean;
2425
disabled?: boolean;
2526
onReady?: (api: ChatInputAPI) => void;
27+
autoCompactionCheck?: AutoCompactionCheckResult; // Computed in parent (AIView) to avoid duplicate calculation
2628
}
2729

2830
// Creation variant: simplified for first message / workspace creation

‎src/browser/components/CompactionWarning.tsx‎

Lines changed: 36 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,36 @@
1+
import React from "react";
2+
3+
/**
4+
* Warning banner shown when context usage is approaching the compaction threshold.
5+
*
6+
* Displays progressive warnings:
7+
* - Below threshold: "Context left until Auto-Compact: X% remaining" (where X = threshold - current)
8+
* - At/above threshold: "Approaching context limit. Next message will trigger auto-compaction."
9+
*
10+
* Displayed above ChatInput when:
11+
* - Token usage >= (threshold - 10%) of model's context window
12+
* - Not currently compacting (user can still send messages)
13+
*
14+
* @param usagePercentage - Current token usage as percentage (0-100)
15+
* @param thresholdPercentage - Auto-compaction trigger threshold (0-100, default 70)
16+
*/
17+
export const CompactionWarning: React.FC<{
18+
usagePercentage: number;
19+
thresholdPercentage: number;
20+
}> = (props) => {
21+
// At threshold or above, next message will trigger compaction
22+
const willCompactNext = props.usagePercentage >= props.thresholdPercentage;
23+
24+
// Calculate remaining percentage until threshold
25+
const remaining = props.thresholdPercentage - props.usagePercentage;
26+
27+
const message = willCompactNext
28+
? "⚠️ Context limit reached. Next message will trigger auto-compaction."
29+
: `Context left until Auto-Compact: ${Math.round(remaining)}%`;
30+
31+
return (
32+
<div className="text-plan-mode bg-plan-mode/10 mx-4 my-4 rounded-sm px-4 py-3 text-center text-xs font-medium">
33+
{message}
34+
</div>
35+
);
36+
};

‎src/browser/hooks/useResumeManager.ts‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -171,7 +171,10 @@ export function useResumeManager() {
171171
if (lastUserMsg?.compactionRequest) {
172172
// Apply compaction overrides using shared function (same as ChatInput)
173173
// This ensures custom model/tokens are preserved across resume
174-
options = applyCompactionOverrides(options, lastUserMsg.compactionRequest.parsed);
174+
options = applyCompactionOverrides(options, {
175+
maxOutputTokens: lastUserMsg.compactionRequest.parsed.maxOutputTokens,
176+
continueMessage: { text: lastUserMsg.compactionRequest.parsed.continueMessage ?? "" },
177+
});
175178
}
176179
}
177180

‎src/browser/stores/WorkspaceStore.ts‎

Lines changed: 20 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -424,27 +424,34 @@ export class WorkspaceStore {
424424
* Extract usage from messages (no tokenization).
425425
* Each usage entry calculated with its own model for accurate costs.
426426
*
427-
* REQUIRES: Workspace must have been added via addWorkspace() first.
427+
* Returns empty state if workspace doesn't exist (e.g., creation mode).
428428
*/
429429
getWorkspaceUsage(workspaceId: string): WorkspaceUsageState {
430430
return this.usageStore.get(workspaceId, () => {
431-
const aggregator = this.assertGet(workspaceId);
431+
const aggregator = this.aggregators.get(workspaceId);
432+
if (!aggregator) {
433+
return { usageHistory: [], totalTokens: 0 };
434+
}
432435

433436
const messages = aggregator.getAllMessages();
434437
const model = aggregator.getCurrentModel();
435438
const usageHistory = collectUsageHistory(messages, model);
436439

437-
// Calculate total from usage history (now includes historical)
438-
const totalTokens = usageHistory.reduce(
439-
(sum, u) =>
440-
sum +
441-
u.input.tokens +
442-
u.cached.tokens +
443-
u.cacheCreate.tokens +
444-
u.output.tokens +
445-
u.reasoning.tokens,
446-
0
447-
);
440+
const messages = aggregator.getAllMessages();
441+
const model = aggregator.getCurrentModel();
442+
const usageHistory = cumUsageHistory(messages, model);
443+
444+
// Use last entry's total (each entry is cumulative, not a delta)
445+
// Each usageHistory entry contains the FULL prompt tokens for that turn,
446+
// so we only need the most recent value, not a sum
447+
const lastEntry = usageHistory[usageHistory.length - 1];
448+
const totalTokens = lastEntry
449+
? lastEntry.input.tokens +
450+
lastEntry.cached.tokens +
451+
lastEntry.cacheCreate.tokens +
452+
lastEntry.output.tokens +
453+
lastEntry.reasoning.tokens
454+
: 0;
448455

449456
return { usageHistory, totalTokens };
450457
});

0 commit comments

Comments
 (0)