fix: 使用 WebSocket 完成设计 Agent Run

问题：客户端虽然保持 Agent Gateway WebSocket，但 Run 生命周期事件被过滤，提交后仍持续轮询 /runs 接口。修复：按 Session 与 Run 关联 WebSocket 终态事件，处理事件抢跑和并发隔离；仅在连接缺失或断开时退化为低频 REST 查询。
2026-08-03 12:02:48 +08:00
parent 20c567fe5f
commit 78e4de91e4
2 changed files with 308 additions and 0 deletions
--- 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<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>(