From 88f5de68d8467d6b9bb8f55c418bf5111bc01841 Mon Sep 17 00:00:00 2001
From: Tyler Gray <wymngray@gmail.com>
Date: Fri, 10 Apr 2026 19:18:23 -0400
Subject: [PATCH] feat: make context window size model-aware

The context usage ring and breakdown popover now dynamically adjust
based on the selected model. Opus shows usage out of 1M tokens,
Sonnet/Haiku out of 200k, etc. Previously hardcoded to 200k for all
models.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
---
 src/components/chat/Composer.test.tsx         | 18 ++++-----
 src/components/chat/Composer.tsx              |  9 +++--
 .../chat/ContextBreakdownPopover.test.tsx     | 38 +++++++++----------
 .../chat/ContextBreakdownPopover.tsx          | 14 +++----
 src/components/chat/ContextUsageIndicator.tsx | 37 +++++++++---------
 src/lib/models.test.ts                        | 29 ++++++++++++++
 src/lib/models.ts                             | 16 ++++++++
 7 files changed, 105 insertions(+), 56 deletions(-)
 create mode 100644 src/lib/models.test.ts
 create mode 100644 src/lib/models.ts

diff --git a/src/components/chat/Composer.test.tsx b/src/components/chat/Composer.test.tsx
index 24a2f20..ca23e93 100644
--- a/src/components/chat/Composer.test.tsx
+++ b/src/components/chat/Composer.test.tsx
@@ -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} />);
@@ -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();
@@ -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} />);
@@ -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" />);
@@ -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} />);
@@ -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} />);
@@ -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();
@@ -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();
@@ -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();
diff --git a/src/components/chat/Composer.tsx b/src/components/chat/Composer.tsx
index 03b0510..46b88ff 100644
--- a/src/components/chat/Composer.tsx
+++ b/src/components/chat/Composer.tsx
@@ -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";
@@ -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
@@ -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
diff --git a/src/components/chat/ContextBreakdownPopover.test.tsx b/src/components/chat/ContextBreakdownPopover.test.tsx
index f1397dc..12c0ee2 100644
--- a/src/components/chat/ContextBreakdownPopover.test.tsx
+++ b/src/components/chat/ContextBreakdownPopover.test.tsx
@@ -148,7 +148,7 @@ describe("ContextBreakdownPopover", () => {
 
   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
@@ -157,7 +157,7 @@ describe("ContextBreakdownPopover", () => {
 
   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();
@@ -166,7 +166,7 @@ describe("ContextBreakdownPopover", () => {
   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} />,
     );
     // 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);
@@ -174,7 +174,7 @@ describe("ContextBreakdownPopover", () => {
 
   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
@@ -188,7 +188,7 @@ describe("ContextBreakdownPopover", () => {
   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();
   });
@@ -197,14 +197,14 @@ describe("ContextBreakdownPopover", () => {
     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();
@@ -218,7 +218,7 @@ describe("ContextBreakdownPopover", () => {
       numTurns: 2,
     };
     render(
-      <ContextBreakdownPopover stats={stats} workspaceId="ws-1" onClose={onClose} />,
+      <ContextBreakdownPopover stats={stats} contextId="ws-1" onClose={onClose} />,
     );
     expect(screen.getByText("0%")).toBeInTheDocument(); // cache hit rate
     expect(screen.getByText("Context Breakdown")).toBeInTheDocument();
@@ -235,26 +235,26 @@ describe("ContextUsageIndicator", () => {
   });
 
   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
@@ -270,7 +270,7 @@ describe("ContextUsageIndicator", () => {
   });
 
   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();
 
@@ -281,7 +281,7 @@ describe("ContextUsageIndicator", () => {
   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>,
     );
@@ -293,7 +293,7 @@ describe("ContextUsageIndicator", () => {
   });
 
   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);
@@ -304,13 +304,13 @@ describe("ContextUsageIndicator", () => {
   });
 
   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();
   });
diff --git a/src/components/chat/ContextBreakdownPopover.tsx b/src/components/chat/ContextBreakdownPopover.tsx
index 7b7410c..0b9e918 100644
--- a/src/components/chat/ContextBreakdownPopover.tsx
+++ b/src/components/chat/ContextBreakdownPopover.tsx
@@ -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
@@ -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
@@ -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(
     () => [
@@ -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}%`} />
diff --git a/src/components/chat/ContextUsageIndicator.tsx b/src/components/chat/ContextUsageIndicator.tsx
index bfc013a..c4d6c62 100644
--- a/src/components/chat/ContextUsageIndicator.tsx
+++ b/src/components/chat/ContextUsageIndicator.tsx
@@ -1,19 +1,19 @@
 import { useState, useRef, useEffect, useCallback } from "react";
 import { formatTokens, formatCost } from "../../lib/format";
 import type { SessionStats } from "../../stores/chatStore";
+import { getContextWindow } from "../../lib/models";
 import { ContextBreakdownPopover } from "./ContextBreakdownPopover";
 
-export const CONTEXT_WINDOW_TOKENS = 200_000;
-
-function ContextTooltip({ stats, pct, color }: {
+function ContextTooltip({ stats, pct, color, contextWindowTokens }: {
   stats: SessionStats;
   pct: number;
   color: string;
+  contextWindowTokens: number;
 }) {
   const rows: Array<{ label: string; value: string }> = [
     {
       label: "Context",
-      value: `${formatTokens(stats.totalInputTokens)} / ${formatTokens(CONTEXT_WINDOW_TOKENS)} (${pct.toFixed(0)}%)`,
+      value: `${formatTokens(stats.totalInputTokens)} / ${formatTokens(contextWindowTokens)} (${pct.toFixed(0)}%)`,
     },
     {
       label: "Output",
@@ -62,15 +62,17 @@ function ContextTooltip({ stats, pct, color }: {
 
 interface ContextUsageIndicatorProps {
   stats: SessionStats;
-  workspaceId?: string;
+  contextId?: string;
+  model?: string;
 }
 
-export function ContextUsageIndicator({ stats, workspaceId }: ContextUsageIndicatorProps) {
+export function ContextUsageIndicator({ stats, contextId, model }: ContextUsageIndicatorProps) {
   const [showTooltip, setShowTooltip] = useState(false);
   const [showPopover, setShowPopover] = useState(false);
   const containerRef = useRef<HTMLDivElement>(null);
 
-  const pct = Math.min(100, (stats.totalInputTokens / CONTEXT_WINDOW_TOKENS) * 100);
+  const contextWindowTokens = getContextWindow(model ?? "");
+  const pct = Math.min(100, (stats.totalInputTokens / contextWindowTokens) * 100);
   const isWarning = pct >= 75;
   const isCritical = pct >= 90;
 
@@ -83,19 +85,19 @@ export function ContextUsageIndicator({ stats, workspaceId }: ContextUsageIndica
   const dashoffset = circumference * (1 - pct / 100);
 
   const handleClick = useCallback(() => {
-    if (workspaceId) {
+    if (contextId) {
       setShowPopover((prev) => !prev);
     }
-  }, [workspaceId]);
+  }, [contextId]);
 
   const handleKeyDown = useCallback(
     (e: React.KeyboardEvent) => {
-      if ((e.key === "Enter" || e.key === " ") && workspaceId) {
+      if ((e.key === "Enter" || e.key === " ") && contextId) {
         e.preventDefault();
         setShowPopover((prev) => !prev);
       }
     },
-    [workspaceId],
+    [contextId],
   );
 
   // Close on Escape and click-outside
@@ -127,10 +129,10 @@ export function ContextUsageIndicator({ stats, workspaceId }: ContextUsageIndica
       onMouseLeave={() => setShowTooltip(false)}
     >
       <div
-        className={`flex items-center justify-center ${workspaceId ? "cursor-pointer" : "cursor-default"}`}
+        className={`flex items-center justify-center ${contextId ? "cursor-pointer" : "cursor-default"}`}
         style={{ width: 24, height: 24 }}
-        role={workspaceId ? "button" : undefined}
-        tabIndex={workspaceId ? 0 : undefined}
+        role={contextId ? "button" : undefined}
+        tabIndex={contextId ? 0 : undefined}
         aria-label={`Context usage: ${pct.toFixed(0)}%`}
         aria-expanded={showPopover}
         onClick={handleClick}
@@ -156,11 +158,12 @@ export function ContextUsageIndicator({ stats, workspaceId }: ContextUsageIndica
           />
         </svg>
       </div>
-      {showTooltip && !showPopover && <ContextTooltip stats={stats} pct={pct} color={color} />}
-      {showPopover && workspaceId && (
+      {showTooltip && !showPopover && <ContextTooltip stats={stats} pct={pct} color={color} contextWindowTokens={contextWindowTokens} />}
+      {showPopover && contextId && (
         <ContextBreakdownPopover
           stats={stats}
-          workspaceId={workspaceId}
+          contextId={contextId}
+          contextWindowTokens={contextWindowTokens}
           onClose={() => setShowPopover(false)}
         />
       )}
diff --git a/src/lib/models.test.ts b/src/lib/models.test.ts
new file mode 100644
index 0000000..d808d04
--- /dev/null
+++ b/src/lib/models.test.ts
@@ -0,0 +1,29 @@
+import { describe, it, expect } from "vitest";
+import { getContextWindow, DEFAULT_CONTEXT_WINDOW } from "./models";
+
+describe("getContextWindow", () => {
+  it("returns 1M for empty string (default/Opus)", () => {
+    expect(getContextWindow("")).toBe(1_000_000);
+  });
+
+  it("returns 1M for 'opus'", () => {
+    expect(getContextWindow("opus")).toBe(1_000_000);
+  });
+
+  it("returns 200k for 'sonnet'", () => {
+    expect(getContextWindow("sonnet")).toBe(200_000);
+  });
+
+  it("returns 200k for 'haiku'", () => {
+    expect(getContextWindow("haiku")).toBe(200_000);
+  });
+
+  it("returns 128k for 'o4-mini'", () => {
+    expect(getContextWindow("o4-mini")).toBe(128_000);
+  });
+
+  it("returns DEFAULT_CONTEXT_WINDOW for unknown model strings", () => {
+    expect(getContextWindow("unknown-model")).toBe(DEFAULT_CONTEXT_WINDOW);
+    expect(getContextWindow("claude-99")).toBe(200_000);
+  });
+});
diff --git a/src/lib/models.ts b/src/lib/models.ts
new file mode 100644
index 0000000..02b3167
--- /dev/null
+++ b/src/lib/models.ts
@@ -0,0 +1,16 @@
+export const MODEL_CONTEXT_WINDOWS: Record<string, number> = {
+  "": 1_000_000,           // Default = Opus 4.6
+  "opus": 1_000_000,       // Claude Opus 4.6
+  "sonnet": 200_000,       // Claude Sonnet
+  "haiku": 200_000,        // Claude Haiku
+  "gpt-5.1-codex": 200_000,
+  "o3": 200_000,
+  "o4-mini": 128_000,
+  "gpt-4.1": 128_000,
+};
+
+export const DEFAULT_CONTEXT_WINDOW = 200_000;
+
+export function getContextWindow(model: string): number {
+  return MODEL_CONTEXT_WINDOWS[model] ?? DEFAULT_CONTEXT_WINDOW;
+}