Skip to content

feat: make context window size model-aware #151

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
tylergraydev wants to merge 1 commit into
base: main
Choose a base branch
from
Open

feat: make context window size model-aware #151

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
18 changes: 9 additions & 9 deletions src/components/chat/Composer.test.tsx
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1134,7 +1134,7 @@ describe("Composer", () => {
it("shows compact warning when context is >= 90% full", () => {
useChatStore.setState({
sessionStats: {
"ws-1": { totalInputTokens: 190_000, totalOutputTokens: 0, totalCostUsd: 1.5, numTurns: 10 },
"ws-1": { totalInputTokens: 950_000, totalOutputTokens: 0, totalCostUsd: 1.5, numTurns: 10 },
},
});
render(<Composer {...defaultProps} />);
Expand All @@ -1145,7 +1145,7 @@ describe("Composer", () => {
const onSend = vi.fn();
useChatStore.setState({
sessionStats: {
"ws-1": { totalInputTokens: 195_000, totalOutputTokens: 0, totalCostUsd: 2.0, numTurns: 15 },
"ws-1": { totalInputTokens: 950_000, totalOutputTokens: 0, totalCostUsd: 2.0, numTurns: 15 },
},
});
const user = userEvent.setup();
Expand All @@ -1157,7 +1157,7 @@ describe("Composer", () => {
it("does not show compact warning when context < 90%", () => {
useChatStore.setState({
sessionStats: {
"ws-1": { totalInputTokens: 100_000, totalOutputTokens: 0, totalCostUsd: 0.5, numTurns: 5 },
"ws-1": { totalInputTokens: 500_000, totalOutputTokens: 0, totalCostUsd: 0.5, numTurns: 5 },
},
});
render(<Composer {...defaultProps} />);
Expand All @@ -1167,7 +1167,7 @@ describe("Composer", () => {
it("does not show compact warning when agent is running", () => {
useChatStore.setState({
sessionStats: {
"ws-1": { totalInputTokens: 195_000, totalOutputTokens: 0, totalCostUsd: 2.0, numTurns: 15 },
"ws-1": { totalInputTokens: 950_000, totalOutputTokens: 0, totalCostUsd: 2.0, numTurns: 15 },
},
});
render(<Composer {...defaultProps} agentStatus="Running" />);
Expand All @@ -1177,7 +1177,7 @@ describe("Composer", () => {
it("shows context usage ring when tokens > 0", () => {
useChatStore.setState({
sessionStats: {
"ws-1": { totalInputTokens: 50_000, totalOutputTokens: 1000, totalCostUsd: 0.25, numTurns: 3 },
"ws-1": { totalInputTokens: 250_000, totalOutputTokens: 1000, totalCostUsd: 0.25, numTurns: 3 },
},
});
render(<Composer {...defaultProps} />);
Expand All @@ -1187,7 +1187,7 @@ describe("Composer", () => {
it("applies warning color to ring when context is 75-89% full", () => {
useChatStore.setState({
sessionStats: {
"ws-1": { totalInputTokens: 160_000, totalOutputTokens: 0, totalCostUsd: 1.0, numTurns: 8 },
"ws-1": { totalInputTokens: 800_000, totalOutputTokens: 0, totalCostUsd: 1.0, numTurns: 8 },
},
});
render(<Composer {...defaultProps} />);
Expand All @@ -1201,7 +1201,7 @@ describe("Composer", () => {
it("shows tooltip with context details on hover", async () => {
useChatStore.setState({
sessionStats: {
"ws-1": { totalInputTokens: 50_000, totalOutputTokens: 2000, totalCostUsd: 0.25, numTurns: 3 },
"ws-1": { totalInputTokens: 250_000, totalOutputTokens: 2000, totalCostUsd: 0.25, numTurns: 3 },
},
});
const user = userEvent.setup();
Expand All @@ -1216,7 +1216,7 @@ describe("Composer", () => {
it("shows cache read row in tooltip when cache tokens > 0", async () => {
useChatStore.setState({
sessionStats: {
"ws-1": { totalInputTokens: 50_000, totalOutputTokens: 2000, totalCostUsd: 0.25, numTurns: 3, totalCacheReadTokens: 10_000 },
"ws-1": { totalInputTokens: 250_000, totalOutputTokens: 2000, totalCostUsd: 0.25, numTurns: 3, totalCacheReadTokens: 10_000 },
},
});
const user = userEvent.setup();
Expand Down Expand Up @@ -1762,7 +1762,7 @@ describe("Composer", () => {
it("hides context tooltip on mouse leave", async () => {
useChatStore.setState({
sessionStats: {
"ws-1": { totalInputTokens: 50_000, totalOutputTokens: 1000, totalCostUsd: 0.25, numTurns: 3 },
"ws-1": { totalInputTokens: 250_000, totalOutputTokens: 1000, totalCostUsd: 0.25, numTurns: 3 },
},
});
const user = userEvent.setup();
Expand Down
9 changes: 5 additions & 4 deletions src/components/chat/Composer.tsx
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -12,7 +12,8 @@ import { useVoiceInput } from "../../hooks/useVoiceInput";
import { useToastStore } from "../../stores/toastStore";
import { ActionBar, ActionBarButton } from "./ActionBar";
import { QuestionCard } from "./QuestionCard";
import { ContextUsageIndicator, CONTEXT_WINDOW_TOKENS } from "./ContextUsageIndicator";
import { ContextUsageIndicator } from "./ContextUsageIndicator";
import { getContextWindow } from "../../lib/models";
import { FileChipIcon } from "./FileChipIcon";
import { useFileDropHandler } from "../../hooks/useFileDropHandler";
import { useSlashCommandAutocomplete } from "../../hooks/useSlashCommandAutocomplete";
Expand Down Expand Up @@ -469,10 +470,10 @@ export function Composer({ contextId, contextType, agentStatus, onSend, onStop,
)}

{/* Context near-full warning with compact button */}
{sessionStats && sessionStats.totalInputTokens > 0 && (sessionStats.totalInputTokens / CONTEXT_WINDOW_TOKENS) >= 0.9 && agentStatus === "Idle" && (
{sessionStats && sessionStats.totalInputTokens > 0 && (sessionStats.totalInputTokens / getContextWindow(selectedModel)) >= 0.9 && agentStatus === "Idle" && (
<ActionBar
icon={<Brain className="h-4 w-4 flex-shrink-0" style={{ color: "var(--error)" }} />}
description={<span>Context window is {Math.round((sessionStats.totalInputTokens / CONTEXT_WINDOW_TOKENS) * 100)}% full. Compact to free space.</span>}
description={<span>Context window is {Math.round((sessionStats.totalInputTokens / getContextWindow(selectedModel)) * 100)}% full. Compact to free space.</span>}
bgStyle={{ backgroundColor: "color-mix(in srgb, var(--error) 10%, transparent)" }}
primaryAction={
<button
Expand Down Expand Up @@ -723,7 +724,7 @@ export function Composer({ contextId, contextType, agentStatus, onSend, onStop,
{/* Right side: Context ring, Plus button, Send button */}
<div className="flex items-center gap-1.5">
{sessionStats && sessionStats.totalInputTokens > 0 && (
<ContextUsageIndicator stats={sessionStats} workspaceId={workspaceId} />
<ContextUsageIndicator stats={sessionStats} contextId={contextId} model={selectedModel} />
)}
<div className="relative" ref={plusMenuRef}>
<button
Expand Down
38 changes: 19 additions & 19 deletions src/components/chat/ContextBreakdownPopover.test.tsx
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -148,7 +148,7 @@

it("renders the popover with header and stats", () => {
render(
<ContextBreakdownPopover stats={baseStats} workspaceId="ws-1" onClose={onClose} />,
<ContextBreakdownPopover stats={baseStats} contextId="ws-1" contextWindowTokens={200_000} onClose={onClose} />,
);
expect(screen.getByText("Context Breakdown")).toBeInTheDocument();
expect(screen.getByText("5")).toBeInTheDocument(); // turns
Expand All @@ -157,7 +157,7 @@

it("shows correct cache hit rate", () => {
render(
<ContextBreakdownPopover stats={baseStats} workspaceId="ws-1" onClose={onClose} />,
<ContextBreakdownPopover stats={baseStats} contextId="ws-1" contextWindowTokens={200_000} onClose={onClose} />,
);
// 40000 / 80000 = 50%
expect(screen.getByText("50%")).toBeInTheDocument();
Expand All @@ -166,15 +166,15 @@
it("shows 0% cache hit rate when no input tokens", () => {
const stats: SessionStats = { ...baseStats, totalInputTokens: 0, totalCacheReadTokens: 0 };
render(
<ContextBreakdownPopover stats={stats} workspaceId="ws-1" onClose={onClose} />,
<ContextBreakdownPopover stats={stats} contextId="ws-1" onClose={onClose} />,

Check failure on line 169 in src/components/chat/ContextBreakdownPopover.test.tsx

View workflow job for this annotation

GitHub Actions / check 

Property 'contextWindowTokens' is missing in type '{ stats: SessionStats; contextId: string; onClose: Mock<Procedure>; }' but required in type 'ContextBreakdownPopoverProps'.
);
// Both the donut center (0% used) and cache hit (0%) show "0%" — just confirm at least one exists
expect(screen.getAllByText("0%").length).toBeGreaterThanOrEqual(1);
});

it("renders donut chart with correct segment values", () => {
render(
<ContextBreakdownPopover stats={baseStats} workspaceId="ws-1" onClose={onClose} />,
<ContextBreakdownPopover stats={baseStats} contextId="ws-1" contextWindowTokens={200_000} onClose={onClose} />,
);
expect(screen.getByTestId("pie-chart")).toBeInTheDocument();
// Fresh input = 80000 - 40000 = 40000
Expand All @@ -188,7 +188,7 @@
it("does not render bar chart with 0 or 1 turn", () => {
mockMessages.push(makeAssistantMsg(5000, 1000));
render(
<ContextBreakdownPopover stats={baseStats} workspaceId="ws-1" onClose={onClose} />,
<ContextBreakdownPopover stats={baseStats} contextId="ws-1" contextWindowTokens={200_000} onClose={onClose} />,
);
expect(screen.queryByTestId("bar-chart")).not.toBeInTheDocument();
});
Expand All @@ -197,14 +197,14 @@
mockMessages.push(makeAssistantMsg(5000, 1000));
mockMessages.push(makeAssistantMsg(10000, 2000));
render(
<ContextBreakdownPopover stats={baseStats} workspaceId="ws-1" onClose={onClose} />,
<ContextBreakdownPopover stats={baseStats} contextId="ws-1" contextWindowTokens={200_000} onClose={onClose} />,
);
expect(screen.getByTestId("bar-chart")).toBeInTheDocument();
});

it("calls onClose when close button clicked", () => {
render(
<ContextBreakdownPopover stats={baseStats} workspaceId="ws-1" onClose={onClose} />,
<ContextBreakdownPopover stats={baseStats} contextId="ws-1" contextWindowTokens={200_000} onClose={onClose} />,
);
fireEvent.click(screen.getByLabelText("Close context breakdown"));
expect(onClose).toHaveBeenCalledOnce();
Expand All @@ -218,7 +218,7 @@
numTurns: 2,
};
render(
<ContextBreakdownPopover stats={stats} workspaceId="ws-1" onClose={onClose} />,
<ContextBreakdownPopover stats={stats} contextId="ws-1" onClose={onClose} />,

Check failure on line 221 in src/components/chat/ContextBreakdownPopover.test.tsx

View workflow job for this annotation

GitHub Actions / check 

Property 'contextWindowTokens' is missing in type '{ stats: SessionStats; contextId: string; onClose: Mock<Procedure>; }' but required in type 'ContextBreakdownPopoverProps'.
);
expect(screen.getByText("0%")).toBeInTheDocument(); // cache hit rate
expect(screen.getByText("Context Breakdown")).toBeInTheDocument();
Expand All @@ -235,26 +235,26 @@
});

it("shows tooltip on hover when popover is closed", () => {
render(<ContextUsageIndicator stats={baseStats} workspaceId="ws-1" />);
render(<ContextUsageIndicator stats={baseStats} contextId="ws-1" />);
const container = screen.getByLabelText(/Context usage/);
fireEvent.mouseEnter(container.closest("[class*='relative']")!);
expect(screen.getByText(/Context/)).toBeInTheDocument();
});

it("opens popover on click when workspaceId is provided", () => {
render(<ContextUsageIndicator stats={baseStats} workspaceId="ws-1" />);
it("opens popover on click when contextId is provided", () => {
render(<ContextUsageIndicator stats={baseStats} contextId="ws-1" />);
fireEvent.click(screen.getByRole("button"));
expect(screen.getByTestId("context-breakdown-popover")).toBeInTheDocument();
});

it("does not open popover when workspaceId is missing", () => {
it("does not open popover when contextId is missing", () => {
render(<ContextUsageIndicator stats={baseStats} />);
// No role="button" when no workspaceId
// No role="button" when no contextId
expect(screen.queryByRole("button")).not.toBeInTheDocument();
});

it("hides tooltip when popover is open", () => {
render(<ContextUsageIndicator stats={baseStats} workspaceId="ws-1" />);
render(<ContextUsageIndicator stats={baseStats} contextId="ws-1" />);
const wrapper = screen.getByLabelText(/Context usage/).closest("[class*='relative']")!;

// Open popover
Expand All @@ -270,7 +270,7 @@
});

it("closes popover on Escape key", () => {
render(<ContextUsageIndicator stats={baseStats} workspaceId="ws-1" />);
render(<ContextUsageIndicator stats={baseStats} contextId="ws-1" />);
fireEvent.click(screen.getByRole("button"));
expect(screen.getByTestId("context-breakdown-popover")).toBeInTheDocument();

Expand All @@ -281,7 +281,7 @@
it("closes popover on click outside", () => {
render(
<div>
<ContextUsageIndicator stats={baseStats} workspaceId="ws-1" />
<ContextUsageIndicator stats={baseStats} contextId="ws-1" />
<div data-testid="outside">Outside</div>
</div>,
);
Expand All @@ -293,7 +293,7 @@
});

it("toggles popover on repeated clicks", () => {
render(<ContextUsageIndicator stats={baseStats} workspaceId="ws-1" />);
render(<ContextUsageIndicator stats={baseStats} contextId="ws-1" />);
const btn = screen.getByRole("button");

fireEvent.click(btn);
Expand All @@ -304,13 +304,13 @@
});

it("opens popover with Enter key", () => {
render(<ContextUsageIndicator stats={baseStats} workspaceId="ws-1" />);
render(<ContextUsageIndicator stats={baseStats} contextId="ws-1" />);
fireEvent.keyDown(screen.getByRole("button"), { key: "Enter" });
expect(screen.getByTestId("context-breakdown-popover")).toBeInTheDocument();
});

it("opens popover with Space key", () => {
render(<ContextUsageIndicator stats={baseStats} workspaceId="ws-1" />);
render(<ContextUsageIndicator stats={baseStats} contextId="ws-1" />);
fireEvent.keyDown(screen.getByRole("button"), { key: " " });
expect(screen.getByTestId("context-breakdown-popover")).toBeInTheDocument();
});
Expand Down
14 changes: 7 additions & 7 deletions src/components/chat/ContextBreakdownPopover.tsx
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,6 @@ import { useThemeColors } from "../../hooks/useThemeColors";
import { useChatStore } from "../../stores/chatStore";
import type { SessionStats } from "../../stores/chatStore";
import type { ChatMessage } from "../../lib/tauri";
import { CONTEXT_WINDOW_TOKENS } from "./ContextUsageIndicator";

// ---------------------------------------------------------------------------
// Per-turn delta computation
Expand Down Expand Up @@ -146,13 +145,14 @@ function StatCell({ label, value }: { label: string; value: string }) {

interface ContextBreakdownPopoverProps {
stats: SessionStats;
workspaceId: string;
contextId: string;
contextWindowTokens: number;
onClose: () => void;
}

export function ContextBreakdownPopover({ stats, workspaceId, onClose }: ContextBreakdownPopoverProps) {
export function ContextBreakdownPopover({ stats, contextId, contextWindowTokens, onClose }: ContextBreakdownPopoverProps) {
const colors = useThemeColors();
const messages = useChatStore((s) => s.messages[workspaceId] ?? []);
const messages = useChatStore((s) => s.messages[contextId] ?? []);
const [visible, setVisible] = useState(false);

// Fade in on mount
Expand All @@ -161,10 +161,10 @@ export function ContextBreakdownPopover({ stats, workspaceId, onClose }: Context
}, []);

// ---- Donut data ----
const pct = Math.min(100, (stats.totalInputTokens / CONTEXT_WINDOW_TOKENS) * 100);
const pct = Math.min(100, (stats.totalInputTokens / contextWindowTokens) * 100);
const cached = stats.totalCacheReadTokens ?? 0;
const freshInput = Math.max(0, stats.totalInputTokens - cached);
const remaining = Math.max(0, CONTEXT_WINDOW_TOKENS - stats.totalInputTokens);
const remaining = Math.max(0, contextWindowTokens - stats.totalInputTokens);

const donutData = useMemo(
() => [
Expand Down Expand Up @@ -289,7 +289,7 @@ export function ContextBreakdownPopover({ stats, workspaceId, onClose }: Context
>
<StatCell
label="Context"
value={`${formatTokens(stats.totalInputTokens)} / ${formatTokens(CONTEXT_WINDOW_TOKENS)}`}
value={`${formatTokens(stats.totalInputTokens)} / ${formatTokens(contextWindowTokens)}`}
/>
<StatCell label="Output" value={formatTokens(stats.totalOutputTokens)} />
<StatCell label="Cache hit" value={`${cacheHitRate}%`} />
Expand Down
Loading
Loading