diff --git a/web/backend/api/gateway_host.go b/web/backend/api/gateway_host.go index a499c1ea2..5ef3ba2c5 100644 --- a/web/backend/api/gateway_host.go +++ b/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" } diff --git a/web/backend/api/gateway_host_test.go b/web/backend/api/gateway_host_test.go index afd600359..43e84ff0e 100644 --- a/web/backend/api/gateway_host_test.go +++ b/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") + } +} diff --git a/web/frontend/src/components/chat/chat-composer.tsx b/web/frontend/src/components/chat/chat-composer.tsx index e8bae89b8..7d696b898 100644 --- a/web/frontend/src/components/chat/chat-composer.tsx +++ b/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} > diff --git a/web/frontend/src/components/chat/chat-empty-state.tsx b/web/frontend/src/components/chat/chat-empty-state.tsx index 624ff9c59..0574c44d1 100644 --- a/web/frontend/src/components/chat/chat-empty-state.tsx +++ b/web/frontend/src/components/chat/chat-empty-state.tsx @@ -34,7 +34,7 @@ export function ChatEmptyState({

{t("chat.empty.noConfiguredModelDescription")}

- diff --git a/web/frontend/src/components/chat/chat-page.tsx b/web/frontend/src/components/chat/chat-page.tsx index 1906a0367..ebcde8981 100644 --- a/web/frontend/src/components/chat/chat-page.tsx +++ b/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() { )} @@ -168,7 +167,7 @@ export function ChatPage() { input={input} onInputChange={setInput} onSend={handleSend} - isConnected={isConnected} + isConnected={isChatConnected} hasDefaultModel={Boolean(defaultModelName)} /> diff --git a/web/frontend/src/lib/pico-chat-controller.ts b/web/frontend/src/features/chat/controller.ts similarity index 53% rename from web/frontend/src/lib/pico-chat-controller.ts rename to web/frontend/src/features/chat/controller.ts index 0e77d1ad0..5e6eb2229 100644 --- a/web/frontend/src/lib/pico-chat-controller.ts +++ b/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 -} +} 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 | null = null let connectionGeneration = 0 +let reconnectTimer: number | null = null +let reconnectAttempts = 0 +let shouldMaintainConnection = false -async function loadSessionMessages(sessionId: string): Promise { - 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() { diff --git a/web/frontend/src/features/chat/history.ts b/web/frontend/src/features/chat/history.ts new file mode 100644 index 000000000..886148184 --- /dev/null +++ b/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 { + 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), + ) +} diff --git a/web/frontend/src/features/chat/protocol.ts b/web/frontend/src/features/chat/protocol.ts new file mode 100644 index 000000000..5e5220c77 --- /dev/null +++ b/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 +} + +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) + } +} diff --git a/web/frontend/src/lib/pico-chat-state.ts b/web/frontend/src/features/chat/state.ts similarity index 100% rename from web/frontend/src/lib/pico-chat-state.ts rename to web/frontend/src/features/chat/state.ts diff --git a/web/frontend/src/features/chat/websocket.ts b/web/frontend/src/features/chat/websocket.ts new file mode 100644 index 000000000..6b132e9a6 --- /dev/null +++ b/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 + ) +} diff --git a/web/frontend/src/hooks/use-gateway.ts b/web/frontend/src/hooks/use-gateway.ts index 848f4d59c..65ec2b776 100644 --- a/web/frontend/src/hooks/use-gateway.ts +++ b/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 { diff --git a/web/frontend/src/hooks/use-pico-chat.ts b/web/frontend/src/hooks/use-pico-chat.ts index 1b97a2a9c..3ac2e1613 100644 --- a/web/frontend/src/hooks/use-pico-chat.ts +++ b/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") } diff --git a/web/frontend/src/hooks/use-websocket.ts b/web/frontend/src/hooks/use-websocket.ts deleted file mode 100644 index c41b5ed34..000000000 --- a/web/frontend/src/hooks/use-websocket.ts +++ /dev/null @@ -1,47 +0,0 @@ -import { useCallback, useEffect, useRef, useState } from "react" - -export function useWebSocket(path: string) { - const [message, setMessage] = useState("No messages yet") - const [connected, setConnected] = useState(false) - const wsRef = useRef(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 } -} diff --git a/web/frontend/src/routes/__root.tsx b/web/frontend/src/routes/__root.tsx index 6431d9490..31fdb7804 100644 --- a/web/frontend/src/routes/__root.tsx +++ b/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(() => { diff --git a/web/frontend/src/store/chat.ts b/web/frontend/src/store/chat.ts index d79a1a93b..da5fa6670 100644 --- a/web/frontend/src/store/chat.ts +++ b/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 diff --git a/web/frontend/src/store/gateway.ts b/web/frontend/src/store/gateway.ts index b7655839c..c5eee8451 100644 --- a/web/frontend/src/store/gateway.ts +++ b/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(