admin/udmin
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat(ws): 新增WebSocket实时通信支持与SSE独立服务

Browse Source 
重构中间件结构，新增ws模块实现WebSocket流程执行实时推送
将SSE服务拆分为独立端口监听，默认8866
优化前端流式模式切换，支持WS/SSE协议选择
统一流式事件处理逻辑，完善错误处理与取消机制
更新Cargo.toml依赖，添加WebSocket相关库
调整代码组织结构，规范导入分组与注释
This commit is contained in:
ayou
2025-09-21 22:15:33 +08:00
parent dd7857940f
commit 30716686ed
23 changed files with 805 additions and 101 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
frontend/src/flows/components/testrun/testrun-panel/index.module.less
View File

@ -11,6 +11,7 @@
justify-content: space-between;
gap: 8px;
margin: 0 12px 8px 0;
flex-wrap: wrap; // 允许在小屏时换行，但尽量保持在一行
.title {
font-size: 15px;
@ -18,6 +19,13 @@
color: #333;
flex: 1;
}
.toggle {
display: inline-flex;
align-items: center;
gap: 6px;
white-space: nowrap;
}
}

101
frontend/src/flows/components/testrun/testrun-panel/index.tsx
View File

@ -50,6 +50,15 @@ export const TestRunSidePanel: FC<TestRunSidePanelProps> = ({ visible, onCancel
_setStreamMode(checked);
localStorage.setItem('testrun-stream-mode', JSON.stringify(checked));
};
// 当启用流式时，选择 WS 或 SSE（默认 SSE）
const [useWS, _setUseWS] = useState<boolean>(() => {
const saved = localStorage.getItem('testrun-ws-mode');
return saved ? JSON.parse(saved) : false;
});
const setUseWS = (checked: boolean) => {
_setUseWS(checked);
localStorage.setItem('testrun-ws-mode', JSON.stringify(checked));
};
// 流式渲染：实时上下文与日志
const [streamCtx, setStreamCtx] = useState<any | undefined>();
@ -95,23 +104,43 @@ export const TestRunSidePanel: FC<TestRunSidePanelProps> = ({ visible, onCancel
}
if (streamMode) {
const { cancel, done } = customService.runStream(values, {
onNode: (evt) => {
if (evt.ctx) setStreamCtx((prev: any) => ({ ...(prev || {}), ...(evt.ctx || {}) }));
if (evt.logs && evt.logs.length) setStreamLogs((prev: string[]) => [...prev, ...evt.logs!]);
},
onError: (evt) => {
const msg = evt.message || I18n.t('Run failed');
setErrors((prev) => [...(prev || []), msg]);
},
onDone: (evt) => {
setResult({ ok: evt.ok, ctx: evt.ctx, logs: evt.logs });
},
onFatal: (err) => {
setErrors((prev) => [...(prev || []), err.message || String(err)]);
setRunning(false);
},
});
const startStream = () => useWS
? customService.runStreamWS(values, {
onNode: (evt) => {
if (evt.ctx) setStreamCtx((prev: any) => ({ ...(prev || {}), ...(evt.ctx || {}) }));
if (evt.logs && evt.logs.length) setStreamLogs((prev: string[]) => [...prev, ...evt.logs!]);
},
onError: (evt) => {
const msg = evt.message || I18n.t('Run failed');
setErrors((prev) => [...(prev || []), msg]);
},
onDone: (evt) => {
setResult({ ok: evt.ok, ctx: evt.ctx, logs: evt.logs });
},
onFatal: (err) => {
setErrors((prev) => [...(prev || []), err.message || String(err)]);
setRunning(false);
},
})
: customService.runStream(values, {
onNode: (evt) => {
if (evt.ctx) setStreamCtx((prev: any) => ({ ...(prev || {}), ...(evt.ctx || {}) }));
if (evt.logs && evt.logs.length) setStreamLogs((prev: string[]) => [...prev, ...evt.logs!]);
},
onError: (evt) => {
const msg = evt.message || I18n.t('Run failed');
setErrors((prev) => [...(prev || []), msg]);
},
onDone: (evt) => {
setResult({ ok: evt.ok, ctx: evt.ctx, logs: evt.logs });
},
onFatal: (err) => {
setErrors((prev) => [...(prev || []), err.message || String(err)]);
setRunning(false);
},
});
const { cancel, done } = startStream();
cancelRef.current = cancel;
@ -212,18 +241,32 @@ export const TestRunSidePanel: FC<TestRunSidePanelProps> = ({ visible, onCancel
<div className={styles['testrun-panel-form']}>
<div className={styles['testrun-panel-input']}>
<div className={styles.title}>{I18n.t('Input Form')}</div>
<div>{I18n.t('JSON Mode')}</div>
<Switch
checked={inputJSONMode}
onChange={(checked: boolean) => setInputJSONMode(checked)}
size="small"
/>
<div>{I18n.t('Streaming Mode')}</div>
<Switch
checked={streamMode}
onChange={(checked: boolean) => setStreamMode(checked)}
size="small"
/>
<div className={styles.toggle}>
<div>{I18n.t('JSON Mode')}</div>
<Switch
checked={inputJSONMode}
onChange={(checked: boolean) => setInputJSONMode(checked)}
size="small"
/>
</div>
<div className={styles.toggle}>
<div>{I18n.t('Streaming Mode')}</div>
<Switch
checked={streamMode}
onChange={(checked: boolean) => setStreamMode(checked)}
size="small"
/>
</div>
{streamMode && (
<div className={styles.toggle}>
<div>WS</div>
<Switch
checked={useWS}
onChange={(checked: boolean) => setUseWS(checked)}
size="small"
/>
</div>
)}
</div>
{renderStatus}
{errors?.map((e) => (

146
frontend/src/flows/services/custom-service.ts
View File

@ -15,6 +15,7 @@ import { I18n } from '@flowgram.ai/free-layout-editor';
import api, { type ApiResp } from '../../utils/axios';
import { stringifyFlowDoc } from '../utils/yaml';
import { postSSE } from '../../utils/sse';
import { getToken } from '../../utils/token'
interface RunResult { ok: boolean; ctx: any; logs: string[] }
@ -124,7 +125,126 @@ export class CustomService {
}
}
// 新增：SSE 流式运行，返回取消函数与完成 Promise
// 新增：WebSocket 流式运行
runStreamWS(
input: any = {},
handlers?: {
onNode?: (e: StreamEvent & { type: 'node' }) => void;
onDone?: (e: StreamEvent & { type: 'done' }) => void;
onError?: (e: StreamEvent & { type: 'error' }) => void;
onFatal?: (err: Error) => void;
}
) {
const id = getFlowIdFromUrl();
if (!id) {
const err = new Error(I18n.t('Flow ID is missing, cannot run'));
handlers?.onFatal?.(err);
return { cancel: () => {}, done: Promise.resolve<RunResult | null>(null) } as const;
}
// 构造 WS URL
const base = (api.defaults.baseURL || '') as string; // 可能是 /api 或 http(s)://host/api
function toWsUrl(httpUrl: string) {
if (httpUrl.startsWith('https://')) return 'wss://' + httpUrl.slice('https://'.length);
if (httpUrl.startsWith('http://')) return 'ws://' + httpUrl.slice('http://'.length);
// 相对路径：拼 window.location
const origin = window.location.origin; // http(s)://host:port
const full = origin.replace(/^http/, 'ws') + (httpUrl.startsWith('/') ? httpUrl : '/' + httpUrl);
return full;
}
const path = `/flows/${id}/run/ws`;
// 取 token 放到查询参数（WS 握手无法自定义 Authorization 头部）
const token = getToken();
// 新增：WS 使用独立端口，默认 8855，可通过 VITE_WS_PORT 覆盖
const wsPort = (import.meta as any).env?.VITE_WS_PORT || '8855';
let wsBase: string;
if (base.startsWith('http://') || base.startsWith('https://')) {
try {
const u = new URL(base);
const proto = u.protocol === 'https:' ? 'wss:' : 'ws:';
u.protocol = proto;
u.port = wsPort; // 改为 WS 端口
wsBase = `${u.protocol}//${u.host}${u.pathname.replace(/\/$/, '')}`;
} catch {
const loc = window.location;
const proto = loc.protocol === 'https:' ? 'wss:' : 'ws:';
wsBase = `${proto}//${loc.hostname}:${wsPort}${base.startsWith('/') ? base : '/' + base}`;
}
} else {
const loc = window.location;
const proto = loc.protocol === 'https:' ? 'wss:' : 'ws:';
wsBase = `${proto}//${loc.hostname}:${wsPort}${base.startsWith('/') ? base : '/' + base}`;
}
const wsUrl = wsBase + path + (token ? (wsBase.includes('?') ? `&access_token=${encodeURIComponent(token)}` : `?access_token=${encodeURIComponent(token)}`) : '');
let ws: WebSocket | null = null;
let resolveDone: (v: RunResult | null) => void;
let rejectDone: (e: any) => void;
const done = new Promise<RunResult | null>((resolve, reject) => { resolveDone = resolve; rejectDone = reject; });
let finished = false;
try {
ws = new WebSocket(wsUrl);
} catch (e: any) {
handlers?.onFatal?.(e);
return { cancel: () => {}, done: Promise.resolve<RunResult | null>(null) } as const;
}
ws.onopen = () => {
try {
ws?.send(JSON.stringify({ input }));
} catch (e: any) {
handlers?.onFatal?.(e);
}
};
ws.onmessage = (ev: MessageEvent) => {
try {
const data = typeof ev.data === 'string' ? ev.data : '' + ev.data;
const evt = JSON.parse(data) as StreamEvent;
if (evt.type === 'node') {
handlers?.onNode?.(evt as any);
return;
}
if (evt.type === 'error') {
handlers?.onError?.(evt as any);
return;
}
if (evt.type === 'done') {
finished = true;
handlers?.onDone?.(evt as any);
resolveDone({ ok: evt.ok, ctx: (evt as any).ctx, logs: (evt as any).logs });
ws?.close();
return;
}
} catch (e: any) {
// 忽略解析错误为致命，仅记录
console.warn('WS message parse error', e);
}
};
ws.onerror = (ev: Event) => {
if (!finished) {
handlers?.onFatal?.(new Error('WebSocket error'));
}
};
ws.onclose = () => {
if (!finished) {
resolveDone(null);
}
};
const cancel = () => {
try { finished = true; ws?.close(); } catch {}
};
return { cancel, done } as const;
}
// 现有：SSE 流式运行
runStream(input: any = {}, handlers?: { onNode?: (e: StreamEvent & { type: 'node' }) => void; onDone?: (e: StreamEvent & { type: 'done' }) => void; onError?: (e: StreamEvent & { type: 'error' }) => void; onFatal?: (err: Error) => void; }) {
const id = getFlowIdFromUrl();
if (!id) {
@ -134,7 +254,25 @@ export class CustomService {
}
const base = (api.defaults.baseURL || '') as string;
const url = base ? `${base}/flows/${id}/run/stream` : `/flows/${id}/run/stream`;
// 参照 WS：SSE 使用独立端口，默认 8866，可通过 VITE_SSE_PORT 覆盖
const ssePort = (import.meta as any).env?.VITE_SSE_PORT || '8866';
let sseBase: string;
if (base.startsWith('http://') || base.startsWith('https://')) {
try {
const u = new URL(base);
// 协议保持与 base 一致，仅替换端口
u.port = ssePort;
sseBase = `${u.protocol}//${u.host}${u.pathname.replace(/\/$/, '')}`;
} catch {
const loc = window.location;
sseBase = `${loc.protocol}//${loc.hostname}:${ssePort}${base.startsWith('/') ? base : '/' + base}`;
}
} else {
const loc = window.location;
sseBase = `${loc.protocol}//${loc.hostname}:${ssePort}${base.startsWith('/') ? base : '/' + base}`;
}
const url = sseBase + `/flows/${id}/run/stream`;
const { cancel, done } = postSSE<RunResult | null>(url, { input }, {
onMessage: (json: any) => {
@ -150,12 +288,12 @@ export class CustomService {
}
if (evt.type === 'done') {
handlers?.onDone?.(evt as any)
return { ok: evt.ok, ctx: evt.ctx, logs: evt.logs }
return { ok: (evt as any).ok, ctx: (evt as any).ctx, logs: (evt as any).logs }
}
} catch (_) {}
return undefined
},
onFatal: (e) => handlers?.onFatal?.(e),
onFatal: (e: any) => handlers?.onFatal?.(e),
})
return { cancel, done } as const;
}

123
frontend/src/utils/__tests__/sse.test.ts Normal file
View File

@ -0,0 +1,123 @@
import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest'
import { postSSE } from '../sse'
// Mocks for axios instance and token utils used inside postSSE
vi.mock('../axios', () => {
return {
default: {
get: vi.fn(async () => ({ data: { code: 0, data: { access_token: 'tok2' } } })),
defaults: { baseURL: '/api' },
},
}
})
let tokenSeq: string[] = ['tok0']
vi.mock('../token', () => {
return { getToken: vi.fn(() => (tokenSeq.length ? tokenSeq.shift() : undefined)) }
})
const encoder = new TextEncoder()
function makeSSEStream(chunks: string[]): ReadableStream<Uint8Array> {
return new ReadableStream<Uint8Array>({
start(controller) {
for (const c of chunks) controller.enqueue(encoder.encode(c))
controller.close()
},
})
}
describe('postSSE', () => {
const originalFetch = global.fetch as any
beforeEach(() => {
vi.useFakeTimers()
tokenSeq = ['tok0']
})
afterEach(() => {
vi.useRealTimers()
vi.restoreAllMocks()
// restore fetch if we replaced it
if (originalFetch) {
// @ts-ignore
global.fetch = originalFetch
}
})
it('streams node and done events (CRLF) and resolves with done payload', async () => {
const bodyStream = makeSSEStream([
'data: {"type":"node","node_id":"n1","logs":["l1"],"server_ts":"2025-01-01T00:00:00.000Z"}\r\n\r\n',
'data: {"type":"done","ok":true,"ctx":{},"logs":[],"server_ts":"2025-01-01T00:00:01.000Z"}\r\n\r\n',
])
const resp = new Response(bodyStream, { status: 200 })
const fetchMock = vi.fn(async () => resp)
// replace global fetch
// @ts-ignore
global.fetch = fetchMock
const events: any[] = []
const { done } = postSSE<{ ok: boolean; ctx: any; logs: string[] }>('http://example.com/sse', { a: 1 }, {
onMessage: (json) => {
events.push(json)
if (json.type === 'done') return { ok: json.ok, ctx: json.ctx, logs: json.logs }
},
})
const result = await done
expect(result).toEqual({ ok: true, ctx: {}, logs: [] })
expect(events.some(e => e.type === 'node')).toBe(true)
expect(events.some(e => e.type === 'done')).toBe(true)
expect(fetchMock).toHaveBeenCalledTimes(1)
const call = fetchMock.mock.calls[0] as any[]
const init = call?.[1]
expect((init as any).headers.Authorization).toBe('Bearer tok0')
})
it('refreshes on 401 then retries with new token and resolves', async () => {
tokenSeq = ['tok1', 'tok2']
const first = new Response(null, { status: 401 })
const sse = makeSSEStream(['data: {"type":"done","ok":true,"ctx":{},"logs":[]}\n\n'])
const second = new Response(sse, { status: 200 })
const fetchMock = vi.fn()
.mockResolvedValueOnce(first)
.mockResolvedValueOnce(second)
// replace global fetch
// @ts-ignore
global.fetch = fetchMock
const { done } = postSSE('http://example.com/sse', {}, {
onMessage: (j) => { if (j.type === 'done') return { ok: j.ok, ctx: j.ctx, logs: j.logs } },
})
const res = await done
expect(res).toEqual({ ok: true, ctx: {}, logs: [] })
expect(fetchMock).toHaveBeenCalledTimes(2)
const call1 = fetchMock.mock.calls[0] as any[]
const call2 = fetchMock.mock.calls[1] as any[]
const init1 = call1?.[1]
const init2 = call2?.[1]
expect((init1 as any).headers.Authorization).toBe('Bearer tok1')
expect((init2 as any).headers.Authorization).toBe('Bearer tok2')
})
it('cancel aborts the stream and resolves to null', async () => {
// a long stream without done
const longStream = new ReadableStream<Uint8Array>({
start(controller) {
controller.enqueue(encoder.encode('data: {"type":"node","logs":["l1"]}\n\n'))
// do not close to simulate long running
},
pull() {},
cancel() {},
})
const fetchMock = vi.fn(async () => new Response(longStream, { status: 200 }))
// replace global fetch
// @ts-ignore
global.fetch = fetchMock
const { cancel, done } = postSSE('http://example.com/sse', {}, { onMessage: () => {} })
// cancel immediately
cancel()
const out = await done
expect(out).toBeNull()
})
})

1
frontend/vite.config.ts
View File

@ -31,6 +31,7 @@ export default defineConfig(({ mode }) => {
'/api': {
target: proxyTarget,
changeOrigin: true,
ws: true,
// 为 SSE 透传加固：禁用超时并保持连接
proxyTimeout: 0,
timeout: 0,