feat: 对接设计 Agent Gateway WebSocket
需求:服务端统一 Agent Gateway 将设计任务状态流切换为 WebSocket,客户端需要实时展示生成任务并支持断线恢复。 实现:Electron Main 管理 Session、一次性 Ticket、WebSocket 心跳与游标续传,按关闭码回收会话;Renderer 继续通过本机 Host API 的 SSE 投影接收任务事件,并保留 REST 降级同步。 验证:typecheck、变更文件 ESLint、37 个聚焦测试及 build:vite 通过。
This commit is contained in:
@@ -49,6 +49,68 @@ function jsonResponse(payload: unknown, status = 200): Response {
|
||||
});
|
||||
}
|
||||
|
||||
type MockSocketScript = {
|
||||
frames?: unknown[];
|
||||
open?: boolean;
|
||||
closeCode?: number;
|
||||
closeReason?: string;
|
||||
};
|
||||
|
||||
class MockAgentWebSocket {
|
||||
readonly url: string;
|
||||
readonly sent: string[] = [];
|
||||
readyState = 0;
|
||||
onopen: (() => void) | null = null;
|
||||
onmessage: ((event: { data: unknown }) => void) | null = null;
|
||||
onerror: ((event: unknown) => void) | null = null;
|
||||
onclose: ((event: { code: number; reason: string }) => void) | null = null;
|
||||
private closed = false;
|
||||
|
||||
constructor(url: string, private readonly script: MockSocketScript) {
|
||||
this.url = url;
|
||||
setTimeout(() => this.runScript(), 0);
|
||||
}
|
||||
|
||||
send(data: string): void {
|
||||
this.sent.push(data);
|
||||
}
|
||||
|
||||
close(code = 1000, reason = ''): void {
|
||||
this.emitClose(code, reason);
|
||||
}
|
||||
|
||||
private runScript(): void {
|
||||
if (this.closed) return;
|
||||
if (this.script.open !== false) {
|
||||
this.readyState = 1;
|
||||
this.onopen?.();
|
||||
for (const frame of this.script.frames ?? []) {
|
||||
this.onmessage?.({ data: JSON.stringify(frame) });
|
||||
}
|
||||
}
|
||||
if (this.script.closeCode !== undefined) {
|
||||
this.emitClose(this.script.closeCode, this.script.closeReason ?? '');
|
||||
}
|
||||
}
|
||||
|
||||
private emitClose(code: number, reason: string): void {
|
||||
if (this.closed) return;
|
||||
this.closed = true;
|
||||
this.readyState = 3;
|
||||
this.onclose?.({ code, reason });
|
||||
}
|
||||
}
|
||||
|
||||
function scriptedSockets(scripts: MockSocketScript[]) {
|
||||
const sockets: MockAgentWebSocket[] = [];
|
||||
const webSocketFactory = vi.fn((url: string) => {
|
||||
const socket = new MockAgentWebSocket(url, scripts[sockets.length] ?? { closeCode: 1000 });
|
||||
sockets.push(socket);
|
||||
return { socket };
|
||||
});
|
||||
return { sockets, webSocketFactory };
|
||||
}
|
||||
|
||||
describe('Works Square AI design adapter', () => {
|
||||
beforeEach(() => {
|
||||
getTokenMock.mockReset();
|
||||
@@ -193,4 +255,356 @@ describe('Works Square AI design adapter', () => {
|
||||
}),
|
||||
);
|
||||
});
|
||||
|
||||
it('reuses one design Agent Session and normalizes matching task events from fresh WebSocket tickets', async () => {
|
||||
const snapshotTask = {
|
||||
task_id: 'task-snapshot',
|
||||
workspace_id: 'workspace-one',
|
||||
medium: 'image',
|
||||
status: 'running',
|
||||
brief_version: 1,
|
||||
brief_summary: 'snapshot brief',
|
||||
quote_id: 'quote-snapshot',
|
||||
quoted_design_points: 1,
|
||||
failure_code: null,
|
||||
result_assets: [],
|
||||
created_at: '2026-08-02T09:59:00Z',
|
||||
updated_at: '2026-08-02T10:00:00Z',
|
||||
};
|
||||
const snapshotEvent = {
|
||||
session_id: 'session-one',
|
||||
sequence: 1,
|
||||
runtime: 'design',
|
||||
type: 'design.workspace.updated',
|
||||
schema_version: 1,
|
||||
payload: {
|
||||
workspace: {
|
||||
workspace_id: 'workspace-one',
|
||||
view_revision: 2,
|
||||
},
|
||||
generation_tasks: [snapshotTask],
|
||||
},
|
||||
};
|
||||
const taskEvent = {
|
||||
session_id: 'session-one',
|
||||
sequence: 3,
|
||||
runtime: 'design',
|
||||
type: 'design.generation_task.updated',
|
||||
command_id: null,
|
||||
run_id: null,
|
||||
client_command_id: null,
|
||||
schema_version: 1,
|
||||
terminal: false,
|
||||
occurred_at: '2026-08-02T10:01:00Z',
|
||||
payload: {
|
||||
workspace_id: 'workspace-one',
|
||||
workspace_view_revision: 4,
|
||||
generation_task: {
|
||||
task_id: 'task-live',
|
||||
workspace_id: 'workspace-one',
|
||||
medium: 'video',
|
||||
status: 'running',
|
||||
brief_version: 2,
|
||||
brief_summary: '海洋公益短片',
|
||||
quote_id: 'quote-live',
|
||||
quoted_design_points: 8,
|
||||
failure_code: null,
|
||||
result_assets: [],
|
||||
created_at: '2026-08-02T10:00:00Z',
|
||||
updated_at: '2026-08-02T10:01:00Z',
|
||||
},
|
||||
},
|
||||
};
|
||||
const fetchMock = vi.fn<typeof fetch>()
|
||||
.mockResolvedValueOnce(jsonResponse({
|
||||
session_id: 'session-one',
|
||||
status: 'active',
|
||||
}, 201))
|
||||
.mockResolvedValueOnce(jsonResponse({
|
||||
ticket: 'secret-ticket-one',
|
||||
transport: 'websocket',
|
||||
stream_url: '/api/agents/sessions/session-one/ws?ticket=secret-ticket-one',
|
||||
expires_at: '2026-08-02T10:02:00Z',
|
||||
}))
|
||||
.mockResolvedValueOnce(jsonResponse({
|
||||
ticket: 'secret-ticket-two',
|
||||
transport: 'websocket',
|
||||
stream_url: '/api/agents/sessions/session-one/ws?ticket=secret-ticket-two',
|
||||
expires_at: '2026-08-02T10:03:00Z',
|
||||
}));
|
||||
const { sockets, webSocketFactory } = scriptedSockets([
|
||||
{
|
||||
frames: [
|
||||
{ type: 'event', event: snapshotEvent },
|
||||
{ type: 'event', event: taskEvent },
|
||||
],
|
||||
closeCode: 1000,
|
||||
},
|
||||
{ closeCode: 1000 },
|
||||
]);
|
||||
const adapter = new WorksSquareDesignWorkspace({
|
||||
apiBaseUrl: 'https://square.example',
|
||||
fetchImpl: fetchMock,
|
||||
webSocketFactory,
|
||||
});
|
||||
|
||||
const first = await adapter.openWorkspaceEvents({ workspaceId: 'workspace-one' });
|
||||
const received = [];
|
||||
for await (const event of first.events) received.push(event);
|
||||
first.close();
|
||||
const second = await adapter.openWorkspaceEvents({
|
||||
workspaceId: 'workspace-one',
|
||||
afterEventId: 'session-one:3',
|
||||
});
|
||||
for await (const _event of second.events) {
|
||||
// The second connection only proves Session reuse and a fresh ticket.
|
||||
}
|
||||
second.close();
|
||||
|
||||
expect(received).toEqual([
|
||||
{
|
||||
id: 'session-one:1',
|
||||
type: 'design.generation_tasks.snapshot',
|
||||
workspaceId: 'workspace-one',
|
||||
workspaceViewRevision: 2,
|
||||
generationTasks: [expect.objectContaining({
|
||||
taskId: 'task-snapshot',
|
||||
medium: 'image',
|
||||
status: 'running',
|
||||
})],
|
||||
},
|
||||
{
|
||||
id: 'session-one:3',
|
||||
type: 'design.generation_task.updated',
|
||||
workspaceId: 'workspace-one',
|
||||
workspaceViewRevision: 4,
|
||||
generationTask: expect.objectContaining({
|
||||
taskId: 'task-live',
|
||||
medium: 'video',
|
||||
status: 'running',
|
||||
}),
|
||||
},
|
||||
]);
|
||||
expect(fetchMock).toHaveBeenNthCalledWith(
|
||||
1,
|
||||
'https://square.example/api/agents/sessions',
|
||||
expect.objectContaining({
|
||||
method: 'POST',
|
||||
body: expect.stringContaining('"runtime":"design"'),
|
||||
}),
|
||||
);
|
||||
expect(fetchMock.mock.calls.filter(([url]) => String(url).endsWith('/api/agents/sessions')))
|
||||
.toHaveLength(1);
|
||||
expect(fetchMock.mock.calls.filter(([url]) => String(url).endsWith('/stream-tickets')))
|
||||
.toHaveLength(2);
|
||||
const ticketCalls = fetchMock.mock.calls.filter(([url]) => (
|
||||
String(url).endsWith('/stream-tickets')
|
||||
));
|
||||
expect(ticketCalls.every(([, init]) => (
|
||||
JSON.parse(String(init?.body)).transport === 'websocket'
|
||||
))).toBe(true);
|
||||
expect(sockets.map((socket) => socket.url)).toEqual([
|
||||
'wss://square.example/api/agents/sessions/session-one/ws?ticket=secret-ticket-one&after_sequence=0',
|
||||
'wss://square.example/api/agents/sessions/session-one/ws?ticket=secret-ticket-two&after_sequence=3',
|
||||
]);
|
||||
});
|
||||
|
||||
it('rotates the Session and client id after an upstream event cursor expires', async () => {
|
||||
const rotate = vi.fn().mockResolvedValue('design-stream-next');
|
||||
const fetchMock = vi.fn<typeof fetch>()
|
||||
.mockResolvedValueOnce(jsonResponse({ session_id: 'session-old', status: 'active' }, 201))
|
||||
.mockResolvedValueOnce(jsonResponse({
|
||||
stream_url: '/api/agents/sessions/session-old/ws?ticket=ticket-old',
|
||||
}))
|
||||
.mockResolvedValueOnce(jsonResponse({
|
||||
session_id: 'session-old', status: 'closed',
|
||||
}))
|
||||
.mockResolvedValueOnce(jsonResponse({ session_id: 'session-new', status: 'active' }, 201))
|
||||
.mockResolvedValueOnce(jsonResponse({
|
||||
stream_url: '/api/agents/sessions/session-new/ws?ticket=ticket-new',
|
||||
}));
|
||||
const { sockets, webSocketFactory } = scriptedSockets([
|
||||
{ open: false, closeCode: 4409, closeReason: 'Agent event cursor expired' },
|
||||
{ closeCode: 1000 },
|
||||
]);
|
||||
const adapter = new WorksSquareDesignWorkspace({
|
||||
apiBaseUrl: 'https://square.example',
|
||||
fetchImpl: fetchMock,
|
||||
webSocketFactory,
|
||||
eventSessionClientIdStore: {
|
||||
getOrCreate: vi.fn().mockResolvedValue('design-stream-current'),
|
||||
rotate,
|
||||
},
|
||||
});
|
||||
|
||||
await expect(adapter.openWorkspaceEvents({
|
||||
workspaceId: 'workspace-one',
|
||||
afterEventId: 'session-old:99',
|
||||
})).rejects.toMatchObject({ status: 410 });
|
||||
const recovered = await adapter.openWorkspaceEvents({
|
||||
workspaceId: 'workspace-one',
|
||||
afterEventId: 'session-old:99',
|
||||
});
|
||||
for await (const _event of recovered.events) {
|
||||
// Empty recovery stream.
|
||||
}
|
||||
|
||||
const sessionCalls = fetchMock.mock.calls.filter(([url]) => (
|
||||
String(url).endsWith('/api/agents/sessions')
|
||||
));
|
||||
expect(sessionCalls).toHaveLength(2);
|
||||
expect(JSON.parse(String(sessionCalls[0][1]?.body)).client_session_id)
|
||||
.not.toBe(JSON.parse(String(sessionCalls[1][1]?.body)).client_session_id);
|
||||
expect(rotate).toHaveBeenCalledWith('workspace-one');
|
||||
expect(fetchMock.mock.calls[2]).toEqual([
|
||||
'https://square.example/api/agents/sessions/session-old',
|
||||
expect.objectContaining({ method: 'DELETE' }),
|
||||
]);
|
||||
expect(sockets.map((socket) => socket.url)).toEqual([
|
||||
'wss://square.example/api/agents/sessions/session-old/ws?ticket=ticket-old&after_sequence=99',
|
||||
'wss://square.example/api/agents/sessions/session-new/ws?ticket=ticket-new&after_sequence=0',
|
||||
]);
|
||||
});
|
||||
|
||||
it('replaces a closed cached Session before opening the task stream', async () => {
|
||||
const fetchMock = vi.fn<typeof fetch>()
|
||||
.mockResolvedValueOnce(jsonResponse({ session_id: 'session-old', status: 'active' }, 201))
|
||||
.mockResolvedValueOnce(jsonResponse({
|
||||
detail: { code: 'agent_session_closed', message: 'closed' },
|
||||
}, 409))
|
||||
.mockResolvedValueOnce(jsonResponse({ session_id: 'session-new', status: 'active' }, 201))
|
||||
.mockResolvedValueOnce(jsonResponse({
|
||||
stream_url: '/api/agents/sessions/session-new/ws?ticket=ticket-new',
|
||||
}));
|
||||
const { webSocketFactory } = scriptedSockets([{ closeCode: 1000 }]);
|
||||
const adapter = new WorksSquareDesignWorkspace({
|
||||
apiBaseUrl: 'https://square.example',
|
||||
fetchImpl: fetchMock,
|
||||
webSocketFactory,
|
||||
});
|
||||
|
||||
const recovered = await adapter.openWorkspaceEvents({ workspaceId: 'workspace-one' });
|
||||
for await (const _event of recovered.events) {
|
||||
// Empty recovery stream.
|
||||
}
|
||||
|
||||
const sessionCalls = fetchMock.mock.calls.filter(([url]) => (
|
||||
String(url).endsWith('/api/agents/sessions')
|
||||
));
|
||||
expect(sessionCalls).toHaveLength(2);
|
||||
expect(JSON.parse(String(sessionCalls[0][1]?.body)).client_session_id)
|
||||
.not.toBe(JSON.parse(String(sessionCalls[1][1]?.body)).client_session_id);
|
||||
});
|
||||
|
||||
it('closes every cached Agent Session during logout or application shutdown', async () => {
|
||||
const fetchMock = vi.fn<typeof fetch>()
|
||||
.mockResolvedValueOnce(jsonResponse({ session_id: 'session-one', status: 'active' }, 201))
|
||||
.mockResolvedValueOnce(jsonResponse({
|
||||
stream_url: '/api/agents/sessions/session-one/ws?ticket=ticket-one',
|
||||
}))
|
||||
.mockResolvedValueOnce(jsonResponse({ session_id: 'session-two', status: 'active' }, 201))
|
||||
.mockResolvedValueOnce(jsonResponse({
|
||||
stream_url: '/api/agents/sessions/session-two/ws?ticket=ticket-two',
|
||||
}))
|
||||
.mockImplementation(() => Promise.resolve(
|
||||
jsonResponse({ session_id: 'closed', status: 'closed' }),
|
||||
));
|
||||
const { webSocketFactory } = scriptedSockets([
|
||||
{ closeCode: 1000 },
|
||||
{ closeCode: 1000 },
|
||||
]);
|
||||
const adapter = new WorksSquareDesignWorkspace({
|
||||
apiBaseUrl: 'https://square.example',
|
||||
fetchImpl: fetchMock,
|
||||
webSocketFactory,
|
||||
});
|
||||
|
||||
const first = await adapter.openWorkspaceEvents({ workspaceId: 'workspace-one' });
|
||||
for await (const _event of first.events) {
|
||||
// Empty stream.
|
||||
}
|
||||
const second = await adapter.openWorkspaceEvents({ workspaceId: 'workspace-two' });
|
||||
for await (const _event of second.events) {
|
||||
// Empty stream.
|
||||
}
|
||||
await adapter.closeEventSessions();
|
||||
|
||||
const closeCalls = fetchMock.mock.calls.filter(([, init]) => init?.method === 'DELETE');
|
||||
expect(closeCalls.map(([url]) => String(url)).sort()).toEqual([
|
||||
'https://square.example/api/agents/sessions/session-one',
|
||||
'https://square.example/api/agents/sessions/session-two',
|
||||
]);
|
||||
});
|
||||
|
||||
it('keeps the persisted Session key when a close result is uncertain', async () => {
|
||||
const rotate = vi.fn().mockResolvedValue('design-stream-next');
|
||||
const fetchMock = vi.fn<typeof fetch>()
|
||||
.mockResolvedValueOnce(jsonResponse({ session_id: 'session-one', status: 'active' }, 201))
|
||||
.mockResolvedValueOnce(jsonResponse({
|
||||
stream_url: '/api/agents/sessions/session-one/ws?ticket=ticket-one',
|
||||
}))
|
||||
.mockResolvedValueOnce(jsonResponse({
|
||||
detail: { code: 'service_unavailable', message: 'offline' },
|
||||
}, 503));
|
||||
const { webSocketFactory } = scriptedSockets([{ closeCode: 1000 }]);
|
||||
const adapter = new WorksSquareDesignWorkspace({
|
||||
apiBaseUrl: 'https://square.example',
|
||||
fetchImpl: fetchMock,
|
||||
webSocketFactory,
|
||||
eventSessionClientIdStore: {
|
||||
getOrCreate: vi.fn().mockResolvedValue('design-stream-current'),
|
||||
rotate,
|
||||
},
|
||||
});
|
||||
|
||||
const stream = await adapter.openWorkspaceEvents({ workspaceId: 'workspace-one' });
|
||||
for await (const _event of stream.events) {
|
||||
// Empty stream.
|
||||
}
|
||||
|
||||
await expect(adapter.closeEventSessions()).rejects.toThrow(
|
||||
'Failed to close 1 AI design Agent Session(s)',
|
||||
);
|
||||
expect(rotate).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it('reuses a stable Session idempotency key after an unclean application restart', async () => {
|
||||
const createFetch = () => vi.fn<typeof fetch>()
|
||||
.mockResolvedValueOnce(jsonResponse({ session_id: 'session-stable', status: 'active' }, 201))
|
||||
.mockResolvedValueOnce(jsonResponse({
|
||||
stream_url: '/api/agents/sessions/session-stable/ws?ticket=ticket-stable',
|
||||
}));
|
||||
const firstFetch = createFetch();
|
||||
const secondFetch = createFetch();
|
||||
const firstSockets = scriptedSockets([{ closeCode: 1000 }]);
|
||||
const restartedSockets = scriptedSockets([{ closeCode: 1000 }]);
|
||||
const first = new WorksSquareDesignWorkspace({
|
||||
apiBaseUrl: 'https://square.example',
|
||||
fetchImpl: firstFetch,
|
||||
clientInstanceId: 'installation-one',
|
||||
webSocketFactory: firstSockets.webSocketFactory,
|
||||
});
|
||||
const restarted = new WorksSquareDesignWorkspace({
|
||||
apiBaseUrl: 'https://square.example',
|
||||
fetchImpl: secondFetch,
|
||||
clientInstanceId: 'installation-one',
|
||||
webSocketFactory: restartedSockets.webSocketFactory,
|
||||
});
|
||||
|
||||
const firstStream = await first.openWorkspaceEvents({ workspaceId: 'workspace-one' });
|
||||
for await (const _event of firstStream.events) {
|
||||
// Empty stream.
|
||||
}
|
||||
const restartedStream = await restarted.openWorkspaceEvents({
|
||||
workspaceId: 'workspace-one',
|
||||
});
|
||||
for await (const _event of restartedStream.events) {
|
||||
// Empty stream.
|
||||
}
|
||||
|
||||
const firstBody = JSON.parse(String(firstFetch.mock.calls[0][1]?.body));
|
||||
const restartedBody = JSON.parse(String(secondFetch.mock.calls[0][1]?.body));
|
||||
expect(firstBody.client_session_id).toBe(restartedBody.client_session_id);
|
||||
expect(firstBody.client_session_id).toMatch(/^design-stream-[a-f0-9]{64}$/);
|
||||
});
|
||||
});
|
||||
|
||||
Reference in New Issue
Block a user