🤖 feat: add cloud icon toggle for Mux Gateway models

ammar-agent · ammar-agent · commit c5ff29fb69e5 · 2025-12-06T15:38:45.000-06:00
- Add whitespace-nowrap to ProviderWithIcon to prevent line breaks
- Create useGatewayModels hook for managing gateway preferences
- Add cloud icon toggle to ModelRow in settings
- Add cloud icon toggle to ModelSelector dropdown
- Transform model to gateway format when sending messages

Users can now click the cloud icon on any model to route it through
Mux Gateway. The preference is stored per-model and the transformation
happens transparently when sending messages.

_Generated with mux_
diff --git a/src/browser/components/ModelSelector.tsx b/src/browser/components/ModelSelector.tsx
@@ -7,9 +7,10 @@ import React, {
   forwardRef,
 } from "react";
 import { cn } from "@/common/lib/utils";
-import { Settings, Star } from "lucide-react";
+import { Cloud, Settings, Star } from "lucide-react";
 import { TooltipWrapper, Tooltip } from "./Tooltip";
 import { useSettings } from "@/browser/contexts/SettingsContext";
+import { useGatewayModels } from "@/browser/hooks/useGatewayModels";
 
 interface ModelSelectorProps {
   value: string;
@@ -27,6 +28,7 @@ export interface ModelSelectorRef {
 export const ModelSelector = forwardRef<ModelSelectorRef, ModelSelectorProps>(
   ({ value, onChange, recentModels, onComplete, defaultModel, onSetDefaultModel }, ref) => {
     const { open: openSettings } = useSettings();
+    const { isEnabled: isGatewayEnabled, toggle: toggleGateway } = useGatewayModels();
     const [isEditing, setIsEditing] = useState(false);
     const [inputValue, setInputValue] = useState(value);
     const [error, setError] = useState<string | null>(null);
@@ -190,8 +192,17 @@ export const ModelSelector = forwardRef<ModelSelectorRef, ModelSelectorProps>(
     }, [highlightedIndex]);
 
     if (!isEditing) {
+      const gatewayEnabled = isGatewayEnabled(value);
       return (
         <div ref={containerRef} className="relative flex items-center gap-1">
+          {gatewayEnabled && (
+            <TooltipWrapper inline>
+              <Cloud className="text-accent h-3 w-3 shrink-0 fill-current" />
+              <Tooltip className="tooltip" align="center">
+                Using Mux Gateway
+              </Tooltip>
+            </TooltipWrapper>
+          )}
           <div
             className="text-muted-light font-monospace dir-rtl hover:bg-hover max-w-36 cursor-pointer truncate rounded-sm px-1 py-0.5 text-left font-mono text-[10px] leading-[11px] transition-colors duration-200"
             onClick={handleClick}
@@ -250,36 +261,67 @@ export const ModelSelector = forwardRef<ModelSelectorRef, ModelSelectorProps>(
                   )}
                   onClick={() => handleSelectModel(model)}
                 >
-                  <div className="grid w-full grid-cols-[1fr_24px] items-center gap-2">
+                  <div className="grid w-full grid-cols-[1fr_auto] items-center gap-2">
                     <span className="min-w-0 truncate">{model}</span>
-                    {onSetDefaultModel && (
+                    <div className="flex items-center gap-0.5">
+                      {/* Gateway toggle */}
                       <TooltipWrapper inline>
                         <button
                           type="button"
                           onMouseDown={(e) => e.preventDefault()}
-                          onClick={(e) => handleSetDefault(e, model)}
+                          onClick={(e) => {
+                            e.preventDefault();
+                            e.stopPropagation();
+                            toggleGateway(model);
+                          }}
                           className={cn(
                             "flex items-center justify-center rounded-sm border px-1 py-0.5 transition-colors duration-150",
-                            defaultModel === model
-                              ? "text-yellow-400 border-yellow-400/40 cursor-default"
+                            isGatewayEnabled(model)
+                              ? "text-accent border-accent/40"
                               : "text-muted-light border-border-light/40 hover:border-foreground/60 hover:text-foreground"
                           )}
                           aria-label={
-                            defaultModel === model
-                              ? "Current default model"
-                              : "Set as default model"
+                            isGatewayEnabled(model) ? "Disable Mux Gateway" : "Enable Mux Gateway"
                           }
-                          disabled={defaultModel === model}
                         >
-                          <Star className="h-3 w-3" />
+                          <Cloud
+                            className={cn("h-3 w-3", isGatewayEnabled(model) && "fill-current")}
+                          />
                         </button>
                         <Tooltip className="tooltip" align="center">
-                          {defaultModel === model
-                            ? "Current default model"
-                            : "Set as default model"}
+                          {isGatewayEnabled(model) ? "Using Mux Gateway" : "Use Mux Gateway"}
                         </Tooltip>
                       </TooltipWrapper>
-                    )}
+                      {/* Default model toggle */}
+                      {onSetDefaultModel && (
+                        <TooltipWrapper inline>
+                          <button
+                            type="button"
+                            onMouseDown={(e) => e.preventDefault()}
+                            onClick={(e) => handleSetDefault(e, model)}
+                            className={cn(
+                              "flex items-center justify-center rounded-sm border px-1 py-0.5 transition-colors duration-150",
+                              defaultModel === model
+                                ? "text-yellow-400 border-yellow-400/40 cursor-default"
+                                : "text-muted-light border-border-light/40 hover:border-foreground/60 hover:text-foreground"
+                            )}
+                            aria-label={
+                              defaultModel === model
+                                ? "Current default model"
+                                : "Set as default model"
+                            }
+                            disabled={defaultModel === model}
+                          >
+                            <Star className="h-3 w-3" />
+                          </button>
+                          <Tooltip className="tooltip" align="center">
+                            {defaultModel === model
+                              ? "Current default model"
+                              : "Set as default model"}
+                          </Tooltip>
+                        </TooltipWrapper>
+                      )}
+                    </div>
                   </div>
                 </div>
               ))
diff --git a/src/browser/components/ProviderIcon.tsx b/src/browser/components/ProviderIcon.tsx
@@ -67,7 +67,7 @@ export function ProviderWithIcon(props: ProviderWithIconProps) {
     : props.provider;
 
   return (
-    <span className={cn("inline-flex items-center gap-1", props.className)}>
+    <span className={cn("inline-flex items-center gap-1 whitespace-nowrap", props.className)}>
       <ProviderIcon provider={props.provider} className={props.iconClassName} />
       <span>{name}</span>
     </span>
diff --git a/src/browser/components/Settings/sections/ModelRow.tsx b/src/browser/components/Settings/sections/ModelRow.tsx
@@ -1,5 +1,5 @@
 import React from "react";
-import { Check, Pencil, Star, Trash2, X } from "lucide-react";
+import { Check, Cloud, Pencil, Star, Trash2, X } from "lucide-react";
 import { cn } from "@/common/lib/utils";
 import { TooltipWrapper, Tooltip } from "@/browser/components/Tooltip";
 import { ProviderWithIcon } from "@/browser/components/ProviderIcon";
@@ -16,12 +16,16 @@ export interface ModelRowProps {
   editError?: string | null;
   saving?: boolean;
   hasActiveEdit?: boolean;
+  /** Whether gateway mode is enabled for this model */
+  isGatewayEnabled?: boolean;
   onSetDefault: () => void;
   onStartEdit?: () => void;
   onSaveEdit?: () => void;
   onCancelEdit?: () => void;
   onEditChange?: (value: string) => void;
   onRemove?: () => void;
+  /** Toggle gateway mode for this model */
+  onToggleGateway?: () => void;
 }
 
 export function ModelRow(props: ModelRowProps) {
@@ -90,6 +94,25 @@ export function ModelRow(props: ModelRowProps) {
           </>
         ) : (
           <>
+            {/* Gateway toggle button */}
+            {props.onToggleGateway && (
+              <TooltipWrapper inline>
+                <button
+                  type="button"
+                  onClick={props.onToggleGateway}
+                  className={cn(
+                    "p-0.5 transition-colors",
+                    props.isGatewayEnabled ? "text-accent" : "text-muted hover:text-accent"
+                  )}
+                  aria-label={props.isGatewayEnabled ? "Disable Mux Gateway" : "Enable Mux Gateway"}
+                >
+                  <Cloud className={cn("h-3.5 w-3.5", props.isGatewayEnabled && "fill-current")} />
+                </button>
+                <Tooltip className="tooltip" align="center">
+                  {props.isGatewayEnabled ? "Using Mux Gateway" : "Use Mux Gateway"}
+                </Tooltip>
+              </TooltipWrapper>
+            )}
             {/* Favorite/default button */}
             <TooltipWrapper inline>
               <button
diff --git a/src/browser/components/Settings/sections/ModelsSection.tsx b/src/browser/components/Settings/sections/ModelsSection.tsx
@@ -3,6 +3,7 @@ import { Plus, Loader2 } from "lucide-react";
 import { SUPPORTED_PROVIDERS, PROVIDER_DISPLAY_NAMES } from "@/common/constants/providers";
 import { KNOWN_MODELS } from "@/common/constants/knownModels";
 import { useModelLRU } from "@/browser/hooks/useModelLRU";
+import { useGatewayModels } from "@/browser/hooks/useGatewayModels";
 import { ModelRow } from "./ModelRow";
 import { useAPI } from "@/browser/contexts/API";
 import { useProvidersConfig } from "@/browser/hooks/useProvidersConfig";
@@ -25,6 +26,7 @@ export function ModelsSection() {
   const [editing, setEditing] = useState<EditingState | null>(null);
   const [error, setError] = useState<string | null>(null);
   const { defaultModel, setDefaultModel } = useModelLRU();
+  const { isEnabled: isGatewayEnabled, toggle: toggleGateway } = useGatewayModels();
 
   // Check if a model already exists (for duplicate prevention)
   const modelExists = useCallback(
@@ -213,6 +215,7 @@ export function ModelsSection() {
               editError={isModelEditing ? error : undefined}
               saving={false}
               hasActiveEdit={editing !== null}
+              isGatewayEnabled={isGatewayEnabled(model.fullId)}
               onSetDefault={() => setDefaultModel(model.fullId)}
               onStartEdit={() => handleStartEdit(model.provider, model.modelId)}
               onSaveEdit={handleSaveEdit}
@@ -221,6 +224,7 @@ export function ModelsSection() {
                 setEditing((prev) => (prev ? { ...prev, newModelId: value } : null))
               }
               onRemove={() => handleRemoveModel(model.provider, model.modelId)}
+              onToggleGateway={() => toggleGateway(model.fullId)}
             />
           );
         })}
@@ -241,7 +245,9 @@ export function ModelsSection() {
             isCustom={false}
             isDefault={defaultModel === model.fullId}
             isEditing={false}
+            isGatewayEnabled={isGatewayEnabled(model.fullId)}
             onSetDefault={() => setDefaultModel(model.fullId)}
+            onToggleGateway={() => toggleGateway(model.fullId)}
           />
         ))}
       </div>
diff --git a/src/browser/hooks/useGatewayModels.ts b/src/browser/hooks/useGatewayModels.ts
@@ -0,0 +1,74 @@
+import { useCallback } from "react";
+import { usePersistedState, readPersistedState, updatePersistedState } from "./usePersistedState";
+
+const GATEWAY_MODELS_KEY = "gateway-models";
+
+/**
+ * Check if a model is gateway-enabled (static read, no reactivity)
+ */
+export function isGatewayEnabled(modelId: string): boolean {
+  const gatewayModels = readPersistedState<string[]>(GATEWAY_MODELS_KEY, []);
+  return gatewayModels.includes(modelId);
+}
+
+/**
+ * Transform a model ID to gateway format if gateway is enabled for it.
+ * Example: "anthropic:claude-opus-4-5" → "mux-gateway:anthropic/claude-opus-4-5"
+ */
+export function toGatewayModel(modelId: string): string {
+  if (!isGatewayEnabled(modelId)) {
+    return modelId;
+  }
+  // Transform provider:model to mux-gateway:provider/model
+  const colonIndex = modelId.indexOf(":");
+  if (colonIndex === -1) {
+    return modelId;
+  }
+  const provider = modelId.slice(0, colonIndex);
+  const model = modelId.slice(colonIndex + 1);
+  return `mux-gateway:${provider}/${model}`;
+}
+
+/**
+ * Toggle gateway mode for a model (static update, no reactivity)
+ */
+export function toggleGatewayModel(modelId: string): void {
+  const gatewayModels = readPersistedState<string[]>(GATEWAY_MODELS_KEY, []);
+  if (gatewayModels.includes(modelId)) {
+    updatePersistedState(
+      GATEWAY_MODELS_KEY,
+      gatewayModels.filter((m) => m !== modelId)
+    );
+  } else {
+    updatePersistedState(GATEWAY_MODELS_KEY, [...gatewayModels, modelId]);
+  }
+}
+
+/**
+ * Hook to manage which models use the Mux Gateway.
+ * Returns reactive state and toggle function.
+ */
+export function useGatewayModels() {
+  const [gatewayModels, setGatewayModels] = usePersistedState<string[]>(GATEWAY_MODELS_KEY, [], {
+    listener: true,
+  });
+
+  const isEnabled = useCallback(
+    (modelId: string) => gatewayModels.includes(modelId),
+    [gatewayModels]
+  );
+
+  const toggle = useCallback(
+    (modelId: string) => {
+      setGatewayModels((prev) => {
+        if (prev.includes(modelId)) {
+          return prev.filter((m) => m !== modelId);
+        }
+        return [...prev, modelId];
+      });
+    },
+    [setGatewayModels]
+  );
+
+  return { gatewayModels, isEnabled, toggle };
+}
diff --git a/src/browser/hooks/useSendMessageOptions.ts b/src/browser/hooks/useSendMessageOptions.ts
@@ -2,6 +2,7 @@ import { useThinkingLevel } from "./useThinkingLevel";
 import { useMode } from "@/browser/contexts/ModeContext";
 import { usePersistedState } from "./usePersistedState";
 import { getDefaultModel } from "./useModelLRU";
+import { toGatewayModel } from "./useGatewayModels";
 import { modeToToolPolicy, PLAN_MODE_INSTRUCTION } from "@/common/utils/ui/modeUtils";
 import { getModelKey } from "@/common/constants/storage";
 import type { SendMessageOptions } from "@/common/orpc/types";
@@ -26,9 +27,12 @@ function constructSendMessageOptions(
   const additionalSystemInstructions = mode === "plan" ? PLAN_MODE_INSTRUCTION : undefined;
 
   // Ensure model is always a valid string (defensive against corrupted localStorage)
-  const model =
+  const baseModel =
     typeof preferredModel === "string" && preferredModel ? preferredModel : fallbackModel;
 
+  // Transform to gateway format if gateway is enabled for this model
+  const model = toGatewayModel(baseModel);
+
   // Enforce thinking policy at the UI boundary as well (e.g., gpt-5-pro → high only)
   const uiThinking = enforceThinkingPolicy(model, thinkingLevel);
 
diff --git a/src/browser/utils/messages/sendOptions.ts b/src/browser/utils/messages/sendOptions.ts
@@ -2,6 +2,7 @@ import { getModelKey, getThinkingLevelKey, getModeKey } from "@/common/constants
 import { modeToToolPolicy, PLAN_MODE_INSTRUCTION } from "@/common/utils/ui/modeUtils";
 import { readPersistedState } from "@/browser/hooks/usePersistedState";
 import { getDefaultModel } from "@/browser/hooks/useModelLRU";
+import { toGatewayModel } from "@/browser/hooks/useGatewayModels";
 import type { SendMessageOptions } from "@/common/orpc/types";
 import type { UIMode } from "@/common/types/mode";
 import type { ThinkingLevel } from "@/common/types/thinking";
@@ -38,7 +39,9 @@ function getProviderOptions(): MuxProviderOptions {
  */
 export function getSendOptionsFromStorage(workspaceId: string): SendMessageOptions {
   // Read model preference (workspace-specific), fallback to LRU default
-  const model = readPersistedState<string>(getModelKey(workspaceId), getDefaultModel());
+  const baseModel = readPersistedState<string>(getModelKey(workspaceId), getDefaultModel());
+  // Transform to gateway format if gateway is enabled for this model
+  const model = toGatewayModel(baseModel);
 
   // Read thinking level (workspace-specific)
   const thinkingLevel = readPersistedState<ThinkingLevel>(
