diff --git a/electron/image-workspace/works-square-workspace.ts b/electron/image-workspace/works-square-workspace.ts index 91a5129..46f91b1 100644 --- a/electron/image-workspace/works-square-workspace.ts +++ b/electron/image-workspace/works-square-workspace.ts @@ -128,6 +128,7 @@ type ServerAgentEvent = { sequence: number; runtime: string; type: string; + run_id?: unknown; schema_version: number; payload: unknown; }; @@ -188,6 +189,8 @@ type TaskEventQueue = { fail(error: unknown): void; }; +type AgentRunEventWaiter = (run: ServerAgentRun) => void; + const AGENT_WEBSOCKET_OPEN = 1; const AGENT_WEBSOCKET_PING_INTERVAL_MS = 20_000; const AGENT_RUN_INITIAL_POLL_INTERVAL_MS = 1_000; @@ -438,6 +441,39 @@ function agentEventFromWebSocketFrame(data: unknown): unknown | null { } } +function normalizeAgentRunEvent(value: unknown, sessionId: string): ServerAgentRun | null { + if (!value || typeof value !== 'object' || Array.isArray(value)) return null; + const event = value as ServerAgentEvent; + if (event.session_id !== sessionId + || !Number.isInteger(event.sequence) + || event.sequence < 1 + || event.runtime !== 'design' + || event.schema_version !== 1 + || typeof event.run_id !== 'string' + || event.run_id.length === 0) { + return null; + } + const status = event.type === 'run.completed' + ? 'succeeded' + : event.type === 'run.failed' + ? 'failed' + : event.type === 'run.cancelled' + ? 'cancelled' + : null; + if (!status) return null; + const payload = event.payload && typeof event.payload === 'object' && !Array.isArray(event.payload) + ? event.payload as Record + : null; + const rawError = payload?.error; + const error = rawError && typeof rawError === 'object' && !Array.isArray(rawError) + && typeof (rawError as Record).code === 'string' + && typeof (rawError as Record).message === 'string' + && typeof (rawError as Record).retryable === 'boolean' + ? rawError as ServerAgentCommandError + : null; + return { run_id: event.run_id, status, error }; +} + function createTaskEventQueue(): TaskEventQueue { const queued: DesignWorkspaceEvent[] = []; const waiters: Array<() => void> = []; @@ -595,6 +631,10 @@ export class WorksSquareDesignWorkspace implements DesignWorkspaceModule { private readonly eventSessions = new Map>(); private readonly eventSessionClientIds = new Map(); private readonly eventSubscriptionClosers = new Map void>>(); + private readonly activeRunEventStreams = new Map(); + private readonly terminalAgentRuns = new Map(); + private readonly agentRunEventWaiters = new Map>(); + private readonly runStreamEndWaiters = new Map void>>(); private eventSessionsEnabled = true; constructor(options: WorksSquareDesignWorkspaceOptions = {}) { @@ -738,6 +778,16 @@ export class WorksSquareDesignWorkspace implements DesignWorkspaceModule { private async waitForAgentRun(sessionId: string, runId: string): Promise { const deadline = Date.now() + AGENT_RUN_TIMEOUT_MS; + const streamedRun = await this.waitForAgentRunEvent(sessionId, runId, deadline); + if (streamedRun) return streamedRun; + return this.pollAgentRun(sessionId, runId, deadline); + } + + private async pollAgentRun( + sessionId: string, + runId: string, + deadline: number, + ): Promise { let pollInterval = AGENT_RUN_INITIAL_POLL_INTERVAL_MS; while (true) { const run = await this.requestJson( @@ -758,6 +808,62 @@ export class WorksSquareDesignWorkspace implements DesignWorkspaceModule { } } + private waitForAgentRunEvent( + sessionId: string, + runId: string, + deadline: number, + ): Promise { + const key = `${sessionId}:${runId}`; + const cached = this.terminalAgentRuns.get(key); + if (cached) { + this.terminalAgentRuns.delete(key); + return Promise.resolve(cached); + } + if (!this.activeRunEventStreams.has(sessionId)) return Promise.resolve(null); + + return new Promise((resolve, reject) => { + let settled = false; + let timeout: ReturnType | null = null; + const cleanup = () => { + if (timeout !== null) clearTimeout(timeout); + const runWaiters = this.agentRunEventWaiters.get(key); + runWaiters?.delete(onRun); + if (runWaiters?.size === 0) this.agentRunEventWaiters.delete(key); + const streamWaiters = this.runStreamEndWaiters.get(sessionId); + streamWaiters?.delete(onStreamEnd); + if (streamWaiters?.size === 0) this.runStreamEndWaiters.delete(sessionId); + }; + const finish = (run: ServerAgentRun | null, error?: unknown) => { + if (settled) return; + settled = true; + cleanup(); + if (error) reject(error); + else resolve(run); + }; + const onRun: AgentRunEventWaiter = (run) => finish(run); + const onStreamEnd = () => finish(null); + const runWaiters = this.agentRunEventWaiters.get(key) ?? new Set(); + runWaiters.add(onRun); + this.agentRunEventWaiters.set(key, runWaiters); + const streamWaiters = this.runStreamEndWaiters.get(sessionId) ?? new Set<() => void>(); + streamWaiters.add(onStreamEnd); + this.runStreamEndWaiters.set(sessionId, streamWaiters); + timeout = setTimeout(() => finish(null, new DesignWorkspaceModuleError( + 504, + 'design_agent_run_timeout', + '设计 Agent 响应超时,请稍后重试', + )), Math.max(0, deadline - Date.now())); + + const racedRun = this.terminalAgentRuns.get(key); + if (racedRun) { + this.terminalAgentRuns.delete(key); + finish(racedRun); + } else if (!this.activeRunEventStreams.has(sessionId)) { + finish(null); + } + }); + } + async listTasks(workspaceId: string): Promise { const tasks = await this.requestJson( `/api/design/workspaces/${encodeURIComponent(workspaceId)}/generation-tasks?limit=100&offset=0`, @@ -841,6 +947,7 @@ export class WorksSquareDesignWorkspace implements DesignWorkspaceModule { if (heartbeat !== null) clearInterval(heartbeat); heartbeat = null; unregister(); + if (didOpen) this.unregisterRunEventStream(session.session_id); try { await connection.dispose?.(); } catch { @@ -883,6 +990,7 @@ export class WorksSquareDesignWorkspace implements DesignWorkspaceModule { socket.onopen = () => { if (ending) return; didOpen = true; + this.registerRunEventStream(session.session_id); heartbeat = setInterval(() => { if (ending || socket.readyState !== AGENT_WEBSOCKET_OPEN) return; try { @@ -899,6 +1007,8 @@ export class WorksSquareDesignWorkspace implements DesignWorkspaceModule { socket.onmessage = ({ data }) => { if (ending) return; const agentEvent = agentEventFromWebSocketFrame(data); + const run = normalizeAgentRunEvent(agentEvent, session.session_id); + if (run) this.publishAgentRun(session.session_id, run); const event = normalizeWorkspaceEvent( agentEvent, session.session_id, @@ -1101,6 +1211,40 @@ export class WorksSquareDesignWorkspace implements DesignWorkspaceModule { }; } + private registerRunEventStream(sessionId: string): void { + this.activeRunEventStreams.set( + sessionId, + (this.activeRunEventStreams.get(sessionId) ?? 0) + 1, + ); + } + + private unregisterRunEventStream(sessionId: string): void { + const remaining = (this.activeRunEventStreams.get(sessionId) ?? 1) - 1; + if (remaining > 0) { + this.activeRunEventStreams.set(sessionId, remaining); + return; + } + this.activeRunEventStreams.delete(sessionId); + const waiters = this.runStreamEndWaiters.get(sessionId); + this.runStreamEndWaiters.delete(sessionId); + for (const resolve of waiters ?? []) resolve(); + } + + private publishAgentRun(sessionId: string, run: ServerAgentRun): void { + const key = `${sessionId}:${run.run_id}`; + const waiters = this.agentRunEventWaiters.get(key); + if (!waiters?.size) { + this.terminalAgentRuns.set(key, run); + if (this.terminalAgentRuns.size > 100) { + const oldest = this.terminalAgentRuns.keys().next().value; + if (oldest) this.terminalAgentRuns.delete(oldest); + } + return; + } + this.agentRunEventWaiters.delete(key); + for (const resolve of waiters) resolve(run); + } + private async closeEventSession(sessionId: string): Promise { try { await this.requestJson( diff --git a/tests/unit/works-square-design-workspace.test.ts b/tests/unit/works-square-design-workspace.test.ts index 8697cdf..85aa5a0 100644 --- a/tests/unit/works-square-design-workspace.test.ts +++ b/tests/unit/works-square-design-workspace.test.ts @@ -75,6 +75,10 @@ class MockAgentWebSocket { this.sent.push(data); } + emitFrame(frame: unknown): void { + this.onmessage?.({ data: JSON.stringify(frame) }); + } + close(code = 1000, reason = ''): void { this.emitClose(code, reason); } @@ -210,6 +214,166 @@ describe('Works Square AI design adapter', () => { ); }); + it('completes a design turn from the connected WebSocket without polling the run endpoint', async () => { + const { sockets, webSocketFactory } = scriptedSockets([{ open: true }]); + const fetchMock = vi.fn(async (input) => { + const url = String(input); + if (url.endsWith('/api/agents/sessions')) { + return jsonResponse({ session_id: 'session-live', status: 'active' }, 201); + } + if (url.endsWith('/stream-tickets')) { + return jsonResponse({ + ticket: 'ticket-live', + transport: 'websocket', + stream_url: '/api/agents/sessions/session-live/ws?ticket=ticket-live', + expires_at: '2026-08-03T03:00:00Z', + }); + } + if (url.endsWith('/api/agents/sessions/session-live/commands')) { + queueMicrotask(() => sockets[0]?.emitFrame({ + type: 'event', + event: { + session_id: 'session-live', + sequence: 6, + runtime: 'design', + type: 'run.completed', + command_id: 'command-live', + run_id: 'run-live', + client_command_id: 'turn-live', + schema_version: 1, + terminal: true, + occurred_at: '2026-08-03T02:00:05Z', + payload: { status: 'succeeded' }, + }, + })); + return jsonResponse({ + command_id: 'command-live', + run_id: 'run-live', + status: 'accepted', + }, 202); + } + if (url.endsWith('/api/design/workspaces/workspace-one')) { + return jsonResponse(serverWorkspace); + } + throw new Error(`Unexpected request: ${url}`); + }); + const adapter = new WorksSquareDesignWorkspace({ + apiBaseUrl: 'https://square.example', + fetchImpl: fetchMock, + webSocketFactory, + }); + const subscription = await adapter.openWorkspaceEvents({ workspaceId: 'workspace-one' }); + + try { + await expect(adapter.submitMessage({ + workspaceId: 'workspace-one', + clientTurnId: 'turn-live', + expectedTurnRevision: 1, + message: '做一张保护海洋的公益海报', + })).resolves.toMatchObject({ workspaceId: 'workspace-one' }); + expect(fetchMock.mock.calls.filter(([url]) => String(url).includes('/runs/'))).toHaveLength(0); + } finally { + subscription.close(); + } + }); + + it('maps a failed Run received from the connected WebSocket without polling', async () => { + const { sockets, webSocketFactory } = scriptedSockets([{ open: true }]); + const fetchMock = vi.fn(async (input) => { + const url = String(input); + if (url.endsWith('/api/agents/sessions')) { + return jsonResponse({ session_id: 'session-failed', status: 'active' }, 201); + } + if (url.endsWith('/stream-tickets')) { + return jsonResponse({ + stream_url: '/api/agents/sessions/session-failed/ws?ticket=ticket-failed', + }); + } + if (url.endsWith('/api/agents/sessions/session-failed/commands')) { + queueMicrotask(() => sockets[0]?.emitFrame({ + type: 'event', + event: { + session_id: 'session-failed', + sequence: 6, + runtime: 'design', + type: 'run.failed', + run_id: 'run-failed', + schema_version: 1, + payload: { + error: { + code: 'agent_command_invalid', + message: 'private validation detail', + retryable: false, + }, + }, + }, + })); + return jsonResponse({ run_id: 'run-failed', status: 'queued', error: null }, 202); + } + throw new Error(`Unexpected request: ${url}`); + }); + const adapter = new WorksSquareDesignWorkspace({ + apiBaseUrl: 'https://square.example', + fetchImpl: fetchMock, + webSocketFactory, + }); + const subscription = await adapter.openWorkspaceEvents({ workspaceId: 'workspace-one' }); + + try { + await expect(adapter.submitMessage({ + workspaceId: 'workspace-one', + clientTurnId: 'turn-failed', + expectedTurnRevision: 1, + message: 'invalid', + })).rejects.toMatchObject({ + status: 422, + code: 'agent_command_invalid', + message: '设计请求内容无效,请检查后重试', + }); + expect(fetchMock.mock.calls.filter(([url]) => String(url).includes('/runs/'))).toHaveLength(0); + } finally { + subscription.close(); + } + }); + + it('falls back to one Run request after the WebSocket disconnects', async () => { + const { webSocketFactory } = scriptedSockets([{ open: true, closeCode: 1006 }]); + const fetchMock = vi.fn() + .mockResolvedValueOnce(jsonResponse({ session_id: 'session-fallback', status: 'active' }, 201)) + .mockResolvedValueOnce(jsonResponse({ + stream_url: '/api/agents/sessions/session-fallback/ws?ticket=ticket-fallback', + })) + .mockResolvedValueOnce(jsonResponse({ + run_id: 'run-fallback', + status: 'queued', + error: null, + }, 202)) + .mockResolvedValueOnce(jsonResponse({ + run_id: 'run-fallback', + status: 'succeeded', + error: null, + })) + .mockResolvedValueOnce(jsonResponse(serverWorkspace)); + const adapter = new WorksSquareDesignWorkspace({ + apiBaseUrl: 'https://square.example', + fetchImpl: fetchMock, + webSocketFactory, + }); + const subscription = await adapter.openWorkspaceEvents({ workspaceId: 'workspace-one' }); + + try { + await expect(adapter.submitMessage({ + workspaceId: 'workspace-one', + clientTurnId: 'turn-fallback', + expectedTurnRevision: 1, + message: '断线后继续完成', + })).resolves.toMatchObject({ workspaceId: 'workspace-one' }); + expect(fetchMock.mock.calls.filter(([url]) => String(url).includes('/runs/'))).toHaveLength(1); + } finally { + subscription.close(); + } + }); + it('waits for a slow design run without flooding the run endpoint', async () => { vi.useFakeTimers(); vi.setSystemTime(new Date('2026-08-03T02:00:00Z'));