Files
makelore/tests/unit/works-square-design-workspace.test.ts
brother7 518d7ab4cd feat: 对接设计 Agent Gateway WebSocket
需求:服务端统一 Agent Gateway 将设计任务状态流切换为 WebSocket,客户端需要实时展示生成任务并支持断线恢复。

实现:Electron Main 管理 Session、一次性 Ticket、WebSocket 心跳与游标续传,按关闭码回收会话;Renderer 继续通过本机 Host API 的 SSE 投影接收任务事件,并保留 REST 降级同步。

验证:typecheck、变更文件 ESLint、37 个聚焦测试及 build:vite 通过。
2026-08-02 20:15:37 +08:00

611 lines
21 KiB
TypeScript

import { beforeEach, describe, expect, it, vi } from 'vitest';
import { WorksSquareDesignWorkspace } from '@electron/image-workspace/works-square-workspace';
import { getValidWorksSquareAccessToken } from '@electron/services/works-square-session';
vi.mock('@electron/services/works-square-session', () => ({
getValidWorksSquareAccessToken: vi.fn(),
}));
const getTokenMock = vi.mocked(getValidWorksSquareAccessToken);
const serverWorkspace = {
workspace_id: 'workspace-one',
title: '海洋公益海报',
turn_revision: 1,
view_revision: 2,
phase: 'awaiting_confirmation',
brief: {
version: 1,
status: 'ready',
medium: 'image',
summary: '保护海洋的竖版公益海报',
ready: true,
missing_decision: null,
},
messages: [{
role: 'assistant',
kind: 'confirmation',
text: '方向已经明确,是否开始生成?',
quick_replies: ['确认生成'],
generation_quote: {
quote_id: 'quote-one',
status: 'active',
medium: 'image',
brief_version: 1,
brief_summary: '保护海洋的竖版公益海报',
quoted_design_points: 1,
expires_at: '2026-07-31T11:00:00Z',
},
turn_revision: 1,
created_at: '2026-07-31T10:00:00Z',
}],
updated_at: '2026-07-31T10:00:00Z',
};
function jsonResponse(payload: unknown, status = 200): Response {
return new Response(JSON.stringify(payload), {
status,
headers: { 'Content-Type': 'application/json' },
});
}
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();
getTokenMock.mockResolvedValue('access-one');
});
it('maps the server snake_case Workspace contract into the shared client model', async () => {
const fetchMock = vi.fn<typeof fetch>()
.mockResolvedValueOnce(jsonResponse({
conversation: true,
generation: true,
image: true,
video: false,
}))
.mockResolvedValueOnce(jsonResponse([serverWorkspace]));
const adapter = new WorksSquareDesignWorkspace({
apiBaseUrl: 'https://square.example',
fetchImpl: fetchMock,
});
await expect(adapter.bootstrap()).resolves.toMatchObject({
capabilities: { conversation: true, image: true, video: false },
workspaces: [{
workspaceId: 'workspace-one',
title: '海洋公益海报',
turnRevision: 1,
viewRevision: 2,
brief: { missingDecision: null },
}],
});
expect(fetchMock).toHaveBeenCalledTimes(2);
expect(fetchMock.mock.calls.every(([, init]) => (
(init?.headers as Record<string, string>).Authorization === 'Bearer access-one'
))).toBe(true);
});
it('keeps task creation behind structured Quote confirmation', async () => {
const fetchMock = vi.fn<typeof fetch>().mockResolvedValue(jsonResponse(serverWorkspace));
const adapter = new WorksSquareDesignWorkspace({
apiBaseUrl: 'https://square.example/',
fetchImpl: fetchMock,
});
const workspace = await adapter.confirmGeneration({
workspaceId: 'workspace-one',
clientTurnId: 'turn-two',
expectedTurnRevision: 1,
quoteId: 'quote-one',
});
expect(workspace.workspaceId).toBe('workspace-one');
expect(fetchMock).toHaveBeenCalledWith(
'https://square.example/api/design/workspaces/workspace-one/turns',
expect.objectContaining({
method: 'POST',
body: JSON.stringify({
client_turn_id: 'turn-two',
expected_turn_revision: 1,
message: '确认生成',
attachment_asset_ids: [],
action: { type: 'confirm_generation', quote_id: 'quote-one' },
}),
}),
);
});
it('maps stable generation tasks and private asset relay paths', async () => {
const fetchMock = vi.fn<typeof fetch>().mockResolvedValue(jsonResponse([{
task_id: 'task-one',
workspace_id: 'workspace-one',
medium: 'image',
status: 'succeeded',
brief_version: 1,
brief_summary: '海洋公益海报',
quote_id: 'quote-one',
quoted_design_points: 1,
failure_code: null,
result_assets: [{
asset_id: 'asset-one',
media_type: 'image',
mime_type: 'image/png',
width: 1024,
height: 1280,
duration_milliseconds: null,
created_at: '2026-07-31T10:01:00Z',
}],
created_at: '2026-07-31T10:00:00Z',
updated_at: '2026-07-31T10:01:00Z',
}]));
const adapter = new WorksSquareDesignWorkspace({
apiBaseUrl: 'https://square.example',
fetchImpl: fetchMock,
});
await expect(adapter.listTasks('workspace-one')).resolves.toMatchObject([{
taskId: 'task-one',
status: 'succeeded',
resultAssets: [{
assetId: 'asset-one',
contentPath: '/api/works/image-workspace/workspaces/workspace-one/assets/asset-one/content',
}],
}]);
});
it('refreshes the Main-owned session once after an upstream 401 and preserves Range', async () => {
getTokenMock
.mockResolvedValueOnce('expired-token')
.mockResolvedValueOnce('fresh-token');
const fetchMock = vi.fn<typeof fetch>()
.mockResolvedValueOnce(new Response(null, { status: 401 }))
.mockResolvedValueOnce(new Response('partial', {
status: 206,
headers: {
'Content-Type': 'video/mp4',
'Content-Range': 'bytes 0-6/100',
},
}));
const adapter = new WorksSquareDesignWorkspace({
apiBaseUrl: 'https://square.example',
fetchImpl: fetchMock,
});
const response = await adapter.openAssetContent(
'workspace-one',
'asset-one',
'bytes=0-6',
);
expect(response.status).toBe(206);
expect(getTokenMock).toHaveBeenNthCalledWith(2, {
fetchImpl: fetchMock,
forceRefresh: true,
});
expect(fetchMock).toHaveBeenNthCalledWith(
2,
expect.stringContaining('/assets/asset-one/content'),
expect.objectContaining({
headers: {
Range: 'bytes=0-6',
Authorization: 'Bearer fresh-token',
},
}),
);
});
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}$/);
});
});