wangxuming/makelore
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

fix: 使用 WebSocket 完成设计 Agent Run

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

修复：按 Session 与 Run 关联 WebSocket 终态事件，处理事件抢跑和并发隔离；仅在连接缺失或断开时退化为低频 REST 查询。
This commit is contained in:
brother7
2026-08-03 12:02:48 +08:00
parent 20c567fe5f
commit 78e4de91e4
2 changed files with 308 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

144
electron/image-workspace/works-square-workspace.ts
View File

@@ -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>(

164
tests/unit/works-square-design-workspace.test.ts
View File

@@ -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'));