Move context limit banner from MessagesList to ChatInput (#2461)

e2e-tests/context_limit_banner.spec.ts

 import { test, Timeout } from "./helpers/test_helper";
 import { expect } from "@playwright/test";
 
-test("context limit banner appears and summarize works", async ({ po }) => {
+test("context limit banner shows 'running out' when near context limit", async ({
+  po,
+}) => {
   await po.setUp();
 
   // Send a message that triggers high token usage (110k tokens)
@@ -9,19 +11,19 @@ test("context limit banner appears and summarize works", async ({ po }) => {
   // which is below the 40k threshold to show the banner
   await po.sendPrompt("tc=context-limit-response [high-tokens=110000]");
 
-  // Verify the context limit banner appears
-  const contextLimitBanner = po.page.getByTestId("context-limit-banner");
+  // Verify the context limit banner appears inside the chat input container
+  const contextLimitBanner = po
+    .getChatInputContainer()
+    .getByTestId("context-limit-banner");
   await expect(contextLimitBanner).toBeVisible({ timeout: Timeout.MEDIUM });
 
-  // Verify banner text
+  // Verify banner text for near-limit case
   await expect(contextLimitBanner).toContainText(
-    "You're close to the context limit for this chat.",
+    "This chat context is running out",
   );
 
   // Click the summarize button
-  await contextLimitBanner
-    .getByRole("button", { name: "Summarize into new chat" })
-    .click();
+  await contextLimitBanner.getByRole("button", { name: "Summarize" }).click();
 
   // Wait for the new chat to load and message to complete
   await po.waitForChatCompletion();
@@ -30,6 +32,39 @@ test("context limit banner appears and summarize works", async ({ po }) => {
   await po.snapshotMessages();
 });
 
+test("context limit banner shows 'costs extra' for long context", async ({
+  po,
+}) => {
+  await po.setUp();
+
+  // Add a custom test model with a 1M context window so 250k tokens isn't "near limit"
+  await po.goToSettingsTab();
+  await po.addCustomTestModel({
+    name: "test-model-large-ctx",
+    contextWindow: 1_000_000,
+  });
+  await po.goToAppsTab();
+  await po.selectModel({
+    provider: "test-provider",
+    model: "test-model-large-ctx",
+  });
+
+  // Send a message with 250k tokens (above 200k threshold)
+  // With 1M context window, 750k tokens remaining > 40k threshold, so not "near limit"
+  await po.sendPrompt("tc=context-limit-response [high-tokens=250000]");
+
+  // Verify the context limit banner appears inside the chat input container
+  const contextLimitBanner = po
+    .getChatInputContainer()
+    .getByTestId("context-limit-banner");
+  await expect(contextLimitBanner).toBeVisible({ timeout: Timeout.MEDIUM });
+
+  // Verify banner text for long context case
+  await expect(contextLimitBanner).toContainText(
+    "Long chat context costs extra",
+  );
+});
+
 test("context limit banner does not appear when within limit", async ({
   po,
 }) => {
@@ -37,10 +72,12 @@ test("context limit banner does not appear when within limit", async ({
 
   // Send a message with low token usage (50k tokens)
   // With a 128k context window, this leaves 78k tokens remaining
-  // which is above the 40k threshold - banner should NOT appear
+  // which is above the 40k threshold AND below 200k - banner should NOT appear
   await po.sendPrompt("tc=context-limit-response [high-tokens=50000]");
 
-  // Verify the context limit banner does NOT appear
-  const contextLimitBanner = po.page.getByTestId("context-limit-banner");
+  // Verify the context limit banner does NOT appear in the chat input container
+  const contextLimitBanner = po
+    .getChatInputContainer()
+    .getByTestId("context-limit-banner");
   await expect(contextLimitBanner).not.toBeVisible();
 });

e2e-tests/helpers/test_helper.ts

@@ -999,26 +999,26 @@ export class PageObject {
   }
 
   async selectModel({ provider, model }: { provider: string; model: string }) {
-    await this.page.getByRole("button", { name: "Model: Auto" }).click();
+    await this.page.getByTestId("model-picker").click();
     await this.page.getByText(provider, { exact: true }).click();
     await this.page.getByText(model, { exact: true }).click();
   }
 
   async selectTestModel() {
-    await this.page.getByRole("button", { name: "Model: Auto" }).click();
+    await this.page.getByTestId("model-picker").click();
     await this.page.getByText("test-provider").click();
     await this.page.getByText("test-model").click();
   }
 
   async selectTestOllamaModel() {
-    await this.page.getByRole("button", { name: "Model: Auto" }).click();
+    await this.page.getByTestId("model-picker").click();
     await this.page.getByText("Local models").click();
     await this.page.getByText("Ollama", { exact: true }).click();
     await this.page.getByText("Testollama", { exact: true }).click();
   }
 
   async selectTestLMStudioModel() {
-    await this.page.getByRole("button", { name: "Model: Auto" }).click();
+    await this.page.getByTestId("model-picker").click();
     await this.page.getByText("Local models").click();
     await this.page.getByText("LM Studio", { exact: true }).click();
     // Both of the elements that match "lmstudio-model-1" are the same button, so we just pick the first.
@@ -1029,7 +1029,7 @@ export class PageObject {
   }
 
   async selectTestAzureModel() {
-    await this.page.getByRole("button", { name: "Model: Auto" }).click();
+    await this.page.getByTestId("model-picker").click();
     await this.page.getByText("Other AI providers").click();
     await this.page.getByText("Azure OpenAI", { exact: true }).click();
     await this.page.getByText("GPT-5", { exact: true }).click();
@@ -1074,6 +1074,24 @@ export class PageObject {
     await this.page.getByRole("button", { name: "Add Model" }).click();
   }
 
+  async addCustomTestModel({
+    name,
+    contextWindow,
+  }: {
+    name: string;
+    contextWindow?: number;
+  }) {
+    await this.page.getByRole("heading", { name: "test-provider" }).click();
+    await this.page.getByRole("button", { name: "Add Custom Model" }).click();
+    await this.page.getByRole("textbox", { name: "Model ID*" }).fill(name);
+    await this.page.getByRole("textbox", { name: "Model ID*" }).press("Tab");
+    await this.page.getByRole("textbox", { name: "Name*" }).fill(name);
+    if (contextWindow) {
+      await this.page.locator("#context-window").fill(String(contextWindow));
+    }
+    await this.page.getByRole("button", { name: "Add Model" }).click();
+  }
+
   async setUpTestProviderApiKey() {
     // Fill in a test API key for the custom provider
     await this.page

e2e-tests/snapshots/context_limit_banner.spec.ts_context-limit-banner-appears-and-summarize-works-1.aria.yml → e2e-tests/snapshots/context_limit_banner.spec.ts_context-limit-banner-shows-running-out-when-near-context-limit-1.aria.yml

@@ -6,9 +6,13 @@
 - img
 - text: file1.txt
 - paragraph: More EOM
-- button:
+- button "Copy":
   - img
 - img
+- text: test-model
+- img
 - text: less than a minute ago
+- button "Undo":
+  - img
 - button "Retry":
   - img
\ No newline at end of file

src/components/ModelPicker.tsx

@@ -167,6 +167,7 @@ export function ModelPicker() {
     <DropdownMenu open={open} onOpenChange={setOpen}>
       <DropdownMenuTrigger
         className="inline-flex items-center justify-center whitespace-nowrap rounded-md text-sm font-medium transition-colors focus-visible:outline-none focus-visible:ring-1 focus-visible:ring-ring disabled:pointer-events-none disabled:opacity-50 border border-input bg-background shadow-sm hover:bg-accent hover:text-accent-foreground h-8 max-w-[130px] px-1.5 text-xs-sm gap-2"
+        data-testid="model-picker"
         title={modelDisplayName}
       >
         <span className="truncate">

src/components/chat/ChatInput.tsx

@@ -75,6 +75,11 @@ import { VisualEditingChangesDialog } from "@/components/preview_panel/VisualEdi
 import { useUserBudgetInfo } from "@/hooks/useUserBudgetInfo";
 import { useQueryClient } from "@tanstack/react-query";
 import { queryKeys } from "@/lib/queryKeys";
+import {
+  ContextLimitBanner,
+  shouldShowContextLimitBanner,
+} from "./ContextLimitBanner";
+import { useCountTokens } from "@/hooks/useCountTokens";
 
 const showTokenBarAtom = atom(false);
 
@@ -157,6 +162,20 @@ export function ChatInput({ chatId }: { chatId?: number }) {
 
   const { userBudget } = useUserBudgetInfo();
 
+  // Token counting for context limit banner
+  const { result: tokenCountResult } = useCountTokens(
+    !isStreaming ? (chatId ?? null) : null,
+    "",
+  );
+
+  const showBanner =
+    !isStreaming &&
+    tokenCountResult &&
+    shouldShowContextLimitBanner({
+      totalTokens: tokenCountResult.actualMaxTokens,
+      contextWindow: tokenCountResult.contextWindow,
+    });
+
   useEffect(() => {
     if (error) {
       setShowError(true);
@@ -313,10 +332,17 @@ export function ChatInput({ chatId }: { chatId?: number }) {
         </div>
       )}
       <div className="p-4" data-testid="chat-input-container">
+        {/* Show context limit banner above chat input for visibility */}
+        {showBanner && tokenCountResult && (
+          <ContextLimitBanner
+            totalTokens={tokenCountResult.actualMaxTokens}
+            contextWindow={tokenCountResult.contextWindow}
+          />
+        )}
         <div
           className={`relative flex flex-col border border-border rounded-lg bg-(--background-lighter) shadow-sm ${
             isDraggingOver ? "ring-2 ring-blue-500 border-blue-500" : ""
-          }`}
+          } ${showBanner ? "rounded-t-none border-t-0" : ""}`}
           onDragOver={handleDragOver}
           onDragLeave={handleDragLeave}
           onDrop={handleDrop}

src/components/chat/ContextLimitBanner.tsx

@@ -3,17 +3,28 @@ import { Button } from "@/components/ui/button";
 import { useSummarizeInNewChat } from "./SummarizeInNewChatButton";
 
 const CONTEXT_LIMIT_THRESHOLD = 40_000;
+const LONG_CONTEXT_THRESHOLD = 200_000;
 
 interface ContextLimitBannerProps {
   totalTokens?: number | null;
   contextWindow?: number;
 }
 
-function formatTokenCount(count: number): string {
-  if (count >= 1000) {
-    return `${(count / 1000).toFixed(1)}k`.replace(".0k", "k");
+/** Check if the context limit banner should be shown */
+export function shouldShowContextLimitBanner({
+  totalTokens,
+  contextWindow,
+}: ContextLimitBannerProps): boolean {
+  if (!totalTokens || !contextWindow) {
+    return false;
   }
-  return count.toString();
+  // Show if long context (costs extra)
+  if (totalTokens > LONG_CONTEXT_THRESHOLD) {
+    return true;
+  }
+  // Show if close to context limit
+  const tokensRemaining = contextWindow - totalTokens;
+  return tokensRemaining <= CONTEXT_LIMIT_THRESHOLD;
 }
 
 export function ContextLimitBanner({
@@ -22,43 +33,34 @@ export function ContextLimitBanner({
 }: ContextLimitBannerProps) {
   const { handleSummarize } = useSummarizeInNewChat();
 
-  // Don't show banner if we don't have the necessary data
-  if (!totalTokens || !contextWindow) {
+  if (!shouldShowContextLimitBanner({ totalTokens, contextWindow })) {
     return null;
   }
 
-  // Check if we're within 40k tokens of the context limit
-  const tokensRemaining = contextWindow - totalTokens;
-  if (tokensRemaining > CONTEXT_LIMIT_THRESHOLD) {
-    return null;
-  }
+  const tokensRemaining = contextWindow! - totalTokens!;
+  const isNearLimit = tokensRemaining <= CONTEXT_LIMIT_THRESHOLD;
+  const message = isNearLimit
+    ? "This chat context is running out"
+    : "Long chat context costs extra";
 
   return (
     <div
-      className="mx-auto max-w-3xl my-3 p-2 rounded-lg border border-amber-500/30 bg-amber-500/10 flex flex-col gap-2"
+      className="mx-auto max-w-3xl px-3 py-1.5 rounded-t-md border-t border-l border-r border-amber-500/30 bg-amber-500/10 flex items-center justify-between gap-3 text-xs text-amber-600 dark:text-amber-500"
       data-testid="context-limit-banner"
     >
-      <div className="flex items-center gap-2 text-amber-600 dark:text-amber-400">
-        <Button
-          variant="ghost"
-          size="icon"
-          className="h-5 w-5 p-0 hover:bg-transparent text-amber-600 dark:text-amber-400 cursor-help"
-          title={`Used: ${formatTokenCount(totalTokens)} / Limit: ${formatTokenCount(contextWindow)}`}
-        >
-          <AlertTriangle className="h-4 w-4 shrink-0" />
-        </Button>
-        <p className="text-sm font-medium">
-          You're close to the context limit for this chat.
-        </p>
-      </div>
+      <span className="flex items-center gap-1.5">
+        <AlertTriangle className="h-3.5 w-3.5 shrink-0" />
+        <span>{message}</span>
+      </span>
       <Button
         onClick={handleSummarize}
         variant="outline"
         size="sm"
-        className="h-8 border-amber-500/50 hover:bg-amber-500/20 hover:border-amber-500 text-amber-600 dark:text-amber-400"
+        className="h-6 px-2 text-xs border-amber-500/40 bg-amber-500/5 text-amber-600 dark:text-amber-500 hover:bg-amber-500/20 hover:border-amber-500/60"
+        title="Summarize to new chat"
       >
-        Summarize into new chat
-        <ArrowRight className="h-3 w-3 ml-2" />
+        Summarize
+        <ArrowRight className="h-3 w-3 ml-1" />
       </Button>
     </div>
   );

src/components/chat/MessagesList.tsx

@@ -19,8 +19,6 @@ import { useLanguageModelProviders } from "@/hooks/useLanguageModelProviders";
 import { useSettings } from "@/hooks/useSettings";
 import { useUserBudgetInfo } from "@/hooks/useUserBudgetInfo";
 import { PromoMessage } from "./PromoMessage";
-import { ContextLimitBanner } from "./ContextLimitBanner";
-import { useCountTokens } from "@/hooks/useCountTokens";
 
 interface MessagesListProps {
   messages: Message[];
@@ -38,7 +36,6 @@ interface FooterContext {
   messages: Message[];
   messagesEndRef: React.RefObject<HTMLDivElement | null>;
   isStreaming: boolean;
-  tokenCountResult: ReturnType<typeof useCountTokens>["result"];
   isUndoLoading: boolean;
   isRetryLoading: boolean;
   setIsUndoLoading: (loading: boolean) => void;
@@ -62,7 +59,6 @@ function FooterComponent({ context }: { context?: FooterContext }) {
     messages,
     messagesEndRef,
     isStreaming,
-    tokenCountResult,
     isUndoLoading,
     isRetryLoading,
     setIsUndoLoading,
@@ -80,14 +76,6 @@ function FooterComponent({ context }: { context?: FooterContext }) {
 
   return (
     <>
-      {/* Show context limit banner when close to token limit */}
-      {!isStreaming && tokenCountResult && (
-        <ContextLimitBanner
-          totalTokens={tokenCountResult.actualMaxTokens}
-          contextWindow={tokenCountResult.contextWindow}
-        />
-      )}
-
       {!isStreaming && (
         <div className="flex max-w-3xl mx-auto gap-2">
           {!!messages.length &&
@@ -275,11 +263,6 @@ export const MessagesList = forwardRef<HTMLDivElement, MessagesListProps>(
     // 2. Tests would need complex scrolling logic to bring elements into view before interaction
     // 3. Race conditions and timing issues occur when waiting for virtualized elements to render after scrolling
     const isTestMode = settings?.isTestMode;
-    // Only fetch token count when not streaming
-    const { result: tokenCountResult } = useCountTokens(
-      !isStreaming ? selectedChatId : null,
-      "",
-    );
 
     // Wrap state setters in useCallback to stabilize references
     const handleSetIsUndoLoading = useCallback((loading: boolean) => {
@@ -335,7 +318,6 @@ export const MessagesList = forwardRef<HTMLDivElement, MessagesListProps>(
         messages,
         messagesEndRef,
         isStreaming,
-        tokenCountResult,
         isUndoLoading,
         isRetryLoading,
         setIsUndoLoading: handleSetIsUndoLoading,
@@ -354,7 +336,6 @@ export const MessagesList = forwardRef<HTMLDivElement, MessagesListProps>(
         messages,
         messagesEndRef,
         isStreaming,
-        tokenCountResult,
         isUndoLoading,
         isRetryLoading,
         handleSetIsUndoLoading,