fix: 使用 WebSocket 完成设计 Agent Run

问题：客户端虽然保持 Agent Gateway WebSocket，但 Run 生命周期事件被过滤，提交后仍持续轮询 /runs 接口。

修复：按 Session 与 Run 关联 WebSocket 终态事件，处理事件抢跑和并发隔离；仅在连接缺失或断开时退化为低频 REST 查询。

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<string, unknown>
+    : null;
+  const rawError = payload?.error;
+  const error = rawError && typeof rawError === 'object' && !Array.isArray(rawError)
+    && typeof (rawError as Record<string, unknown>).code === 'string'
+    && typeof (rawError as Record<string, unknown>).message === 'string'
+    && typeof (rawError as Record<string, unknown>).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<string, Promise<ServerAgentSession>>();
   private readonly eventSessionClientIds = new Map<string, string>();
   private readonly eventSubscriptionClosers = new Map<string, Set<() => void>>();
+  private readonly activeRunEventStreams = new Map<string, number>();
+  private readonly terminalAgentRuns = new Map<string, ServerAgentRun>();
+  private readonly agentRunEventWaiters = new Map<string, Set<AgentRunEventWaiter>>();
+  private readonly runStreamEndWaiters = new Map<string, Set<() => void>>();
   private eventSessionsEnabled = true;
 
   constructor(options: WorksSquareDesignWorkspaceOptions = {}) {
@@ -738,6 +778,16 @@ export class WorksSquareDesignWorkspace implements DesignWorkspaceModule {
 
   private async waitForAgentRun(sessionId: string, runId: string): Promise<ServerAgentRun> {
     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<ServerAgentRun> {
     let pollInterval = AGENT_RUN_INITIAL_POLL_INTERVAL_MS;
     while (true) {
       const run = await this.requestJson<ServerAgentRun>(
@@ -758,6 +808,62 @@ export class WorksSquareDesignWorkspace implements DesignWorkspaceModule {
     }
   }
 
+  private waitForAgentRunEvent(
+    sessionId: string,
+    runId: string,
+    deadline: number,
+  ): Promise<ServerAgentRun | null> {
+    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<ServerAgentRun | null>((resolve, reject) => {
+      let settled = false;
+      let timeout: ReturnType<typeof setTimeout> | 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<AgentRunEventWaiter>();
+      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<DesignGenerationTask[]> {
     const tasks = await this.requestJson<ServerTask[]>(
       `/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<void> {
     try {
       await this.requestJson<ServerAgentSession>(

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<typeof fetch>(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<typeof fetch>(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<typeof fetch>()
+      .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'));