fix(web): refactor pico chat flow and fix proxied websocket URLs (#1639)

- move chat controller, state, protocol, history, and websocket logic into a dedicated chat feature module
- improve chat reconnection, session hydration, and send gating based on actual websocket state
- preserve gateway status during transient SSE disconnects and update stop state immediately
- generate wss websocket URLs behind HTTPS proxies and add backend tests for forwarded proto handling

web/backend/api/gateway_host.go

@@ -57,10 +57,28 @@ func requestHostName(r *http.Request) string {
 	return "127.0.0.1"
 }
 
+func requestWSScheme(r *http.Request) string {
+	if forwarded := strings.TrimSpace(r.Header.Get("X-Forwarded-Proto")); forwarded != "" {
+		proto := strings.ToLower(strings.TrimSpace(strings.Split(forwarded, ",")[0]))
+		if proto == "https" || proto == "wss" {
+			return "wss"
+		}
+		if proto == "http" || proto == "ws" {
+			return "ws"
+		}
+	}
+
+	if r.TLS != nil {
+		return "wss"
+	}
+
+	return "ws"
+}
+
 func (h *Handler) buildWsURL(r *http.Request, cfg *config.Config) string {
 	host := h.effectiveGatewayBindHost(cfg)
 	if host == "" || host == "0.0.0.0" {
 		host = requestHostName(r)
 	}
-	return "ws://" + net.JoinHostPort(host, strconv.Itoa(cfg.Gateway.Port)) + "/pico/ws"
+	return requestWSScheme(r) + "://" + net.JoinHostPort(host, strconv.Itoa(cfg.Gateway.Port)) + "/pico/ws"
 }

web/backend/api/gateway_host_test.go

@@ -1,6 +1,7 @@
 package api
 
 import (
+	"crypto/tls"
 	"net/http/httptest"
 	"path/filepath"
 	"testing"
@@ -57,3 +58,55 @@ func TestGatewayProbeHostUsesLoopbackForWildcardBind(t *testing.T) {
 		t.Fatalf("gatewayProbeHost() = %q, want %q", got, "127.0.0.1")
 	}
 }
+
+func TestBuildWsURLUsesWSSWhenForwardedProtoIsHTTPS(t *testing.T) {
+	configPath := filepath.Join(t.TempDir(), "config.json")
+	h := NewHandler(configPath)
+
+	cfg := config.DefaultConfig()
+	cfg.Gateway.Host = "0.0.0.0"
+	cfg.Gateway.Port = 18790
+
+	req := httptest.NewRequest("GET", "http://launcher.local/api/pico/token", nil)
+	req.Host = "chat.example.com"
+	req.Header.Set("X-Forwarded-Proto", "https")
+
+	if got := h.buildWsURL(req, cfg); got != "wss://chat.example.com:18790/pico/ws" {
+		t.Fatalf("buildWsURL() = %q, want %q", got, "wss://chat.example.com:18790/pico/ws")
+	}
+}
+
+func TestBuildWsURLUsesWSSWhenRequestIsTLS(t *testing.T) {
+	configPath := filepath.Join(t.TempDir(), "config.json")
+	h := NewHandler(configPath)
+
+	cfg := config.DefaultConfig()
+	cfg.Gateway.Host = "0.0.0.0"
+	cfg.Gateway.Port = 18790
+
+	req := httptest.NewRequest("GET", "https://launcher.local/api/pico/token", nil)
+	req.Host = "secure.example.com"
+	req.TLS = &tls.ConnectionState{}
+
+	if got := h.buildWsURL(req, cfg); got != "wss://secure.example.com:18790/pico/ws" {
+		t.Fatalf("buildWsURL() = %q, want %q", got, "wss://secure.example.com:18790/pico/ws")
+	}
+}
+
+func TestBuildWsURLPrefersForwardedHTTPOverTLS(t *testing.T) {
+	configPath := filepath.Join(t.TempDir(), "config.json")
+	h := NewHandler(configPath)
+
+	cfg := config.DefaultConfig()
+	cfg.Gateway.Host = "0.0.0.0"
+	cfg.Gateway.Port = 18790
+
+	req := httptest.NewRequest("GET", "https://launcher.local/api/pico/token", nil)
+	req.Host = "chat.example.com"
+	req.TLS = &tls.ConnectionState{}
+	req.Header.Set("X-Forwarded-Proto", "http")
+
+	if got := h.buildWsURL(req, cfg); got != "ws://chat.example.com:18790/pico/ws" {
+		t.Fatalf("buildWsURL() = %q, want %q", got, "ws://chat.example.com:18790/pico/ws")
+	}
+}

web/frontend/src/components/chat/chat-composer.tsx

@@ -42,7 +42,7 @@ export function ChatComposer({
           placeholder={t("chat.placeholder")}
           disabled={!canInput}
           className={cn(
-            "max-h-[200px] min-h-[60px] resize-none border-0 bg-transparent px-2 py-1 text-[15px] shadow-none transition-colors focus-visible:ring-0 focus-visible:outline-none dark:bg-transparent",
+            "placeholder:text-muted-foreground max-h-[200px] min-h-[60px] resize-none border-0 bg-transparent px-2 py-1 text-[15px] shadow-none transition-colors focus-visible:ring-0 focus-visible:outline-none dark:bg-transparent",
             !canInput && "cursor-not-allowed",
           )}
           minRows={1}
@@ -56,7 +56,7 @@ export function ChatComposer({
             size="icon"
             className="size-8 rounded-full bg-violet-500 text-white transition-transform hover:bg-violet-600 active:scale-95"
             onClick={onSend}
-            disabled={!input.trim() || !isConnected}
+            disabled={!input.trim() || !canInput}
           >
             <IconArrowUp className="size-4" />
           </Button>

web/frontend/src/components/chat/chat-empty-state.tsx

@@ -34,7 +34,7 @@ export function ChatEmptyState({
         <p className="text-muted-foreground mb-4 text-center text-sm">
           {t("chat.empty.noConfiguredModelDescription")}
         </p>
-        <Button asChild variant="secondary" size="sm" className="px-4">
+        <Button asChild variant="outline" size="sm" className="px-4">
           <Link to="/models">{t("chat.empty.goToModels")}</Link>
         </Button>
       </div>

web/frontend/src/components/chat/chat-page.tsx

@@ -15,7 +15,6 @@ import { useChatModels } from "@/hooks/use-chat-models"
 import { useGateway } from "@/hooks/use-gateway"
 import { usePicoChat } from "@/hooks/use-pico-chat"
 import { useSessionHistory } from "@/hooks/use-session-history"
-import { hydrateActiveSession } from "@/lib/pico-chat-controller"
 
 export function ChatPage() {
   const { t } = useTranslation()
@@ -26,6 +25,7 @@ export function ChatPage() {
 
   const {
     messages,
+    connectionState,
     isTyping,
     activeSessionId,
     sendMessage,
@@ -34,7 +34,8 @@ export function ChatPage() {
   } = usePicoChat()
 
   const { state: gwState } = useGateway()
-  const isConnected = gwState === "running"
+  const isGatewayRunning = gwState === "running"
+  const isChatConnected = connectionState === "connected"
 
   const {
     defaultModelName,
@@ -43,7 +44,8 @@ export function ChatPage() {
     oauthModels,
     localModels,
     handleSetDefault,
-  } = useChatModels({ isConnected })
+  } = useChatModels({ isConnected: isGatewayRunning })
+  const canSend = isChatConnected && Boolean(defaultModelName)
 
   const {
     sessions,
@@ -68,10 +70,6 @@ export function ChatPage() {
     syncScrollState(e.currentTarget)
   }
 
-  useEffect(() => {
-    void hydrateActiveSession()
-  }, [])
-
   useEffect(() => {
     if (scrollRef.current) {
       if (isAtBottom) {
@@ -82,9 +80,10 @@ export function ChatPage() {
   }, [messages, isTyping, isAtBottom])
 
   const handleSend = () => {
-    if (!input.trim() || !isConnected) return
-    sendMessage(input.trim())
-    setInput("")
+    if (!input.trim() || !canSend) return
+    if (sendMessage(input.trim())) {
+      setInput("")
+    }
   }
 
   return (
@@ -143,7 +142,7 @@ export function ChatPage() {
             <ChatEmptyState
               hasConfiguredModels={hasConfiguredModels}
               defaultModelName={defaultModelName}
-              isConnected={isConnected}
+              isConnected={isGatewayRunning}
             />
           )}
 
@@ -168,7 +167,7 @@ export function ChatPage() {
         input={input}
         onInputChange={setInput}
         onSend={handleSend}
-        isConnected={isConnected}
+        isConnected={isChatConnected}
         hasDefaultModel={Boolean(defaultModelName)}
       />
     </div>

web/frontend/src/features/chat/controller.ts

@@ -2,24 +2,24 @@ import { getDefaultStore } from "jotai"
 import { toast } from "sonner"
 
 import { getPicoToken } from "@/api/pico"
-import { getSessionHistory } from "@/api/sessions"
-import i18n from "@/i18n"
+import {
+  loadSessionMessages,
+  mergeHistoryMessages,
+} from "@/features/chat/history"
+import { type PicoMessage, handlePicoMessage } from "@/features/chat/protocol"
 import {
   clearStoredSessionId,
   generateSessionId,
-  normalizeUnixTimestamp,
   readStoredSessionId,
-} from "@/lib/pico-chat-state"
-import { type ChatMessage, getChatState, updateChatStore } from "@/store/chat"
-import { gatewayAtom } from "@/store/gateway"
-
-interface PicoMessage {
-  type: string
-  id?: string
-  session_id?: string
-  timestamp?: number | string
-  payload?: Record<string, unknown>
-}
+} from "@/features/chat/state"
+import {
+  invalidateSocket,
+  isCurrentSocket,
+  normalizeWsUrlForBrowser,
+} from "@/features/chat/websocket"
+import i18n from "@/i18n"
+import { getChatState, updateChatStore } from "@/store/chat"
+import { type GatewayState, gatewayAtom } from "@/store/gateway"
 
 const store = getDefaultStore()
 
@@ -31,81 +31,51 @@ let initialized = false
 let unsubscribeGateway: (() => void) | null = null
 let hydratePromise: Promise<void> | null = null
 let connectionGeneration = 0
+let reconnectTimer: number | null = null
+let reconnectAttempts = 0
+let shouldMaintainConnection = false
 
-async function loadSessionMessages(sessionId: string): Promise<ChatMessage[]> {
-  const detail = await getSessionHistory(sessionId)
-  const fallbackTime = detail.updated
-
-  return detail.messages.map((message, index) => ({
-    id: `hist-${index}-${Date.now()}`,
-    role: message.role,
-    content: message.content,
-    timestamp: fallbackTime,
-  }))
+function clearReconnectTimer() {
+  if (reconnectTimer !== null) {
+    window.clearTimeout(reconnectTimer)
+    reconnectTimer = null
+  }
 }
 
-function handlePicoMessage(message: PicoMessage) {
-  const payload = message.payload || {}
+function shouldReconnectFor(generation: number, sessionId: string): boolean {
+  return (
+    shouldMaintainConnection &&
+    generation === connectionGeneration &&
+    sessionId === activeSessionIdRef &&
+    store.get(gatewayAtom).status === "running"
+  )
+}
 
-  switch (message.type) {
-    case "message.create": {
-      const content = (payload.content as string) || ""
-      const messageId = (payload.message_id as string) || `pico-${Date.now()}`
-      const timestamp =
-        message.timestamp !== undefined &&
-        Number.isFinite(Number(message.timestamp))
-          ? normalizeUnixTimestamp(Number(message.timestamp))
-          : Date.now()
-
-      updateChatStore((prev) => ({
-        messages: [
-          ...prev.messages,
-          {
-            id: messageId,
-            role: "assistant",
-            content,
-            timestamp,
-          },
-        ],
-        isTyping: false,
-      }))
-      break
-    }
-
-    case "message.update": {
-      const content = (payload.content as string) || ""
-      const messageId = payload.message_id as string
-      if (!messageId) {
-        break
-      }
-
-      updateChatStore((prev) => ({
-        messages: prev.messages.map((msg) =>
-          msg.id === messageId ? { ...msg, content } : msg,
-        ),
-      }))
-      break
-    }
-
-    case "typing.start":
-      updateChatStore({ isTyping: true })
-      break
-
-    case "typing.stop":
-      updateChatStore({ isTyping: false })
-      break
-
-    case "error":
-      console.error("Pico error:", payload)
-      updateChatStore({ isTyping: false })
-      break
-
-    case "pong":
-      break
-
-    default:
-      console.log("Unknown pico message type:", message.type)
+function scheduleReconnect(generation: number, sessionId: string) {
+  if (!shouldReconnectFor(generation, sessionId) || reconnectTimer !== null) {
+    return
   }
+
+  const delay = Math.min(1000 * 2 ** reconnectAttempts, 5000)
+  reconnectAttempts += 1
+  reconnectTimer = window.setTimeout(() => {
+    reconnectTimer = null
+    if (!shouldReconnectFor(generation, sessionId)) {
+      return
+    }
+    void connectChat()
+  }, delay)
+}
+
+function needsActiveSessionHydration(): boolean {
+  const state = getChatState()
+  const storedSessionId = readStoredSessionId()
+
+  return Boolean(
+    storedSessionId &&
+    storedSessionId === state.activeSessionId &&
+    !state.hasHydratedActiveSession,
+  )
 }
 
 function setActiveSessionId(sessionId: string) {
@@ -113,8 +83,35 @@ function setActiveSessionId(sessionId: string) {
   updateChatStore({ activeSessionId: sessionId })
 }
 
+function disconnectChatInternal({
+  clearDesiredConnection,
+}: {
+  clearDesiredConnection: boolean
+}) {
+  connectionGeneration += 1
+  clearReconnectTimer()
+
+  if (clearDesiredConnection) {
+    shouldMaintainConnection = false
+  }
+
+  const socket = wsRef
+  wsRef = null
+  isConnecting = false
+
+  invalidateSocket(socket)
+
+  updateChatStore({
+    connectionState: "disconnected",
+    isTyping: false,
+  })
+}
+
 export async function connectChat() {
-  if (store.get(gatewayAtom).status !== "running") {
+  if (
+    store.get(gatewayAtom).status !== "running" ||
+    needsActiveSessionHydration()
+  ) {
     return
   }
 
@@ -130,12 +127,15 @@ export async function connectChat() {
   const generation = connectionGeneration + 1
   connectionGeneration = generation
   isConnecting = true
+  clearReconnectTimer()
   updateChatStore({ connectionState: "connecting" })
 
   try {
     const { token, ws_url } = await getPicoToken()
+    const sessionId = activeSessionIdRef
 
     if (generation !== connectionGeneration) {
+      isConnecting = false
       return
     }
 
@@ -143,56 +143,71 @@ export async function connectChat() {
       console.error("No pico token available")
       updateChatStore({ connectionState: "error" })
       isConnecting = false
+      scheduleReconnect(generation, sessionId)
       return
     }
 
-    let finalWsUrl = ws_url
-    try {
-      const parsedUrl = new URL(ws_url)
-      const isLocalHost =
-        parsedUrl.hostname === "localhost" ||
-        parsedUrl.hostname === "127.0.0.1" ||
-        parsedUrl.hostname === "0.0.0.0"
-      const isBrowserLocal =
-        window.location.hostname === "localhost" ||
-        window.location.hostname === "127.0.0.1"
-
-      if (isLocalHost && !isBrowserLocal) {
-        parsedUrl.hostname = window.location.hostname
-        finalWsUrl = parsedUrl.toString()
-      }
-    } catch (error) {
-      console.warn("Could not parse ws_url:", error)
-    }
-
-    const url = `${finalWsUrl}?session_id=${encodeURIComponent(activeSessionIdRef)}`
-    // Send token as a subprotocol so it doesn't end up in the URL.
+    const finalWsUrl = normalizeWsUrlForBrowser(ws_url)
+    const url = `${finalWsUrl}?session_id=${encodeURIComponent(sessionId)}`
     const socket = new WebSocket(url, [`token.${token}`])
 
     if (generation !== connectionGeneration) {
-      socket.close()
+      isConnecting = false
+      invalidateSocket(socket)
       return
     }
 
     socket.onopen = () => {
-      if (wsRef !== socket) {
+      if (
+        !isCurrentSocket({
+          socket,
+          currentSocket: wsRef,
+          generation,
+          currentGeneration: connectionGeneration,
+          sessionId,
+          currentSessionId: activeSessionIdRef,
+        })
+      ) {
         return
       }
       updateChatStore({ connectionState: "connected" })
       isConnecting = false
+      reconnectAttempts = 0
     }
 
     socket.onmessage = (event) => {
+      if (
+        !isCurrentSocket({
+          socket,
+          currentSocket: wsRef,
+          generation,
+          currentGeneration: connectionGeneration,
+          sessionId,
+          currentSessionId: activeSessionIdRef,
+        })
+      ) {
+        return
+      }
+
       try {
-        const message: PicoMessage = JSON.parse(event.data)
-        handlePicoMessage(message)
+        const message = JSON.parse(event.data) as PicoMessage
+        handlePicoMessage(message, sessionId)
       } catch {
         console.warn("Non-JSON message from pico:", event.data)
       }
     }
 
     socket.onclose = () => {
-      if (wsRef !== socket) {
+      if (
+        !isCurrentSocket({
+          socket,
+          currentSocket: wsRef,
+          generation,
+          currentGeneration: connectionGeneration,
+          sessionId,
+          currentSessionId: activeSessionIdRef,
+        })
+      ) {
         return
       }
       wsRef = null
@@ -201,42 +216,42 @@ export async function connectChat() {
         connectionState: "disconnected",
         isTyping: false,
       })
+      scheduleReconnect(generation, sessionId)
     }
 
     socket.onerror = () => {
-      if (wsRef !== socket) {
+      if (
+        !isCurrentSocket({
+          socket,
+          currentSocket: wsRef,
+          generation,
+          currentGeneration: connectionGeneration,
+          sessionId,
+          currentSessionId: activeSessionIdRef,
+        })
+      ) {
         return
       }
       isConnecting = false
       updateChatStore({ connectionState: "error" })
+      scheduleReconnect(generation, sessionId)
     }
 
     wsRef = socket
   } catch (error) {
     if (generation !== connectionGeneration) {
+      isConnecting = false
       return
     }
     console.error("Failed to connect to pico:", error)
     updateChatStore({ connectionState: "error" })
     isConnecting = false
+    scheduleReconnect(generation, activeSessionIdRef)
   }
 }
 
 export function disconnectChat() {
-  connectionGeneration += 1
-
-  const socket = wsRef
-  wsRef = null
-  isConnecting = false
-
-  if (socket) {
-    socket.close()
-  }
-
-  updateChatStore({
-    connectionState: "disconnected",
-    isTyping: false,
-  })
+  disconnectChatInternal({ clearDesiredConnection: true })
 }
 
 export async function hydrateActiveSession() {
@@ -250,7 +265,6 @@ export async function hydrateActiveSession() {
   if (
     !storedSessionId ||
     state.hasHydratedActiveSession ||
-    state.messages.length > 0 ||
     storedSessionId !== state.activeSessionId
   ) {
     if (!state.hasHydratedActiveSession) {
@@ -267,7 +281,13 @@ export async function hydrateActiveSession() {
       }
 
       if (currentState.messages.length > 0) {
-        updateChatStore({ hasHydratedActiveSession: true })
+        updateChatStore({
+          messages: mergeHistoryMessages(
+            historyMessages,
+            currentState.messages,
+          ),
+          hasHydratedActiveSession: true,
+        })
         return
       }
 
@@ -307,9 +327,10 @@ export async function hydrateActiveSession() {
 export function sendChatMessage(content: string) {
   if (!wsRef || wsRef.readyState !== WebSocket.OPEN) {
     console.warn("WebSocket not connected")
-    return
+    return false
   }
 
+  const socket = wsRef
   const id = `msg-${++msgIdCounter}-${Date.now()}`
 
   updateChatStore((prev) => ({
@@ -320,13 +341,23 @@ export function sendChatMessage(content: string) {
     isTyping: true,
   }))
 
-  wsRef.send(
-    JSON.stringify({
-      type: "message.send",
-      id,
-      payload: { content },
-    }),
-  )
+  try {
+    socket.send(
+      JSON.stringify({
+        type: "message.send",
+        id,
+        payload: { content },
+      }),
+    )
+    return true
+  } catch (error) {
+    console.error("Failed to send pico message:", error)
+    updateChatStore((prev) => ({
+      messages: prev.messages.filter((message) => message.id !== id),
+      isTyping: false,
+    }))
+    return false
+  }
 }
 
 export async function switchChatSession(sessionId: string) {
@@ -337,7 +368,7 @@ export async function switchChatSession(sessionId: string) {
   try {
     const historyMessages = await loadSessionMessages(sessionId)
 
-    disconnectChat()
+    disconnectChatInternal({ clearDesiredConnection: false })
     setActiveSessionId(sessionId)
     updateChatStore({
       messages: historyMessages,
@@ -346,6 +377,7 @@ export async function switchChatSession(sessionId: string) {
     })
 
     if (store.get(gatewayAtom).status === "running") {
+      shouldMaintainConnection = true
       await connectChat()
     }
   } catch (error) {
@@ -359,7 +391,7 @@ export async function newChatSession() {
     return
   }
 
-  disconnectChat()
+  disconnectChatInternal({ clearDesiredConnection: false })
   setActiveSessionId(generateSessionId())
   updateChatStore({
     messages: [],
@@ -368,6 +400,7 @@ export async function newChatSession() {
   })
 
   if (store.get(gatewayAtom).status === "running") {
+    shouldMaintainConnection = true
     await connectChat()
   }
 }
@@ -379,23 +412,43 @@ export function initializeChatStore() {
 
   initialized = true
   activeSessionIdRef = getChatState().activeSessionId
+  let lastGatewayStatus: GatewayState | null = null
 
-  const syncConnectionWithGateway = () => {
-    if (store.get(gatewayAtom).status === "running") {
+  const syncConnectionWithGateway = (force: boolean = false) => {
+    const gatewayStatus = store.get(gatewayAtom).status
+    if (!force && gatewayStatus === lastGatewayStatus) {
+      return
+    }
+    lastGatewayStatus = gatewayStatus
+
+    if (gatewayStatus === "running") {
+      shouldMaintainConnection = true
+      if (needsActiveSessionHydration()) {
+        return
+      }
       void connectChat()
       return
     }
 
-    disconnectChat()
+    if (gatewayStatus === "stopped" || gatewayStatus === "error") {
+      disconnectChatInternal({ clearDesiredConnection: true })
+    }
   }
 
   unsubscribeGateway = store.sub(gatewayAtom, syncConnectionWithGateway)
 
   if (!readStoredSessionId()) {
     updateChatStore({ hasHydratedActiveSession: true })
+    syncConnectionWithGateway(true)
+    return
   }
 
-  syncConnectionWithGateway()
+  void hydrateActiveSession().finally(() => {
+    if (!initialized) {
+      return
+    }
+    syncConnectionWithGateway(true)
+  })
 }
 
 export function teardownChatStore() {

web/frontend/src/features/chat/history.ts

@@ -0,0 +1,68 @@
+import { getSessionHistory } from "@/api/sessions"
+import { normalizeUnixTimestamp } from "@/features/chat/state"
+import type { ChatMessage } from "@/store/chat"
+
+export async function loadSessionMessages(
+  sessionId: string,
+): Promise<ChatMessage[]> {
+  const detail = await getSessionHistory(sessionId)
+  const fallbackTime = detail.updated
+
+  return detail.messages.map((message, index) => ({
+    id: `hist-${index}-${Date.now()}`,
+    role: message.role,
+    content: message.content,
+    timestamp: fallbackTime,
+  }))
+}
+
+function normalizeMessageTimestamp(timestamp: number | string): string {
+  if (typeof timestamp === "number") {
+    return String(normalizeUnixTimestamp(timestamp))
+  }
+
+  const trimmed = timestamp.trim()
+  if (/^-?\d+(\.\d+)?$/.test(trimmed)) {
+    return String(normalizeUnixTimestamp(Number(trimmed)))
+  }
+
+  const parsed = Date.parse(trimmed)
+  return Number.isNaN(parsed) ? trimmed : String(parsed)
+}
+
+function messageSignature(message: ChatMessage): string {
+  return `${message.role}\u0000${message.content}\u0000${normalizeMessageTimestamp(
+    message.timestamp,
+  )}`
+}
+
+function comparableTimestamp(timestamp: number | string): number {
+  const normalized = normalizeMessageTimestamp(timestamp)
+  const numeric = Number(normalized)
+  return Number.isFinite(numeric) ? numeric : 0
+}
+
+export function mergeHistoryMessages(
+  historyMessages: ChatMessage[],
+  currentMessages: ChatMessage[],
+): ChatMessage[] {
+  const currentIds = new Set(currentMessages.map((message) => message.id))
+  const currentSignatures = new Set(
+    currentMessages.map((message) => messageSignature(message)),
+  )
+
+  const merged = [
+    ...historyMessages.filter(
+      (message) =>
+        !currentIds.has(message.id) &&
+        !currentSignatures.has(messageSignature(message)),
+    ),
+    ...currentMessages,
+  ]
+
+  return merged.sort(
+    (left, right) =>
+      comparableTimestamp(left.timestamp) -
+      comparableTimestamp(right.timestamp),
+  )
+}

web/frontend/src/features/chat/protocol.ts

@@ -0,0 +1,81 @@
+import { normalizeUnixTimestamp } from "@/features/chat/state"
+import { updateChatStore } from "@/store/chat"
+
+export interface PicoMessage {
+  type: string
+  id?: string
+  session_id?: string
+  timestamp?: number | string
+  payload?: Record<string, unknown>
+}
+
+export function handlePicoMessage(
+  message: PicoMessage,
+  expectedSessionId: string,
+) {
+  if (message.session_id && message.session_id !== expectedSessionId) {
+    return
+  }
+
+  const payload = message.payload || {}
+
+  switch (message.type) {
+    case "message.create": {
+      const content = (payload.content as string) || ""
+      const messageId = (payload.message_id as string) || `pico-${Date.now()}`
+      const timestamp =
+        message.timestamp !== undefined &&
+        Number.isFinite(Number(message.timestamp))
+          ? normalizeUnixTimestamp(Number(message.timestamp))
+          : Date.now()
+
+      updateChatStore((prev) => ({
+        messages: [
+          ...prev.messages,
+          {
+            id: messageId,
+            role: "assistant",
+            content,
+            timestamp,
+          },
+        ],
+        isTyping: false,
+      }))
+      break
+    }
+
+    case "message.update": {
+      const content = (payload.content as string) || ""
+      const messageId = payload.message_id as string
+      if (!messageId) {
+        break
+      }
+
+      updateChatStore((prev) => ({
+        messages: prev.messages.map((msg) =>
+          msg.id === messageId ? { ...msg, content } : msg,
+        ),
+      }))
+      break
+    }
+
+    case "typing.start":
+      updateChatStore({ isTyping: true })
+      break
+
+    case "typing.stop":
+      updateChatStore({ isTyping: false })
+      break
+
+    case "error":
+      console.error("Pico error:", payload)
+      updateChatStore({ isTyping: false })
+      break
+
+    case "pong":
+      break
+
+    default:
+      console.log("Unknown pico message type:", message.type)
+  }
+}

web/frontend/src/features/chat/websocket.ts

@@ -0,0 +1,57 @@
+export function normalizeWsUrlForBrowser(wsUrl: string): string {
+  let finalWsUrl = wsUrl
+
+  try {
+    const parsedUrl = new URL(wsUrl)
+    const isLocalHost =
+      parsedUrl.hostname === "localhost" ||
+      parsedUrl.hostname === "127.0.0.1" ||
+      parsedUrl.hostname === "0.0.0.0"
+    const isBrowserLocal =
+      window.location.hostname === "localhost" ||
+      window.location.hostname === "127.0.0.1"
+
+    if (isLocalHost && !isBrowserLocal) {
+      parsedUrl.hostname = window.location.hostname
+      finalWsUrl = parsedUrl.toString()
+    }
+  } catch (error) {
+    console.warn("Could not parse ws_url:", error)
+  }
+
+  return finalWsUrl
+}
+
+export function invalidateSocket(socket: WebSocket | null) {
+  if (!socket) {
+    return
+  }
+
+  socket.onopen = null
+  socket.onmessage = null
+  socket.onclose = null
+  socket.onerror = null
+  socket.close()
+}
+
+export function isCurrentSocket({
+  socket,
+  currentSocket,
+  generation,
+  currentGeneration,
+  sessionId,
+  currentSessionId,
+}: {
+  socket: WebSocket
+  currentSocket: WebSocket | null
+  generation: number
+  currentGeneration: number
+  sessionId: string
+  currentSessionId: string
+}): boolean {
+  return (
+    currentSocket === socket &&
+    generation === currentGeneration &&
+    sessionId === currentSessionId
+  )
+}

web/frontend/src/hooks/use-gateway.ts

@@ -67,10 +67,9 @@ export function useGateway() {
     }
 
     es.onerror = () => {
-      // EventSource will auto-reconnect
-      updateGatewayStore((prev) =>
-        prev.status === "restarting" ? {} : { status: "unknown" },
-      )
+      // EventSource will auto-reconnect. Preserve the last known gateway
+      // status so transient SSE disconnects do not suppress chat websocket
+      // reconnects while polling catches up.
     }
 
     return () => {
@@ -105,6 +104,11 @@ export function useGateway() {
     setLoading(true)
     try {
       await stopGateway()
+      updateGatewayStore({
+        status: "stopped",
+        canStart: true,
+        restartRequired: false,
+      })
     } catch (err) {
       console.error("Failed to stop gateway:", err)
     } finally {

web/frontend/src/hooks/use-pico-chat.ts

@@ -5,7 +5,7 @@ import {
   newChatSession,
   sendChatMessage,
   switchChatSession,
-} from "@/lib/pico-chat-controller"
+} from "@/features/chat/controller"
 import { chatAtom } from "@/store/chat"
 
 const UNIX_MS_THRESHOLD = 1e12
@@ -33,7 +33,6 @@ function parseTimestamp(dateRaw: number | string | Date) {
   return dayjs(dateRaw)
 }
 
-// Helper to format message timestamps
 export function formatMessageTime(dateRaw: number | string | Date): string {
   const date = parseTimestamp(dateRaw)
   if (!date.isValid()) {
@@ -48,7 +47,6 @@ export function formatMessageTime(dateRaw: number | string | Date): string {
     return date.format("LT")
   }
 
-  // Cross-day formatting
   if (isThisYear) {
     return date.format("MMM D LT")
   }

web/frontend/src/hooks/use-websocket.ts

@@ -1,47 +0,0 @@
-import { useCallback, useEffect, useRef, useState } from "react"
-
-export function useWebSocket(path: string) {
-  const [message, setMessage] = useState<string>("No messages yet")
-  const [connected, setConnected] = useState(false)
-  const wsRef = useRef<WebSocket | null>(null)
-
-  const connect = useCallback(() => {
-    if (wsRef.current) {
-      wsRef.current.close()
-    }
-
-    const protocol = window.location.protocol === "https:" ? "wss:" : "ws:"
-    const url = `${protocol}//${window.location.host}${path}`
-    const socket = new WebSocket(url)
-
-    socket.onopen = () => {
-      setConnected(true)
-      setMessage("Connected to WebSocket server.")
-    }
-
-    socket.onmessage = (event) => {
-      setMessage(event.data)
-    }
-
-    socket.onclose = () => {
-      setConnected(false)
-      setMessage("WebSocket connection closed.")
-    }
-
-    socket.onerror = (error) => {
-      setConnected(false)
-      setMessage("WebSocket error occurred.")
-      console.error("WebSocket Error:", error)
-    }
-
-    wsRef.current = socket
-  }, [path])
-
-  useEffect(() => {
-    return () => {
-      wsRef.current?.close()
-    }
-  }, [])
-
-  return { message, connected, connect }
-}

web/frontend/src/routes/__root.tsx

@@ -3,7 +3,7 @@ import { TanStackRouterDevtools } from "@tanstack/react-router-devtools"
 import { useEffect } from "react"
 
 import { AppLayout } from "@/components/app-layout"
-import { initializeChatStore } from "@/lib/pico-chat-controller"
+import { initializeChatStore } from "@/features/chat/controller"
 
 const RootLayout = () => {
   useEffect(() => {

web/frontend/src/store/chat.ts

@@ -3,7 +3,7 @@ import { atom, getDefaultStore } from "jotai"
 import {
   getInitialActiveSessionId,
   writeStoredSessionId,
-} from "@/lib/pico-chat-state"
+} from "@/features/chat/state"
 
 export interface ChatMessage {
   id: string

web/frontend/src/store/gateway.ts

@@ -31,7 +31,17 @@ function normalizeGatewayStoreState(
   prev: GatewayStoreState,
   patch: GatewayStorePatch,
 ) {
-  return { ...prev, ...patch }
+  const next = { ...prev, ...patch }
+
+  if (
+    next.status === prev.status &&
+    next.canStart === prev.canStart &&
+    next.restartRequired === prev.restartRequired
+  ) {
+    return prev
+  }
+
+  return next
 }
 
 export function updateGatewayStore(