diff --git a/src/langbot/templates/embed/widget.js b/src/langbot/templates/embed/widget.js index 63d00f737..6cd9643a4 100644 --- a/src/langbot/templates/embed/widget.js +++ b/src/langbot/templates/embed/widget.js @@ -535,18 +535,15 @@ return; } - var protocol = CONFIG.baseUrl.indexOf("https") === 0 ? "wss:" : "ws:"; - var host = CONFIG.baseUrl.replace(/^https?:\/\//, ""); - var url = - protocol + - "//" + - host + - "/api/v1/embed/" + - CONFIG.botUuid + - "/ws/connect?session_type=" + - CONFIG.sessionType + - "&session_id=" + - encodeURIComponent(state.sessionId); + var endpoint = new URL(CONFIG.baseUrl || "/", window.location.href); + endpoint.protocol = endpoint.protocol === "https:" || endpoint.protocol === "wss:" ? "wss:" : "ws:"; + endpoint.pathname = endpoint.pathname.replace(/\/$/, "") + + "/api/v1/embed/" + encodeURIComponent(CONFIG.botUuid) + "/ws/connect"; + endpoint.search = ""; + endpoint.hash = ""; + endpoint.searchParams.set("session_type", CONFIG.sessionType); + endpoint.searchParams.set("session_id", state.sessionId); + var url = endpoint.toString(); try { state.ws = new WebSocket(url); diff --git a/tests/unit_tests/platform/test_websocket_session_isolation.py b/tests/unit_tests/platform/test_websocket_session_isolation.py index 18dad8a29..d6c123695 100644 --- a/tests/unit_tests/platform/test_websocket_session_isolation.py +++ b/tests/unit_tests/platform/test_websocket_session_isolation.py @@ -596,7 +596,8 @@ def test_widget_sends_stable_session_id_to_all_conversation_endpoints(): assert 'window.sessionStorage' in widget assert 'window.localStorage' not in widget assert 'session_id=' in widget - assert widget.count('encodeURIComponent(state.sessionId)') >= 3 + assert widget.count('encodeURIComponent(state.sessionId)') >= 2 + assert 'endpoint.searchParams.set("session_id", state.sessionId)' in widget assert 'loadHistory(true)' in widget assert 'messageVersion !== state.messageVersion' in widget assert 'scheduleHistoryReload();' in widget diff --git a/web/src/app/home/pipelines/components/debug-dialog/DebugDialog.tsx b/web/src/app/home/pipelines/components/debug-dialog/DebugDialog.tsx index 9548a16ca..4b8c7be09 100644 --- a/web/src/app/home/pipelines/components/debug-dialog/DebugDialog.tsx +++ b/web/src/app/home/pipelines/components/debug-dialog/DebugDialog.tsx @@ -231,6 +231,7 @@ export default function DebugDialog({ wsClient .onConnected(() => { if (wsClientRef.current !== wsClient) return; + errorReported = false; setIsConnected(true); isInitializingRef.current = false; }) @@ -261,11 +262,12 @@ export default function DebugDialog({ }) .onError((error) => { if (wsClientRef.current !== wsClient) return; + const notify = !errorReported; errorReported = true; console.error('WebSocket error:', error); setIsConnected(false); isInitializingRef.current = false; - toast.error(t('pipelines.debugDialog.connectionError')); + if (notify) toast.error(t('pipelines.debugDialog.connectionError')); }) .onClose(() => { if (wsClientRef.current !== wsClient) return; @@ -332,7 +334,22 @@ export default function DebugDialog({ loadMessages(selectedPipelineId); initWebSocket(selectedPipelineId); } - }, [sessionType, selectedPipelineId, open, loadMessages, initWebSocket]); + return () => { + invalidateHistoryRequests(); + const client = wsClientRef.current; + wsClientRef.current = null; + client?.disconnect(); + isInitializingRef.current = false; + setIsConnected(false); + }; + }, [ + invalidateHistoryRequests, + sessionType, + selectedPipelineId, + open, + loadMessages, + initWebSocket, + ]); // Notify parent of connection status changes useEffect(() => { diff --git a/web/src/app/infra/websocket/WebSocketClient.ts b/web/src/app/infra/websocket/WebSocketClient.ts index 9e930b723..2d79041af 100644 --- a/web/src/app/infra/websocket/WebSocketClient.ts +++ b/web/src/app/infra/websocket/WebSocketClient.ts @@ -3,6 +3,7 @@ * 用于管理WebSocket连接和消息处理 */ import { getActiveWorkspaceUuid } from '@/app/infra/http/workspaceContext'; +import { backendWebSocketURL } from './url'; import type { MessageChainComponent } from '@/app/infra/entities/message'; export interface WebSocketMessage { @@ -17,12 +18,7 @@ export interface WebSocketMessage { export interface WebSocketResponse { type: - | 'connected' - | 'response' - | 'user_message' - | 'pong' - | 'broadcast' - | 'error'; + 'connected' | 'response' | 'user_message' | 'pong' | 'broadcast' | 'error'; connection_id?: string; pipeline_uuid?: string; session_type?: string; @@ -43,6 +39,25 @@ export class WebSocketClient { private isConnecting = false; // 防止重复连接 private shouldReconnect = true; private disconnectedByUser = false; + private workspaceUuid = getActiveWorkspaceUuid(); + private handshakeTimeout: ReturnType | null = null; + private recoveryListenersAttached = false; + private recover = () => { + if ( + document.visibilityState === 'hidden' || + this.disconnectedByUser || + !this.shouldReconnect || + this.isConnecting || + this.isConnected() + ) + return; + if (getActiveWorkspaceUuid() !== this.workspaceUuid) return; + void this.connect().catch(() => {}); + }; + private clearHandshakeTimeout() { + if (this.handshakeTimeout) clearTimeout(this.handshakeTimeout); + this.handshakeTimeout = null; + } // 事件回调 private onConnectedCallback?: (data: WebSocketResponse) => void; @@ -89,17 +104,32 @@ export class WebSocketClient { this.isConnecting = true; - // 构建WebSocket URL - const protocol = window.location.protocol === 'https:' ? 'wss:' : 'ws:'; - // extract host from import.meta.env.VITE_API_BASE_URL - // If env var is undefined, use current page host (for production) - const host = - import.meta.env.VITE_API_BASE_URL?.split('://')[1] || - window.location.host; - const url = `${protocol}//${host}/api/v1/pipelines/${this.pipelineId}/ws/connect?session_type=${this.sessionType}`; - - const socket = new WebSocket(url); + if (getActiveWorkspaceUuid() !== this.workspaceUuid) { + this.isConnecting = false; + this.shouldReconnect = false; + reject(new Error('Workspace changed')); + return; + } + const url = backendWebSocketURL( + import.meta.env.VITE_API_BASE_URL, + window.location.href, + `api/v1/pipelines/${encodeURIComponent(this.pipelineId)}/ws/connect`, + ); + url.searchParams.set('session_type', this.sessionType); + const socket = new WebSocket(url.toString()); this.ws = socket; + if (!this.recoveryListenersAttached) { + window.addEventListener('online', this.recover); + window.addEventListener('focus', this.recover); + document.addEventListener('visibilitychange', this.recover); + this.recoveryListenersAttached = true; + } + this.clearHandshakeTimeout(); + this.handshakeTimeout = setTimeout(() => { + if (this.ws !== socket || this.connectionId) return; + reject(new Error('WebSocket authentication timed out')); + socket.close(); + }, 15000); // 连接打开 socket.onopen = () => { @@ -107,9 +137,8 @@ export class WebSocketClient { socket.close(); return; } - this.isConnecting = false; const token = this.token || localStorage.getItem('token'); - const workspaceUuid = getActiveWorkspaceUuid(); + const workspaceUuid = this.workspaceUuid; if (!token || !workspaceUuid) { const error = new Error('WebSocket认证信息缺失'); this.shouldReconnect = false; @@ -144,6 +173,8 @@ export class WebSocketClient { // Only a fully authenticated runtime connection should reset // the retry budget. Resetting on TCP open makes server-side // errors (for example a pipeline still loading) retry forever. + this.clearHandshakeTimeout(); + this.isConnecting = false; this.reconnectAttempts = 0; this.connectionId = data.connection_id; this.startHeartbeat(); @@ -157,6 +188,9 @@ export class WebSocketClient { // 连接关闭 socket.onclose = () => { + reject(new Error('WebSocket connection closed')); + if (this.ws !== socket) return; + this.clearHandshakeTimeout(); if (this.ws === socket) { this.ws = null; this.connectionId = null; @@ -169,9 +203,12 @@ export class WebSocketClient { // 自动重连 if ( this.shouldReconnect && - this.reconnectAttempts < this.maxReconnectAttempts + getActiveWorkspaceUuid() === this.workspaceUuid ) { - this.reconnectAttempts++; + this.reconnectAttempts = Math.min( + this.reconnectAttempts + 1, + this.maxReconnectAttempts, + ); this.reconnectTimeout = setTimeout(() => { this.reconnectTimeout = null; if (!this.shouldReconnect || this.disconnectedByUser) return; @@ -320,6 +357,11 @@ export class WebSocketClient { */ public disconnect() { this.disconnectedByUser = true; + this.clearHandshakeTimeout(); + window.removeEventListener('online', this.recover); + window.removeEventListener('focus', this.recover); + document.removeEventListener('visibilitychange', this.recover); + this.recoveryListenersAttached = false; this.shouldReconnect = false; this.reconnectAttempts = this.maxReconnectAttempts; if (this.reconnectTimeout) { @@ -354,7 +396,11 @@ export class WebSocketClient { * 获取连接状态 */ public isConnected(): boolean { - return this.ws !== null && this.ws.readyState === WebSocket.OPEN; + return ( + this.ws !== null && + this.ws.readyState === WebSocket.OPEN && + this.connectionId !== null + ); } // ===== 事件回调设置 ===== diff --git a/web/src/app/infra/websocket/url.ts b/web/src/app/infra/websocket/url.ts new file mode 100644 index 000000000..60b92f2ff --- /dev/null +++ b/web/src/app/infra/websocket/url.ts @@ -0,0 +1,15 @@ +/** Match the HTTP backend base, including reverse-proxy path prefixes. */ +export function backendWebSocketURL( + base: string | undefined, + pageURL: string, + path: string, +): URL { + const url = new URL(base || '/', pageURL); + url.protocol = + url.protocol === 'https:' || url.protocol === 'wss:' ? 'wss:' : 'ws:'; + url.pathname = + url.pathname.replace(/\/$/, '') + '/' + path.replace(/^\//, ''); + url.search = ''; + url.hash = ''; + return url; +} diff --git a/web/tests/unit/embed-widget-assistant-dedupe.test.mjs b/web/tests/unit/embed-widget-assistant-dedupe.test.mjs index 42f172d3d..b15d00a2c 100644 --- a/web/tests/unit/embed-widget-assistant-dedupe.test.mjs +++ b/web/tests/unit/embed-widget-assistant-dedupe.test.mjs @@ -129,11 +129,13 @@ function launchWidget() { FakeWebSocket.instances = []; const document = new FakeDocument(); const window = { + location: {href: "https://embed.test/"}, crypto: { randomUUID: () => '00000000-0000-4000-8000-000000000000' }, sessionStorage: { getItem: () => null, setItem: () => {} }, }; const context = vm.createContext({ document, + URL, window, navigator: { clipboard: { writeText: () => Promise.resolve() } }, WebSocket: FakeWebSocket, diff --git a/web/tests/unit/websocket-client.test.mjs b/web/tests/unit/websocket-client.test.mjs new file mode 100644 index 000000000..6f6a1f1a7 --- /dev/null +++ b/web/tests/unit/websocket-client.test.mjs @@ -0,0 +1,52 @@ +import assert from 'node:assert/strict'; +import test from 'node:test'; +import fs from 'node:fs'; +import ts from 'typescript'; +function load(file,deps={},globals={}) { + const source=fs.readFileSync(new URL('../../src/app/infra/websocket/'+file,import.meta.url),'utf8').replaceAll('import.meta.env.VITE_API_BASE_URL','undefined'); + const code=ts.transpileModule(source,{compilerOptions:{module:ts.ModuleKind.CommonJS,target:ts.ScriptTarget.ES2022}}).outputText; + const m={exports:{}};new Function('require','module','exports',...Object.keys(globals),code)(id=>deps[id],m,m.exports,...Object.values(globals));return m.exports; +} +const urls=load('url.ts'); +test('OSS and Cloud URLs follow HTTP backend scheme, port and prefix',()=>{ + for(const [base,page,expected] of [ + [undefined,'http://localhost:3000/home','ws://localhost:3000/api/ws'], + [undefined,'https://cloud.test/home','wss://cloud.test/api/ws'], + ['http://localhost:5300','http://localhost:3000','ws://localhost:5300/api/ws'], + ['https://api.test/root/','https://ui.test','wss://api.test/root/api/ws'], + ['/backend/','https://ui.test/home','wss://ui.test/backend/api/ws'], + ['https://api.test','http://ui.test','wss://api.test/api/ws'], + ]) assert.equal(urls.backendWebSocketURL(base,page,'/api/ws').toString(),expected); +}); +function setup(){ + let workspace='one',next=0;const timers=new Map(),win=new EventTarget(),doc=new EventTarget();win.location={href:'https://cloud.test/home'};doc.visibilityState='visible'; + class WS{ + static CONNECTING=0;static OPEN=1;static instances=[];readyState=0; + constructor(url){this.url=url;WS.instances.push(this)} + send(p){this.sent=JSON.parse(p)} close(){this.readyState=3} closed(){this.readyState=3;this.onclose?.()} + open(){this.readyState=1;this.onopen?.()} receive(data){this.onmessage?.({data:JSON.stringify(data)})} + } + const {WebSocketClient}=load('WebSocketClient.ts',{'./url':urls,'@/app/infra/http/workspaceContext':{getActiveWorkspaceUuid:()=>workspace}},{ + window:win,document:doc,WebSocket:WS,localStorage:{getItem:()=> 'token'}, + setTimeout:fn=>{timers.set(++next,fn);return next},clearTimeout:id=>timers.delete(id),setInterval:()=>1,clearInterval:()=>{},console:{warn(){},error(){},debug(){}}, + }); + return {client:new WebSocketClient('pipeline'),WS,win,timers,setWorkspace:v=>workspace=v}; +} +test('only authenticated sockets are connected; disposal removes recovery',async()=>{ + const {client,WS,win}=setup(),pending=client.connect(),ws=WS.instances[0];ws.open();assert.equal(client.isConnected(),false);assert.equal(ws.sent.workspace_uuid,'one'); + ws.receive({type:'connected',connection_id:'id'});assert.equal(await pending,'id');assert(client.isConnected());client.disconnect();ws.closed();win.dispatchEvent(new Event('focus'));assert.equal(WS.instances.length,1); +}); +test('closed handshakes settle and transient retries survive more than five failures',async()=>{ + const {client,WS,timers}=setup();const first=client.connect().catch(()=>{});WS.instances[0].closed();await first; + for(let i=0;i<7;i++){const [id,fn]=[...timers][0];timers.delete(id);fn();WS.instances.at(-1).closed();await Promise.resolve()} + assert.equal(WS.instances.length,8);assert.equal(timers.size,1);client.disconnect();assert.equal(timers.size,0); +}); +test('stale closes do not affect replacement sockets and workspace changes block recovery',async()=>{ + const {client,WS,win,setWorkspace}=setup();const pending=client.connect().catch(()=>{}),old=WS.instances[0];client.disconnect(); + const current=client.connect(),ws=WS.instances[1];old.closed();await pending;ws.open();ws.receive({type:'connected',connection_id:'new'});await current; + assert(client.isConnected());setWorkspace('two');ws.closed();win.dispatchEvent(new Event('focus'));assert.equal(WS.instances.length,2);client.disconnect(); +}); +test('authentication rejection does not reconnect',async()=>{ + const {client,WS,timers,win}=setup();const pending=client.connect().catch(()=>{}),ws=WS.instances[0];ws.open();ws.receive({type:'error',message:'Unauthorized'});ws.closed();await pending; + win.dispatchEvent(new Event('focus'));assert.equal(timers.size,0);assert.equal(WS.instances.length,1);client.disconnect(); +});