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() .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).Authorization === 'Bearer access-one' ))).toBe(true); }); it('submits a conversation turn through the persistent Agent Gateway Session', async () => { const fetchMock = vi.fn() .mockResolvedValueOnce(jsonResponse({ session_id: 'session-one', status: 'active', }, 201)) .mockResolvedValueOnce(jsonResponse({ run_id: 'run-one', status: 'queued', error: null, }, 202)) .mockResolvedValueOnce(jsonResponse({ run_id: 'run-one', status: 'succeeded', error: null, })) .mockResolvedValueOnce(jsonResponse(serverWorkspace)); const adapter = new WorksSquareDesignWorkspace({ apiBaseUrl: 'https://square.example', fetchImpl: fetchMock, clientInstanceId: 'installation-one', }); await expect(adapter.submitMessage({ workspaceId: 'workspace-one', clientTurnId: 'turn-two', expectedTurnRevision: 1, message: '做一张保护海洋的公益海报', attachmentAssetIds: ['asset-reference'], })).resolves.toMatchObject({ workspaceId: 'workspace-one', turnRevision: 1, }); expect(fetchMock).toHaveBeenNthCalledWith( 2, 'https://square.example/api/agents/sessions/session-one/commands', expect.objectContaining({ method: 'POST', body: JSON.stringify({ client_command_id: 'turn-two', name: 'turn.submit', input: { expected_turn_revision: 1, message: '做一张保护海洋的公益海报', attachment_asset_ids: ['asset-reference'], action: null, }, }), }), ); expect(fetchMock).toHaveBeenNthCalledWith( 3, 'https://square.example/api/agents/sessions/session-one/runs/run-one', expect.objectContaining({ headers: expect.any(Object) }), ); expect(fetchMock).toHaveBeenNthCalledWith( 4, 'https://square.example/api/design/workspaces/workspace-one', expect.objectContaining({ headers: expect.any(Object) }), ); }); it('maps an invalid Runtime command to a user input error', async () => { const fetchMock = vi.fn() .mockResolvedValueOnce(jsonResponse({ session_id: 'session-one', status: 'active', }, 201)) .mockResolvedValueOnce(jsonResponse({ run_id: 'run-invalid', status: 'queued', error: null, }, 202)) .mockResolvedValueOnce(jsonResponse({ run_id: 'run-invalid', status: 'failed', error: { code: 'agent_command_invalid', message: 'private validation detail', retryable: false, }, })); const adapter = new WorksSquareDesignWorkspace({ apiBaseUrl: 'https://square.example', fetchImpl: fetchMock, }); await expect(adapter.submitMessage({ workspaceId: 'workspace-one', clientTurnId: 'turn-invalid', expectedTurnRevision: 1, message: 'invalid', })).rejects.toMatchObject({ status: 422, code: 'agent_command_invalid', message: '设计请求内容无效,请检查后重试', }); }); it('keeps task creation behind structured Quote confirmation', async () => { const fetchMock = vi.fn() .mockResolvedValueOnce(jsonResponse({ session_id: 'session-one', status: 'active', }, 201)) .mockResolvedValueOnce(jsonResponse({ run_id: 'run-confirm', status: 'queued', error: null, }, 202)) .mockResolvedValueOnce(jsonResponse({ run_id: 'run-confirm', status: 'succeeded', error: null, })) .mockResolvedValueOnce(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).toHaveBeenNthCalledWith( 2, 'https://square.example/api/agents/sessions/session-one/commands', expect.objectContaining({ method: 'POST', body: JSON.stringify({ client_command_id: 'turn-two', name: 'turn.submit', input: { expected_turn_revision: 1, message: '确认生成', attachment_asset_ids: [], action: { type: 'confirm_generation', quote_id: 'quote-one' }, }, }), }), ); expect(fetchMock).toHaveBeenNthCalledWith( 3, 'https://square.example/api/agents/sessions/session-one/runs/run-confirm', expect.objectContaining({ headers: expect.any(Object) }), ); expect(fetchMock).toHaveBeenNthCalledWith( 4, 'https://square.example/api/design/workspaces/workspace-one', expect.objectContaining({ headers: expect.any(Object) }), ); }); it('maps stable generation tasks and private asset relay paths', async () => { const fetchMock = vi.fn().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() .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: serverWorkspace, generation_tasks: [snapshotTask], }, }; const assistantDeltaEvent = { session_id: 'session-one', sequence: 2, runtime: 'design', type: 'design.assistant.delta', schema_version: 1, payload: { workspace_id: 'workspace-one', client_turn_id: 'turn-two', turn_revision: 2, chunk_index: 0, delta: '方向已经明确', }, }; const malformedDeltaEvent = { ...assistantDeltaEvent, sequence: 99, payload: { ...assistantDeltaEvent.payload, chunk_index: -1, }, }; 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() .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: malformedDeltaEvent }, { type: 'event', event: assistantDeltaEvent }, { 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, workspace: expect.objectContaining({ workspaceId: 'workspace-one', title: serverWorkspace.title, messages: [expect.objectContaining({ text: serverWorkspace.messages[0].text, })], }), generationTasks: [expect.objectContaining({ taskId: 'task-snapshot', medium: 'image', status: 'running', })], }, { id: 'session-one:2', type: 'design.assistant.delta', workspaceId: 'workspace-one', clientTurnId: 'turn-two', turnRevision: 2, chunkIndex: 0, delta: '方向已经明确', }, { 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() .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() .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() .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() .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() .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}$/); }); });