diff --git a/e2e-tests/helpers/test_helper.ts b/e2e-tests/helpers/test_helper.ts
index b9fbbe5..ec3a239 100644
--- a/e2e-tests/helpers/test_helper.ts
+++ b/e2e-tests/helpers/test_helper.ts
@@ -132,10 +132,7 @@ export class PageObject {
   }
 
   async openContextFilesPicker() {
-    const contextButton = this.page.getByRole("button", {
-      name: "Context",
-      exact: true,
-    });
+    const contextButton = this.page.getByTestId("codebase-context-button");
     await contextButton.click();
     return new ContextFilesPickerDialog(this.page, async () => {
       await contextButton.click();
diff --git a/src/components/ChatInputControls.tsx b/src/components/ChatInputControls.tsx
index 1997638..3928118 100644
--- a/src/components/ChatInputControls.tsx
+++ b/src/components/ChatInputControls.tsx
@@ -8,10 +8,17 @@ export function ChatInputControls({
   showContextFilesPicker?: boolean;
 }) {
   return (
-    <div className="pb-2 flex gap-2">
+    <div className="flex">
       <ModelPicker />
-      {showContextFilesPicker && <ContextFilesPicker />}
+      <div className="w-2"></div>
       <ProModeSelector />
+      <div className="w-1"></div>
+      {showContextFilesPicker && (
+        <>
+          <ContextFilesPicker />
+          <div className="w-0.5"></div>
+        </>
+      )}
     </div>
   );
 }
diff --git a/src/components/ContextFilesPicker.tsx b/src/components/ContextFilesPicker.tsx
index 289b5d2..839fd8e 100644
--- a/src/components/ContextFilesPicker.tsx
+++ b/src/components/ContextFilesPicker.tsx
@@ -6,7 +6,7 @@ import {
   PopoverTrigger,
 } from "@/components/ui/popover";
 
-import { FileCode, InfoIcon, Trash2 } from "lucide-react";
+import { InfoIcon, Settings2, Trash2 } from "lucide-react";
 import { useState } from "react";
 import {
   Tooltip,
@@ -84,12 +84,22 @@ export function ContextFilesPicker() {
 
   return (
     <Popover open={isOpen} onOpenChange={setIsOpen}>
-      <PopoverTrigger asChild>
-        <Button variant="ghost" className="gap-2">
-          <FileCode className="size-4" />
-          <span>Context</span>
-        </Button>
-      </PopoverTrigger>
+      <Tooltip>
+        <TooltipTrigger asChild>
+          <PopoverTrigger asChild>
+            <Button
+              variant="ghost"
+              className="has-[>svg]:px-2"
+              size="sm"
+              data-testid="codebase-context-button"
+            >
+              <Settings2 className="size-4" />
+            </Button>
+          </PopoverTrigger>
+        </TooltipTrigger>
+        <TooltipContent>Codebase Context</TooltipContent>
+      </Tooltip>
+
       <PopoverContent className="w-96" align="start">
         <div className="relative space-y-4">
           <div>
diff --git a/src/components/ModelPicker.tsx b/src/components/ModelPicker.tsx
index c4fe2cf..c7be1c9 100644
--- a/src/components/ModelPicker.tsx
+++ b/src/components/ModelPicker.tsx
@@ -123,25 +123,34 @@ export function ModelPicker() {
 
   return (
     <DropdownMenu open={open} onOpenChange={setOpen}>
-      <DropdownMenuTrigger asChild>
-        <Button
-          variant="outline"
-          size="sm"
-          className="flex items-center gap-2 h-8"
-        >
-          <span>
-            {modelDisplayName === "Auto" && (
-              <>
-                <span className="text-xs text-muted-foreground">
-                  Model:
-                </span>{" "}
-              </>
-            )}
-            {modelDisplayName}
-          </span>
-        </Button>
-      </DropdownMenuTrigger>
-      <DropdownMenuContent className="w-64" align="start">
+      <Tooltip>
+        <TooltipTrigger asChild>
+          <DropdownMenuTrigger asChild>
+            <Button
+              variant="outline"
+              size="sm"
+              className="flex items-center gap-2 h-8 max-w-[160px] px-2"
+            >
+              <span className="truncate">
+                {modelDisplayName === "Auto" && (
+                  <>
+                    <span className="text-xs text-muted-foreground">
+                      Model:
+                    </span>{" "}
+                  </>
+                )}
+                {modelDisplayName}
+              </span>
+            </Button>
+          </DropdownMenuTrigger>
+        </TooltipTrigger>
+        <TooltipContent>{modelDisplayName}</TooltipContent>
+      </Tooltip>
+      <DropdownMenuContent
+        className="w-64"
+        align="start"
+        onCloseAutoFocus={(e) => e.preventDefault()}
+      >
         <DropdownMenuLabel>Cloud Models</DropdownMenuLabel>
         <DropdownMenuSeparator />
 
diff --git a/src/components/ProModeSelector.tsx b/src/components/ProModeSelector.tsx
index 1a95a37..c3255c9 100644
--- a/src/components/ProModeSelector.tsx
+++ b/src/components/ProModeSelector.tsx
@@ -40,7 +40,7 @@ export function ProModeSelector() {
             <Button
               variant="outline"
               size="sm"
-              className="flex items-center gap-2 h-8 border-primary/50 bg-primary/10 hover:bg-primary/20 font-medium shadow-sm shadow-primary/10 transition-all hover:shadow-md hover:shadow-primary/15"
+              className="has-[>svg]:px-2 flex items-center gap-1.5 h-8 border-primary/50 hover:bg-primary/10 font-medium shadow-sm shadow-primary/10 transition-all hover:shadow-md hover:shadow-primary/15"
             >
               <Sparkles className="h-4 w-4 text-primary" />
               <span className="text-primary font-medium">Pro</span>
diff --git a/src/components/chat/ChatInput.tsx b/src/components/chat/ChatInput.tsx
index c85ca2f..80e2fcf 100644
--- a/src/components/chat/ChatInput.tsx
+++ b/src/components/chat/ChatInput.tsx
@@ -1,5 +1,4 @@
 import {
-  SendIcon,
   StopCircleIcon,
   X,
   ChevronDown,
@@ -15,8 +14,9 @@ import {
   Database,
   ChevronsUpDown,
   ChevronsDownUp,
-  BarChart2,
   Paperclip,
+  ChartColumnIncreasing,
+  SendHorizontalIcon,
 } from "lucide-react";
 import type React from "react";
 import { useCallback, useEffect, useRef, useState } from "react";
@@ -313,28 +313,10 @@ export function ChatInput({ chatId }: { chatId?: number }) {
               disabled={isStreaming}
             />
 
-            {/* File attachment button */}
-            <button
-              onClick={handleAttachmentClick}
-              className="px-2 py-2 mt-1 mr-1 hover:bg-(--background-darkest) text-(--sidebar-accent-fg) rounded-lg disabled:opacity-50"
-              disabled={isStreaming}
-              title="Attach files"
-            >
-              <Paperclip size={20} />
-            </button>
-            <input
-              type="file"
-              ref={fileInputRef}
-              onChange={handleFileChange}
-              className="hidden"
-              multiple
-              accept=".jpg,.jpeg,.png,.gif,.webp,.txt,.md,.js,.ts,.html,.css,.json,.csv"
-            />
-
             {isStreaming ? (
               <button
                 onClick={handleCancel}
-                className="px-2 py-2 mt-1 mr-2 hover:bg-(--background-darkest) text-(--sidebar-accent-fg) rounded-lg"
+                className="px-2 py-2 mt-1 mr-1 hover:bg-(--background-darkest) text-(--sidebar-accent-fg) rounded-lg"
                 title="Cancel generation"
               >
                 <StopCircleIcon size={20} />
@@ -343,23 +325,62 @@ export function ChatInput({ chatId }: { chatId?: number }) {
               <button
                 onClick={handleSubmit}
                 disabled={!inputValue.trim() && attachments.length === 0}
-                className="px-2 py-2 mt-1 mr-2 hover:bg-(--background-darkest) text-(--sidebar-accent-fg) rounded-lg disabled:opacity-50"
+                className="px-2 py-2 mt-1 mr-1 hover:bg-(--background-darkest) text-(--sidebar-accent-fg) rounded-lg disabled:opacity-50"
                 title="Send message"
               >
-                <SendIcon size={20} />
+                <SendHorizontalIcon size={20} />
               </button>
             )}
           </div>
-          <div className="pl-2 pr-1 flex items-center justify-between">
-            <ChatInputControls showContextFilesPicker={true} />
-            <button
-              onClick={() => setShowTokenBar(!showTokenBar)}
-              className="flex items-center px-2 py-1 text-xs text-muted-foreground hover:bg-muted rounded"
-              title={showTokenBar ? "Hide token usage" : "Show token usage"}
-            >
-              <BarChart2 size={14} className="mr-1" />
-              {showTokenBar ? "Hide tokens" : "Tokens"}
-            </button>
+          <div className="pl-2 pr-1 flex items-center justify-between pb-2">
+            <div className="flex items-center">
+              <ChatInputControls showContextFilesPicker={true} />
+              {/* File attachment button */}
+              <TooltipProvider>
+                <Tooltip>
+                  <TooltipTrigger asChild>
+                    <Button
+                      variant="ghost"
+                      onClick={handleAttachmentClick}
+                      disabled={isStreaming}
+                      title="Attach files"
+                      size="sm"
+                    >
+                      <Paperclip size={20} />
+                    </Button>
+                  </TooltipTrigger>
+                  <TooltipContent>Attach files</TooltipContent>
+                </Tooltip>
+              </TooltipProvider>
+              <input
+                type="file"
+                ref={fileInputRef}
+                onChange={handleFileChange}
+                className="hidden"
+                multiple
+                accept=".jpg,.jpeg,.png,.gif,.webp,.txt,.md,.js,.ts,.html,.css,.json,.csv"
+              />
+            </div>
+
+            <TooltipProvider>
+              <Tooltip>
+                <TooltipTrigger asChild>
+                  <Button
+                    onClick={() => setShowTokenBar(!showTokenBar)}
+                    variant="ghost"
+                    className={`has-[>svg]:px-2 ${
+                      showTokenBar ? "text-purple-500 bg-purple-100" : ""
+                    }`}
+                    size="sm"
+                  >
+                    <ChartColumnIncreasing size={14} />
+                  </Button>
+                </TooltipTrigger>
+                <TooltipContent>
+                  {showTokenBar ? "Hide token usage" : "Show token usage"}
+                </TooltipContent>
+              </Tooltip>
+            </TooltipProvider>
           </div>
           {/* TokenBar is only displayed when showTokenBar is true */}
           {showTokenBar && <TokenBar chatId={chatId} />}