Skip to content

Commit da7df6d

Browse files
committed
🤖 feat: add storybook story for auto-compact warning
- Add contextUsage parameter to createAssistantMessage for testing high token usage - Add AutoCompactWarning story showing the indicator at 65% context usage
1 parent 9584805 commit da7df6d

File tree

3 files changed

+120
-3
lines changed

3 files changed

+120
-3
lines changed

src/browser/components/CompactionWarning.tsx

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -47,7 +47,7 @@ export const CompactionWarning: React.FC<{
4747
}
4848

4949
return (
50-
<div className="mx-auto -mb-1 w-full max-w-4xl px-[15px] text-right text-[10px]">
50+
<div className="mx-auto mb-1 w-full max-w-4xl px-[15px] text-right text-[10px]">
5151
<button
5252
type="button"
5353
onClick={props.onCompactClick}

src/browser/stories/App.chat.stories.tsx

Lines changed: 110 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -398,6 +398,116 @@ export const GenericTool: AppStory = {
398398
},
399399
};
400400

401+
/** Shows auto-compact warning when context usage approaches threshold */
402+
export const AutoCompactWarning: AppStory = {
403+
render: () => (
404+
<AppWithMocks
405+
setup={() =>
406+
setupSimpleChatStory({
407+
workspaceId: "ws-compaction-warning",
408+
messages: [
409+
createUserMessage("msg-1", "Help me set up authentication for this app", {
410+
historySequence: 1,
411+
timestamp: STABLE_TIMESTAMP - 600000,
412+
}),
413+
createAssistantMessage(
414+
"msg-2",
415+
"I'll help you set up authentication. Let me first examine your current project structure and dependencies.",
416+
{
417+
historySequence: 2,
418+
timestamp: STABLE_TIMESTAMP - 580000,
419+
toolCalls: [
420+
createFileReadTool(
421+
"call-1",
422+
"package.json",
423+
'{\n "name": "my-app",\n "dependencies": {\n "express": "^4.18.0",\n "mongoose": "^7.0.0"\n }\n}'
424+
),
425+
],
426+
}
427+
),
428+
createUserMessage("msg-3", "I want to use JWT tokens for the auth", {
429+
historySequence: 3,
430+
timestamp: STABLE_TIMESTAMP - 500000,
431+
}),
432+
createAssistantMessage(
433+
"msg-4",
434+
"Great choice! I'll implement JWT-based authentication. Let me create the auth middleware and user routes.",
435+
{
436+
historySequence: 4,
437+
timestamp: STABLE_TIMESTAMP - 480000,
438+
toolCalls: [
439+
createFileEditTool(
440+
"call-2",
441+
"src/middleware/auth.ts",
442+
"+ import jwt from 'jsonwebtoken';\n+ export const authenticate = (req, res, next) => {\n+ const token = req.headers.authorization?.split(' ')[1];\n+ // ... verification logic\n+ };"
443+
),
444+
],
445+
}
446+
),
447+
createUserMessage("msg-5", "Add password hashing too", {
448+
historySequence: 5,
449+
timestamp: STABLE_TIMESTAMP - 400000,
450+
}),
451+
createAssistantMessage(
452+
"msg-6",
453+
"I'll add bcrypt for secure password hashing. Here's the updated user model with password hashing on save.",
454+
{
455+
historySequence: 6,
456+
timestamp: STABLE_TIMESTAMP - 380000,
457+
toolCalls: [
458+
createTerminalTool("call-3", "bun add bcrypt", "added bcrypt@5.1.0"),
459+
createFileEditTool(
460+
"call-4",
461+
"src/models/user.ts",
462+
"+ import bcrypt from 'bcrypt';\n+ userSchema.pre('save', async function() {\n+ this.password = await bcrypt.hash(this.password, 10);\n+ });"
463+
),
464+
],
465+
}
466+
),
467+
createUserMessage("msg-7", "Now add refresh token support", {
468+
historySequence: 7,
469+
timestamp: STABLE_TIMESTAMP - 300000,
470+
}),
471+
createAssistantMessage(
472+
"msg-8",
473+
"I've implemented refresh token support with secure rotation. The access token expires in 15 minutes and the refresh token in 7 days.",
474+
{
475+
historySequence: 8,
476+
timestamp: STABLE_TIMESTAMP - 280000,
477+
// High context usage to trigger warning (65% of 200k)
478+
contextUsage: { inputTokens: 130000, outputTokens: 5000, totalTokens: 135000 },
479+
}
480+
),
481+
],
482+
})
483+
}
484+
/>
485+
),
486+
play: async ({ canvasElement }: { canvasElement: HTMLElement }) => {
487+
const canvas = within(canvasElement);
488+
489+
// Wait for last assistant message to render (unique text from assistant)
490+
await waitFor(() => {
491+
canvas.getByText(/access token expires in 15 minutes/);
492+
});
493+
494+
// Wait a bit for auto-scroll to settle, then scroll up
495+
await new Promise((r) => setTimeout(r, 100));
496+
const messageWindow = canvasElement.querySelector('[data-testid="message-window"]');
497+
if (messageWindow) {
498+
messageWindow.scrollTop = messageWindow.scrollHeight - messageWindow.clientHeight - 150;
499+
}
500+
},
501+
parameters: {
502+
docs: {
503+
description: {
504+
story:
505+
"Displays the auto-compact warning indicator above the chat input when context usage is approaching the compaction threshold. The indicator shows remaining headroom and is clickable to insert /compact command.",
506+
},
507+
},
508+
},
509+
};
510+
401511
/** Streaming compaction with shimmer effect - tests GPU-accelerated animation */
402512
export const StreamingCompaction: AppStory = {
403513
render: () => (

src/browser/stories/mockFactory.ts

Lines changed: 9 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -203,6 +203,8 @@ export function createAssistantMessage(
203203
model?: string;
204204
reasoning?: string;
205205
toolCalls?: MuxPart[];
206+
/** Custom context usage for testing compaction warning (default: small values) */
207+
contextUsage?: { inputTokens: number; outputTokens: number; totalTokens: number };
206208
}
207209
): ChatMuxMessage {
208210
const parts: MuxPart[] = [];
@@ -213,6 +215,11 @@ export function createAssistantMessage(
213215
if (opts.toolCalls) {
214216
parts.push(...opts.toolCalls);
215217
}
218+
const contextUsage = opts.contextUsage ?? {
219+
inputTokens: 100,
220+
outputTokens: 50,
221+
totalTokens: 150,
222+
};
216223
return {
217224
type: "message",
218225
id,
@@ -222,8 +229,8 @@ export function createAssistantMessage(
222229
historySequence: opts.historySequence,
223230
timestamp: opts.timestamp ?? STABLE_TIMESTAMP,
224231
model: opts.model ?? DEFAULT_MODEL,
225-
usage: { inputTokens: 100, outputTokens: 50, totalTokens: 150 },
226-
contextUsage: { inputTokens: 100, outputTokens: 50, totalTokens: 150 },
232+
usage: contextUsage,
233+
contextUsage,
227234
duration: 1000,
228235
},
229236
};

0 commit comments

Comments
 (0)