From ce983ade4370f0ba682793bcdf97bbbfa9aa6293 Mon Sep 17 00:00:00 2001
From: Thomas Kosiewski <tk@coder.com>
Date: Fri, 12 Dec 2025 11:36:52 +0100
Subject: [PATCH] =?UTF-8?q?=F0=9F=A4=96=20fix:=20gateway=20toggle=20double?=
 =?UTF-8?q?-toggle=20race=20condition?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

- Fix GatewayIcon not forwarding refs (React warning)
- Make gateway icon clickable on main page to toggle gateway
- Remove redundant updatePersistedState calls that caused double-toggle

The bug: setEnabledModels() from usePersistedState already writes to
localStorage synchronously, but updatePersistedState() was also called,
reading the already-toggled value and toggling again, reverting changes.
---
 src/browser/components/ModelSelector.tsx     | 26 ++++++++++++++------
 src/browser/components/icons/GatewayIcon.tsx |  9 ++++---
 src/browser/hooks/useGatewayModels.ts        | 11 +++------
 3 files changed, 27 insertions(+), 19 deletions(-)

diff --git a/src/browser/components/ModelSelector.tsx b/src/browser/components/ModelSelector.tsx
index 1a56038d62..432b46d388 100644
--- a/src/browser/components/ModelSelector.tsx
+++ b/src/browser/components/ModelSelector.tsx
@@ -205,17 +205,27 @@ export const ModelSelector = forwardRef<ModelSelectorRef, ModelSelectorProps>(
 
       return (
         <div ref={containerRef} className="relative flex items-center gap-1">
-          {gatewayActive && (
+          {gateway.canToggleModel(value) && (
             <Tooltip>
               <TooltipTrigger asChild>
-                <GatewayIcon
-                  className="text-accent h-3 w-3 shrink-0"
-                  active
-                  role="img"
-                  aria-label="Using Mux Gateway"
-                />
+                <button
+                  type="button"
+                  onClick={() => gateway.toggleModelGateway(value)}
+                  className="flex items-center justify-center"
+                  aria-label={gatewayActive ? "Disable Mux Gateway" : "Enable Mux Gateway"}
+                >
+                  <GatewayIcon
+                    className={cn(
+                      "h-3 w-3 shrink-0 transition-colors",
+                      gatewayActive ? "text-accent" : "text-muted-light hover:text-accent"
+                    )}
+                    active={gatewayActive}
+                  />
+                </button>
               </TooltipTrigger>
-              <TooltipContent align="center">Using Mux Gateway</TooltipContent>
+              <TooltipContent align="center">
+                {gatewayActive ? "Using Mux Gateway (click to disable)" : "Enable Mux Gateway"}
+              </TooltipContent>
             </Tooltip>
           )}
           <Tooltip>
diff --git a/src/browser/components/icons/GatewayIcon.tsx b/src/browser/components/icons/GatewayIcon.tsx
index 170970a6c8..0886b54c7a 100644
--- a/src/browser/components/icons/GatewayIcon.tsx
+++ b/src/browser/components/icons/GatewayIcon.tsx
@@ -1,4 +1,4 @@
-import React from "react";
+import React, { forwardRef } from "react";
 
 interface GatewayIconProps extends React.SVGProps<SVGSVGElement> {
   className?: string;
@@ -10,11 +10,12 @@ interface GatewayIconProps extends React.SVGProps<SVGSVGElement> {
  * Gateway icon - represents routing through Mux Gateway.
  * Circle with M logo. Active state adds outer ring.
  */
-export function GatewayIcon(props: GatewayIconProps) {
+export const GatewayIcon = forwardRef<SVGSVGElement, GatewayIconProps>((props, ref) => {
   const { active, ...svgProps } = props;
 
   return (
     <svg
+      ref={ref}
       xmlns="http://www.w3.org/2000/svg"
       viewBox="0 0 24 24"
       fill="none"
@@ -32,4 +33,6 @@ export function GatewayIcon(props: GatewayIconProps) {
       <path d="M8 16V8l4 5 4-5v8" />
     </svg>
   );
-}
+});
+
+GatewayIcon.displayName = "GatewayIcon";
diff --git a/src/browser/hooks/useGatewayModels.ts b/src/browser/hooks/useGatewayModels.ts
index 172379de6a..6953c38998 100644
--- a/src/browser/hooks/useGatewayModels.ts
+++ b/src/browser/hooks/useGatewayModels.ts
@@ -165,9 +165,8 @@ export function useGateway(): GatewayState {
   const isActive = isConfigured && isEnabled;
 
   const toggleEnabled = useCallback(() => {
+    // setIsEnabled (from usePersistedState) already writes to localStorage synchronously
     setIsEnabled((prev) => !prev);
-    // Also update localStorage synchronously
-    updatePersistedState<boolean>(GATEWAY_ENABLED_KEY, (prev) => !prev);
   }, [setIsEnabled]);
 
   const modelUsesGateway = useCallback(
@@ -177,15 +176,11 @@ export function useGateway(): GatewayState {
 
   const toggleModelGateway = useCallback(
     (modelId: string) => {
-      // Update React state for UI
+      // setEnabledModels (from usePersistedState) already writes to localStorage synchronously,
+      // so we don't need to call updatePersistedState separately.
       setEnabledModels((prev) =>
         prev.includes(modelId) ? prev.filter((m) => m !== modelId) : [...prev, modelId]
       );
-      // Also update localStorage synchronously so toGatewayModel() sees it immediately
-      // (usePersistedState batches writes in microtask, which can cause race conditions)
-      updatePersistedState<string[]>(GATEWAY_MODELS_KEY, (prev) =>
-        prev.includes(modelId) ? prev.filter((m) => m !== modelId) : [...prev, modelId]
-      );
     },
     [setEnabledModels]
   );