diff --git a/README.md b/README.md index 42fb5bf..c4a969f 100644 --- a/README.md +++ b/README.md @@ -16,7 +16,7 @@ Makelore 是一个面向软件与视觉创作的 AI 桌面工作台。当前版 - AI 编程运行时:Electron Main 管理项目内声明的 `opencode-ai` 依赖,Renderer 不直接启动或调用运行时。 - 共享开发浏览器:AI 编程右侧提供项目级浏览器,用户与 Agent 查看并调试同一实时页面、Console 和 Network,支持本地与公网开发地址。 - 后端边界:Renderer 通过 Main 所有的 Host API 访问认证、模型、同步、更新、语音、图像与运行时能力。 -- AI 绘画:每个设计项目固定一个设计 Agent。Agent 先通过持续对话收敛 Brief,用户确认服务端 Quote 后才创建图片或视频任务;任务进度与对话互不阻塞,并在统一列表中展示。生产环境使用 Works Square 云端 Workspace 契约,上游不可用时明确报错。 +- AI 绘画:每个设计项目固定一个设计 Agent。Agent 先通过持续对话收敛 Brief,用户确认服务端 Quote 后才创建图片或视频任务;任务进度通过 Agent Gateway 实时推送并在统一列表中展示,断流时使用低频 REST 同步。生产环境使用 Works Square 云端 Workspace 契约,上游不可用时明确报错。 - 视觉系统:单一浅色主题,品牌蓝 `#3A5578`、星火橙 `#F26A3D`、白色画布与低饱和蓝灰层级。 - 品牌资产:生产 SVG、PNG、应用图标、托盘图标与安装器视觉位于 [`resources/brand/`](resources/brand/README.md) 和 `resources/icons/`。 @@ -80,7 +80,7 @@ pnpm run package:linux - Renderer 不直接调用 Electron IPC 或本地运行时 HTTP 地址。 - Electron Main 负责认证、秘密存储、运行时生命周期、代理、同步和系统集成。 - AI 编程项目配置以项目内 `.niancode/project.json` 为准;项目文件和会话保持本地。 -- AI 绘画 Renderer 只调用 Main-owned Host API;Main 负责 Works Square Token 刷新、snake_case 契约映射、结构化错误和私有媒体 Range 转发。 +- AI 绘画 Renderer 只调用 Main-owned Host API;Main 负责 Works Square Token 刷新、可跨异常重启幂等复用的 Agent Session、单次 WebSocket ticket、断点续传与事件契约映射,并通过本机 Host API 的 SSE 投影同步任务状态;注销/退出时回收 Session,远端 Token 与 ticket 不进入 Renderer。 - AI 绘画使用独立的云端 Workspace 边界,不回退到 AI 编程项目数据,也不向 Renderer 暴露 Provider、模型、Prompt、存储 URI 或远端登录 Token。 - AI 编程的 Agent 配置是项目所有的;稳定 id 用于保持会话兼容,显示名称可以修改。AI 绘画的设计 Agent 是固定产品能力,不作为用户可增删的项目实体。 diff --git a/electron/api/routes/auth.ts b/electron/api/routes/auth.ts index 55ca5df..8bce6b7 100644 --- a/electron/api/routes/auth.ts +++ b/electron/api/routes/auth.ts @@ -365,6 +365,12 @@ async function handleSessionSync(req: IncomingMessage, res: ServerResponse): Pro async function clearManagedWorksSquareRuntime(ctx: HostApiContext): Promise { const errors: string[] = []; + try { + await ctx.imageWorkspace?.closeEventSessions?.(); + } catch (error) { + errors.push(error instanceof Error ? error.message : String(error)); + } + try { await ctx.opencodeManager.stop(); } catch (error) { diff --git a/electron/api/routes/image-workspace.ts b/electron/api/routes/image-workspace.ts index d8332ae..4d5b1e2 100644 --- a/electron/api/routes/image-workspace.ts +++ b/electron/api/routes/image-workspace.ts @@ -12,7 +12,12 @@ import { } from '../../../shared/image-workspace'; import { DesignWorkspaceModuleError } from '../../image-workspace/module'; import type { HostApiContext } from '../context'; -import { parseJsonBody, sendJson } from '../route-utils'; +import { + flushStreamingHeaders, + parseJsonBody, + sendJson, + writeStreamingChunk, +} from '../route-utils'; function decodedSegments(pathname: string): string[] | null { const suffix = pathname.slice(IMAGE_WORKSPACE_API_PATH.length).replace(/^\/+/, ''); @@ -105,6 +110,55 @@ async function relayAssetContent( await pipeline(Readable.fromWeb(content.body), res); } +async function relayWorkspaceEvents( + req: IncomingMessage, + res: ServerResponse, + ctx: HostApiContext, + workspaceId: string, +): Promise { + if (!ctx.imageWorkspace?.openWorkspaceEvents) { + throw new DesignWorkspaceModuleError( + 501, + 'DESIGN_EVENT_STREAM_UNAVAILABLE', + 'AI 设计任务实时状态暂时不可用', + ); + } + const header = req.headers['last-event-id']; + const afterEventId = Array.isArray(header) ? header[0] : header; + const subscription = await ctx.imageWorkspace.openWorkspaceEvents({ + workspaceId, + ...(afterEventId ? { afterEventId } : {}), + }); + let closed = false; + const close = () => { + if (closed) return; + closed = true; + subscription.close(); + }; + res.once('close', close); + try { + res.statusCode = 200; + res.setHeader('Content-Type', 'text/event-stream; charset=utf-8'); + res.setHeader('Cache-Control', 'no-cache, no-transform'); + res.setHeader('Connection', 'keep-alive'); + res.setHeader('X-Accel-Buffering', 'no'); + flushStreamingHeaders(res); + if (!await writeStreamingChunk(res, ': connected\n\n')) return; + for await (const event of subscription.events) { + if (!await writeStreamingChunk( + res, + `id: ${event.id}\nevent: ${event.type}\ndata: ${JSON.stringify(event)}\n\n`, + )) { + return; + } + } + if (!res.writableEnded) res.end(); + } finally { + res.off('close', close); + close(); + } +} + export async function handleImageWorkspaceRoutes( req: IncomingMessage, res: ServerResponse, @@ -199,6 +253,14 @@ export async function handleImageWorkspaceRoutes( return true; } + if (segments.length === 3 + && segments[0] === 'workspaces' + && segments[2] === 'events' + && req.method === 'GET') { + await relayWorkspaceEvents(req, res, ctx, segments[1]); + return true; + } + if (segments.length === 5 && segments[0] === 'workspaces' && segments[2] === 'quotes' diff --git a/electron/image-workspace/module.ts b/electron/image-workspace/module.ts index e180a95..fbc2c00 100644 --- a/electron/image-workspace/module.ts +++ b/electron/image-workspace/module.ts @@ -7,8 +7,19 @@ import type { DesignSubmitMessageInput, DesignWorkspace, DesignWorkspaceBootstrap, + DesignWorkspaceEvent, } from '../../shared/image-workspace'; +export type DesignWorkspaceEventSubscription = { + events: AsyncIterable; + close(): void; +}; + +export type DesignWorkspaceEventSubscriptionInput = { + workspaceId: string; + afterEventId?: string; +}; + export class DesignWorkspaceModuleError extends Error { readonly status: number; readonly code: string; @@ -30,6 +41,10 @@ export interface DesignWorkspaceModule { submitMessage(input: DesignSubmitMessageInput): Promise; confirmGeneration(input: DesignConfirmGenerationInput): Promise; listTasks(workspaceId: string): Promise; + openWorkspaceEvents?( + input: DesignWorkspaceEventSubscriptionInput, + ): Promise; + closeEventSessions?(): Promise; openAssetContent(workspaceId: string, assetId: string, range?: string): Promise; reset?(): Promise; } diff --git a/electron/image-workspace/works-square-workspace.ts b/electron/image-workspace/works-square-workspace.ts index 47d5321..5660320 100644 --- a/electron/image-workspace/works-square-workspace.ts +++ b/electron/image-workspace/works-square-workspace.ts @@ -6,18 +6,28 @@ import type { DesignCreateWorkspaceInput, DesignGenerationQuote, DesignGenerationTask, + DesignGenerationTasksSnapshotEvent, + DesignGenerationTaskUpdatedEvent, DesignMessage, DesignRenameWorkspaceInput, DesignSubmitMessageInput, DesignWorkspace, DesignWorkspaceBootstrap, + DesignWorkspaceEvent, DesignWorkspaceSummary, } from '../../shared/image-workspace'; +import { createHash } from 'node:crypto'; +import WebSocket from 'ws'; import { designAssetContentPath } from '../../shared/image-workspace'; import { WORKS_SQUARE_CONFIG } from '../api/works-config'; import { proxyAwareFetch } from '../utils/proxy-fetch'; import { getValidWorksSquareAccessToken } from '../services/works-square-session'; -import { DesignWorkspaceModuleError, type DesignWorkspaceModule } from './module'; +import { + DesignWorkspaceModuleError, + type DesignWorkspaceEventSubscription, + type DesignWorkspaceEventSubscriptionInput, + type DesignWorkspaceModule, +} from './module'; type ServerBrief = { version: number; @@ -95,8 +105,61 @@ type ServerErrorDetail = { type WorksSquareDesignWorkspaceOptions = { apiBaseUrl?: string; fetchImpl?: typeof fetch; + clientInstanceId?: string; + webSocketFactory?: AgentWebSocketFactory; + eventSessionClientIdStore?: { + getOrCreate(workspaceId: string): Promise; + rotate(workspaceId: string): Promise; + }; }; +type ServerAgentSession = { + session_id: string; + status: 'active' | 'closed'; +}; + +type ServerAgentStreamTicket = { + stream_url: string; +}; + +type ServerAgentEvent = { + session_id: string; + sequence: number; + runtime: string; + type: string; + schema_version: number; + payload: unknown; +}; + +type AgentWebSocket = { + readyState: number; + onopen: (() => void) | null; + onmessage: ((event: { data: unknown }) => void) | null; + onerror: ((event: unknown) => void) | null; + onclose: ((event: { code: number; reason: string }) => void) | null; + send(data: string): void; + close(code?: number, reason?: string): void; +}; + +type AgentWebSocketConnection = { + socket: AgentWebSocket; + dispose?: () => void | Promise; +}; + +type AgentWebSocketFactory = ( + url: string, +) => AgentWebSocketConnection | Promise; + +type TaskEventQueue = { + events: AsyncIterable; + push(event: DesignWorkspaceEvent): void; + finish(): void; + fail(error: unknown): void; +}; + +const AGENT_WEBSOCKET_OPEN = 1; +const AGENT_WEBSOCKET_PING_INTERVAL_MS = 20_000; + function mapBrief(brief: ServerBrief): DesignBrief { return { version: brief.version, @@ -176,6 +239,198 @@ function mapTask(task: ServerTask): DesignGenerationTask { }; } +function createClientId(prefix: string): string { + const id = globalThis.crypto?.randomUUID?.() + ?? `${Date.now()}-${Math.random().toString(16).slice(2)}`; + return `${prefix}-${id}`; +} + +function stableWorkspaceSessionClientId(clientInstanceId: string, workspaceId: string): string { + const digest = createHash('sha256') + .update(`${clientInstanceId}\0${workspaceId}`) + .digest('hex'); + return `design-stream-${digest}`; +} + +function isServerTask(value: unknown): value is ServerTask { + if (!value || typeof value !== 'object' || Array.isArray(value)) return false; + const task = value as Record; + return typeof task.task_id === 'string' + && typeof task.workspace_id === 'string' + && (task.medium === 'image' || task.medium === 'video') + && ['queued', 'running', 'succeeded', 'failed', 'cancelled'].includes(String(task.status)) + && typeof task.brief_version === 'number' + && typeof task.brief_summary === 'string' + && Array.isArray(task.result_assets) + && typeof task.created_at === 'string' + && typeof task.updated_at === 'string'; +} + +function normalizeWorkspaceEvent( + value: unknown, + sessionId: string, + workspaceId: string, +): DesignWorkspaceEvent | null { + try { + 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 + || !event.payload + || typeof event.payload !== 'object' + || Array.isArray(event.payload)) { + return null; + } + const payload = event.payload as Record; + if (event.type === 'design.generation_task.updated') { + if (payload.workspace_id !== workspaceId + || !Number.isInteger(payload.workspace_view_revision) + || Number(payload.workspace_view_revision) < 1 + || !isServerTask(payload.generation_task) + || payload.generation_task.workspace_id !== workspaceId) { + return null; + } + return { + id: `${sessionId}:${event.sequence}`, + type: 'design.generation_task.updated', + workspaceId, + workspaceViewRevision: Number(payload.workspace_view_revision), + generationTask: mapTask(payload.generation_task), + } satisfies DesignGenerationTaskUpdatedEvent; + } + + if (event.type !== 'design.workspace.updated' + || !payload.workspace + || typeof payload.workspace !== 'object' + || Array.isArray(payload.workspace) + || !Array.isArray(payload.generation_tasks)) { + return null; + } + const workspace = payload.workspace as Record; + if (workspace.workspace_id !== workspaceId + || !Number.isInteger(workspace.view_revision) + || Number(workspace.view_revision) < 0 + || !payload.generation_tasks.every( + (task) => isServerTask(task) && task.workspace_id === workspaceId, + )) { + return null; + } + return { + id: `${sessionId}:${event.sequence}`, + type: 'design.generation_tasks.snapshot', + workspaceId, + workspaceViewRevision: Number(workspace.view_revision), + generationTasks: payload.generation_tasks.map((task) => mapTask(task as ServerTask)), + } satisfies DesignGenerationTasksSnapshotEvent; + } catch { + return null; + } +} + +function agentEventFromWebSocketFrame(data: unknown): unknown | null { + if (typeof data !== 'string') return null; + try { + const frame = JSON.parse(data) as Record; + return frame + && typeof frame === 'object' + && !Array.isArray(frame) + && frame.type === 'event' + ? frame.event ?? null + : null; + } catch { + return null; + } +} + +function createTaskEventQueue(): TaskEventQueue { + const queued: DesignWorkspaceEvent[] = []; + const waiters: Array<() => void> = []; + let finished = false; + let failed = false; + let failure: unknown; + const wake = () => { + for (const resolve of waiters.splice(0)) resolve(); + }; + + return { + events: { + async *[Symbol.asyncIterator]() { + while (true) { + const event = queued.shift(); + if (event) { + yield event; + continue; + } + if (failed) throw failure; + if (finished) return; + await new Promise((resolve) => waiters.push(resolve)); + } + }, + }, + push(event) { + if (finished || failed) return; + queued.push(event); + wake(); + }, + finish() { + if (finished || failed) return; + finished = true; + wake(); + }, + fail(error) { + if (finished || failed) return; + failed = true; + failure = error; + wake(); + }, + }; +} + +function webSocketCloseError(code: number): DesignWorkspaceModuleError | null { + if (code === 1000 || code === 1001) return null; + if (code === 4401) { + return new DesignWorkspaceModuleError( + 401, + 'DESIGN_EVENT_TICKET_INVALID', + 'AI 设计任务连接凭证已失效', + ); + } + if (code === 4404) { + return new DesignWorkspaceModuleError( + 404, + 'DESIGN_EVENT_SESSION_NOT_FOUND', + 'AI 设计任务状态会话不存在', + ); + } + if (code === 4409) { + return new DesignWorkspaceModuleError( + 410, + 'DESIGN_EVENT_CURSOR_EXPIRED', + 'AI 设计任务状态断点已过期', + ); + } + return new DesignWorkspaceModuleError( + 502, + 'DESIGN_EVENT_STREAM_UNAVAILABLE', + 'AI 设计任务状态连接已断开', + ); +} + +function defaultAgentWebSocketFactory(url: string): AgentWebSocketConnection { + return { + socket: new WebSocket(url) as unknown as AgentWebSocket, + }; +} + +function eventSequence(afterEventId: string | undefined, sessionId: string): number { + if (!afterEventId?.startsWith(`${sessionId}:`)) return 0; + const sequence = Number(afterEventId.slice(sessionId.length + 1)); + return Number.isSafeInteger(sequence) && sequence >= 0 ? sequence : 0; +} + async function readPayload(response: Response): Promise { const text = await response.text(); if (!text.trim()) return null; @@ -216,10 +471,26 @@ function userFacingErrorMessage(code: string, fallback: string): string { export class WorksSquareDesignWorkspace implements DesignWorkspaceModule { private readonly apiBaseUrl: string; private readonly fetchImpl: typeof fetch; + private readonly webSocketFactory: AgentWebSocketFactory; + private readonly eventSessionClientIdStore: NonNullable< + WorksSquareDesignWorkspaceOptions['eventSessionClientIdStore'] + >; + private readonly eventSessions = new Map>(); + private readonly eventSessionClientIds = new Map(); + private readonly eventSubscriptionClosers = new Map void>>(); + private eventSessionsEnabled = true; constructor(options: WorksSquareDesignWorkspaceOptions = {}) { this.apiBaseUrl = (options.apiBaseUrl ?? WORKS_SQUARE_CONFIG.apiBaseUrl).replace(/\/+$/, ''); this.fetchImpl = options.fetchImpl ?? proxyAwareFetch; + this.webSocketFactory = options.webSocketFactory ?? defaultAgentWebSocketFactory; + const clientInstanceId = options.clientInstanceId?.trim() || createClientId('process'); + this.eventSessionClientIdStore = options.eventSessionClientIdStore ?? { + getOrCreate: async (workspaceId) => ( + stableWorkspaceSessionClientId(clientInstanceId, workspaceId) + ), + rotate: async () => createClientId('design-stream-rotated'), + }; } async bootstrap(): Promise { @@ -227,6 +498,7 @@ export class WorksSquareDesignWorkspace implements DesignWorkspaceModule { this.getCapabilities(), this.requestJson('/api/design/workspaces?limit=100&offset=0'), ]); + this.eventSessionsEnabled = true; return { capabilities, workspaces: workspaces.map(mapWorkspaceSummary), @@ -310,6 +582,212 @@ export class WorksSquareDesignWorkspace implements DesignWorkspaceModule { return tasks.map(mapTask); } + async openWorkspaceEvents( + input: DesignWorkspaceEventSubscriptionInput, + ): Promise { + let session = await this.ensureEventSession(input.workspaceId); + let ticket: ServerAgentStreamTicket; + try { + ticket = await this.createEventStreamTicket(session.session_id); + } catch (error) { + if (!(error instanceof DesignWorkspaceModuleError) + || (error.status !== 404 && error.status !== 409)) { + throw error; + } + await this.invalidateEventSession(input.workspaceId); + session = await this.ensureEventSession(input.workspaceId); + ticket = await this.createEventStreamTicket(session.session_id); + } + const streamUrl = new URL(ticket.stream_url, `${this.apiBaseUrl}/`); + if (streamUrl.origin !== new URL(this.apiBaseUrl).origin) { + throw new DesignWorkspaceModuleError( + 502, + 'DESIGN_EVENT_STREAM_INVALID', + 'AI 设计任务状态流地址无效', + ); + } + if (streamUrl.protocol !== 'http:' && streamUrl.protocol !== 'https:') { + throw new DesignWorkspaceModuleError( + 502, + 'DESIGN_EVENT_STREAM_INVALID', + 'AI 设计任务状态流协议无效', + ); + } + streamUrl.searchParams.set( + 'after_sequence', + String(eventSequence(input.afterEventId, session.session_id)), + ); + streamUrl.protocol = streamUrl.protocol === 'https:' ? 'wss:' : 'ws:'; + + let connection: AgentWebSocketConnection; + try { + connection = await this.webSocketFactory(streamUrl.toString()); + } catch { + throw new DesignWorkspaceModuleError( + 502, + 'DESIGN_EVENT_STREAM_UNAVAILABLE', + 'AI 设计任务状态连接失败', + ); + } + + const { socket } = connection; + const queue = createTaskEventQueue(); + let didOpen = false; + let ending = false; + let settled = false; + let heartbeat: ReturnType | null = null; + let unregister = () => undefined; + let resolveOpened: () => void = () => undefined; + let rejectOpened: (error: unknown) => void = () => undefined; + const opened = new Promise((resolve, reject) => { + resolveOpened = resolve; + rejectOpened = reject; + }); + const unavailableError = () => new DesignWorkspaceModuleError( + 502, + 'DESIGN_EVENT_STREAM_UNAVAILABLE', + 'AI 设计任务状态连接已断开', + ); + const settle = async (error: unknown | null): Promise => { + if (settled) return; + settled = true; + socket.onopen = null; + socket.onmessage = null; + socket.onerror = null; + socket.onclose = null; + if (heartbeat !== null) clearInterval(heartbeat); + heartbeat = null; + unregister(); + try { + await connection.dispose?.(); + } catch { + // The stream outcome is authoritative; dispatcher cleanup is best effort. + } + if (!didOpen) { + rejectOpened(error ?? unavailableError()); + } else if (error) { + queue.fail(error); + } else { + queue.finish(); + } + }; + const beginEnd = ( + error: unknown | null, + cleanup?: () => void | Promise, + ): void => { + if (ending) return; + ending = true; + void (async () => { + let finalError = error; + try { + await cleanup?.(); + } catch (cleanupError) { + finalError ??= cleanupError; + } + await settle(finalError); + })(); + }; + const close = (): void => { + if (ending) return; + beginEnd(null); + try { + socket.close(1000, 'Client closed design event stream'); + } catch { + // The local queue is already closed. + } + }; + + socket.onopen = () => { + if (ending) return; + didOpen = true; + heartbeat = setInterval(() => { + if (ending || socket.readyState !== AGENT_WEBSOCKET_OPEN) return; + try { + socket.send(JSON.stringify({ + type: 'ping', + request_id: `design-ping-${Date.now()}`, + })); + } catch { + beginEnd(unavailableError()); + } + }, AGENT_WEBSOCKET_PING_INTERVAL_MS); + resolveOpened(); + }; + socket.onmessage = ({ data }) => { + if (ending) return; + const agentEvent = agentEventFromWebSocketFrame(data); + const event = normalizeWorkspaceEvent( + agentEvent, + session.session_id, + input.workspaceId, + ); + if (event) queue.push(event); + }; + socket.onerror = () => { + setTimeout(() => { + if (!ending) beginEnd(unavailableError()); + }, 0); + }; + socket.onclose = ({ code }) => { + const error = webSocketCloseError(code); + if (code === 4409) { + beginEnd(error, async () => { + try { + await this.closeEventSession(session.session_id); + } finally { + await this.invalidateEventSession(input.workspaceId); + } + }); + return; + } + if (code === 4404) { + beginEnd(error, () => this.invalidateEventSession(input.workspaceId)); + return; + } + beginEnd(error); + }; + unregister = this.registerEventSubscription(input.workspaceId, close); + await opened; + return { + events: queue.events, + close, + }; + } + + async closeEventSessions(): Promise { + this.eventSessionsEnabled = false; + const activeSubscriptions = [...this.eventSubscriptionClosers.values()] + .flatMap((closers) => [...closers]); + this.eventSubscriptionClosers.clear(); + for (const close of activeSubscriptions) close(); + const pendingSessions = [...this.eventSessions.entries()]; + this.eventSessions.clear(); + const sessions = await Promise.allSettled( + pendingSessions.map(async ([workspaceId, pending]) => ({ + workspaceId, + session: await pending, + })), + ); + const closeResults = await Promise.allSettled( + sessions.flatMap((result) => ( + result.status === 'fulfilled' + ? [this.closeEventSession(result.value.session.session_id) + .then(() => this.rotateEventSession(result.value.workspaceId))] + : [] + )), + ); + const uncertainCreations = sessions.filter((result) => ( + result.status === 'rejected' + && !(result.reason instanceof DesignWorkspaceModuleError + && result.reason.code === 'DESIGN_EVENT_SESSION_CLOSED') + )).length; + const failedCloses = closeResults.filter((result) => result.status === 'rejected').length; + const failed = uncertainCreations + failedCloses; + if (failed > 0) { + throw new Error(`Failed to close ${failed} AI design Agent Session(s)`); + } + } + openAssetContent(workspaceId: string, assetId: string, range?: string): Promise { return this.authorizedFetch( `/api/design/workspaces/${encodeURIComponent(workspaceId)}/assets/${encodeURIComponent(assetId)}/content`, @@ -373,4 +851,94 @@ export class WorksSquareDesignWorkspace implements DesignWorkspaceModule { }, }); } + + private ensureEventSession(workspaceId: string): Promise { + if (!this.eventSessionsEnabled) { + throw new DesignWorkspaceModuleError( + 503, + 'DESIGN_EVENT_STREAM_PAUSED', + 'AI 设计任务状态流已暂停', + ); + } + const existing = this.eventSessions.get(workspaceId); + if (existing) return existing; + const pending = (async () => { + const clientSessionId = this.eventSessionClientIds.get(workspaceId) + ?? await this.eventSessionClientIdStore.getOrCreate(workspaceId); + this.eventSessionClientIds.set(workspaceId, clientSessionId); + const session = await this.requestJson('/api/agents/sessions', { + method: 'POST', + body: JSON.stringify({ + client_session_id: clientSessionId, + runtime: 'design', + runtime_version: 'v1', + binding: { + kind: 'design_workspace', + key: workspaceId, + }, + }), + }); + if (session.status !== 'active') { + throw new DesignWorkspaceModuleError( + 409, + 'DESIGN_EVENT_SESSION_CLOSED', + 'AI 设计任务状态会话已关闭', + ); + } + return session; + })().catch(async (error) => { + this.eventSessions.delete(workspaceId); + if (error instanceof DesignWorkspaceModuleError + && error.code === 'DESIGN_EVENT_SESSION_CLOSED') { + await this.rotateEventSession(workspaceId); + } + throw error; + }); + this.eventSessions.set(workspaceId, pending); + return pending; + } + + private createEventStreamTicket(sessionId: string): Promise { + return this.requestJson( + `/api/agents/sessions/${encodeURIComponent(sessionId)}/stream-tickets`, + { + method: 'POST', + body: JSON.stringify({ transport: 'websocket' }), + }, + ); + } + + private registerEventSubscription(workspaceId: string, close: () => void): () => void { + const closers = this.eventSubscriptionClosers.get(workspaceId) ?? new Set<() => void>(); + closers.add(close); + this.eventSubscriptionClosers.set(workspaceId, closers); + return () => { + closers.delete(close); + if (closers.size === 0) this.eventSubscriptionClosers.delete(workspaceId); + }; + } + + private async closeEventSession(sessionId: string): Promise { + try { + await this.requestJson( + `/api/agents/sessions/${encodeURIComponent(sessionId)}`, + { method: 'DELETE' }, + ); + } catch (error) { + if (error instanceof DesignWorkspaceModuleError + && (error.status === 404 || error.status === 409)) return; + throw error; + } + } + + private invalidateEventSession(workspaceId: string): Promise { + return this.rotateEventSession(workspaceId).then(() => undefined); + } + + private async rotateEventSession(workspaceId: string): Promise { + this.eventSessions.delete(workspaceId); + const clientSessionId = await this.eventSessionClientIdStore.rotate(workspaceId); + this.eventSessionClientIds.set(workspaceId, clientSessionId); + return clientSessionId; + } } diff --git a/electron/main/index.ts b/electron/main/index.ts index 925100c..5fce5e4 100644 --- a/electron/main/index.ts +++ b/electron/main/index.ts @@ -32,6 +32,10 @@ import { logger } from '../utils/logger'; import { warmupNetworkOptimization } from '../utils/uv-env'; import { resolvePythonRuntime } from '../utils/python-runtime'; import { initTelemetry } from '../utils/telemetry'; +import { + getOrCreateAgentGatewaySessionClientId, + rotateAgentGatewaySessionClientId, +} from '../utils/store'; import { isQuitting, setQuitting } from './app-state'; import { applyProxySettings } from './proxy'; @@ -73,6 +77,7 @@ import { LocalImageWorkspace, } from '../image-workspace/local-workspace'; import { WorksSquareDesignWorkspace } from '../image-workspace/works-square-workspace'; +import type { DesignWorkspaceModule } from '../image-workspace/module'; const WINDOWS_APP_USER_MODEL_ID = 'app.niancode.desktop'; const isE2EMode = process.env.NIANCODE_E2E === '1'; @@ -152,6 +157,7 @@ let hostApiServer: Server | null = null; let projectProgressSync: ReturnType | null = null; let worksCloudDeployment: ReturnType | null = null; let agentBrowser: AgentBrowserModule | null = null; +let imageWorkspaceModule: DesignWorkspaceModule | null = null; const mainWindowFocusState = createMainWindowFocusState(); const quitLifecycleState = createQuitLifecycleState(); const launchDeepLinkUrl = findNianCodeDeepLinkUrl(process.argv); @@ -419,7 +425,13 @@ async function initialize(): Promise { }); const imageWorkspace = localImageWorkspaceEnabled ? new LocalImageWorkspace({ userDataDir: app.getPath('userData') }) - : new WorksSquareDesignWorkspace(); + : new WorksSquareDesignWorkspace({ + eventSessionClientIdStore: { + getOrCreate: getOrCreateAgentGatewaySessionClientId, + rotate: rotateAgentGatewaySessionClientId, + }, + }); + imageWorkspaceModule = imageWorkspace; if (localImageWorkspaceEnabled) { logger.info('AI painting workspace is using local development storage'); } else { @@ -647,9 +659,13 @@ if (gotTheLock) { const stopAgentBrowserPromise = agentBrowser?.dispose().catch((err) => { logger.warn('agentBrowser.dispose() error during quit:', err); }) ?? Promise.resolve(); + const closeImageWorkspacePromise = imageWorkspaceModule?.closeEventSessions?.().catch((err) => { + logger.warn('imageWorkspace.closeEventSessions() error during quit:', err); + }) ?? Promise.resolve(); const stopPromise = Promise.allSettled([ stopOpencodePromise, stopAgentBrowserPromise, + closeImageWorkspacePromise, ]); const timeoutPromise = new Promise<'timeout'>((resolve) => { setTimeout(() => resolve('timeout'), 5000); diff --git a/electron/utils/store.ts b/electron/utils/store.ts index 74f0e75..af832a1 100644 --- a/electron/utils/store.ts +++ b/electron/utils/store.ts @@ -4,6 +4,7 @@ */ import { app } from 'electron'; +import { randomUUID } from 'node:crypto'; import { resolveSupportedLanguage } from '../../shared/language'; // Lazy-load electron-store (ESM module) @@ -21,6 +22,7 @@ export interface AppSettings { launchAtStartup: boolean; telemetryEnabled: boolean; machineId: string; + agentGatewaySessionClientIds: Record; hasReportedInstall: boolean; proxyEnabled: boolean; @@ -68,6 +70,7 @@ function createDefaultSettings(): AppSettings { launchAtStartup: false, telemetryEnabled: true, machineId: '', + agentGatewaySessionClientIds: {}, hasReportedInstall: false, proxyEnabled: false, @@ -127,6 +130,39 @@ export async function setSetting( store.set(key, value); } +let agentGatewaySessionIdMutation: Promise = Promise.resolve(); + +async function mutateAgentGatewaySessionClientIds( + mutation: (current: Record) => Promise | T, +): Promise { + const result = agentGatewaySessionIdMutation.then(async () => { + const current = await getSetting('agentGatewaySessionClientIds'); + return await mutation({ ...current }); + }); + agentGatewaySessionIdMutation = result.then(() => undefined, () => undefined); + return await result; +} + +export function getOrCreateAgentGatewaySessionClientId(workspaceId: string): Promise { + return mutateAgentGatewaySessionClientIds(async (current) => { + const existing = current[workspaceId]?.trim(); + if (existing) return existing; + const created = `design-stream-${randomUUID()}`; + current[workspaceId] = created; + await setSetting('agentGatewaySessionClientIds', current); + return created; + }); +} + +export function rotateAgentGatewaySessionClientId(workspaceId: string): Promise { + return mutateAgentGatewaySessionClientIds(async (current) => { + const created = `design-stream-${randomUUID()}`; + current[workspaceId] = created; + await setSetting('agentGatewaySessionClientIds', current); + return created; + }); +} + /** * Get all settings */ diff --git a/shared/image-workspace.ts b/shared/image-workspace.ts index 25c39c1..f6ca53d 100644 --- a/shared/image-workspace.ts +++ b/shared/image-workspace.ts @@ -85,6 +85,26 @@ export type DesignGenerationTask = { updatedAt: string; }; +export type DesignGenerationTaskUpdatedEvent = { + id: string; + type: 'design.generation_task.updated'; + workspaceId: string; + workspaceViewRevision: number; + generationTask: DesignGenerationTask; +}; + +export type DesignGenerationTasksSnapshotEvent = { + id: string; + type: 'design.generation_tasks.snapshot'; + workspaceId: string; + workspaceViewRevision: number; + generationTasks: DesignGenerationTask[]; +}; + +export type DesignWorkspaceEvent = + | DesignGenerationTaskUpdatedEvent + | DesignGenerationTasksSnapshotEvent; + export type DesignWorkspaceBootstrap = { capabilities: DesignCapabilities; workspaces: DesignWorkspaceSummary[]; diff --git a/src/lib/image-workspace.ts b/src/lib/image-workspace.ts index 4a3e2f8..48d4d98 100644 --- a/src/lib/image-workspace.ts +++ b/src/lib/image-workspace.ts @@ -1,5 +1,6 @@ import { AppError } from '@/lib/error-model'; import { + createHostEventSource, ensureHostApiToken, getHostApiBase, hostApiFetch, @@ -159,6 +160,13 @@ export function fetchImageWorkspaceTasks( ); } +export async function openImageWorkspaceTaskEvents(workspaceId: string): Promise { + await ensureHostApiToken(); + return createHostEventSource( + `${IMAGE_WORKSPACE_API_PATH}/workspaces/${encodeURIComponent(workspaceId)}/events`, + ); +} + export async function resolveImageWorkspaceAssetUrl(contentPath: string): Promise { const token = await ensureHostApiToken(); const separator = contentPath.includes('?') ? '&' : '?'; diff --git a/src/pages/ImageCanvas/index.tsx b/src/pages/ImageCanvas/index.tsx index e4c864a..a334899 100644 --- a/src/pages/ImageCanvas/index.tsx +++ b/src/pages/ImageCanvas/index.tsx @@ -194,6 +194,8 @@ export function ImageCanvas() { const load = useImageWorkspaceStore((state) => state.load); const refreshWorkspace = useImageWorkspaceStore((state) => state.refreshWorkspace); const refreshTasks = useImageWorkspaceStore((state) => state.refreshTasks); + const connectTaskStream = useImageWorkspaceStore((state) => state.connectTaskStream); + const disconnectTaskStream = useImageWorkspaceStore((state) => state.disconnectTaskStream); const sendMessage = useImageWorkspaceStore((state) => state.sendMessage); const confirmGeneration = useImageWorkspaceStore((state) => state.confirmGeneration); const [prompt, setPrompt] = useState(''); @@ -206,6 +208,7 @@ export function ImageCanvas() { () => workspace ? activeQuote(workspace.messages) : null, [workspace], ); + const taskWorkspaceId = workspace?.workspaceId ?? null; const hasActiveTasks = tasks.some((task) => ACTIVE_TASK_STATUSES.has(task.status)); useEffect(() => { @@ -213,12 +216,10 @@ export function ImageCanvas() { }, [authenticated, load, status]); useEffect(() => { - if (!hasActiveTasks || !workspace) return; - const timer = window.setInterval(() => { - void refreshTasks().catch(() => undefined); - }, 2_500); - return () => window.clearInterval(timer); - }, [hasActiveTasks, refreshTasks, workspace]); + if (!taskWorkspaceId) return; + connectTaskStream(); + return disconnectTaskStream; + }, [connectTaskStream, disconnectTaskStream, taskWorkspaceId]); useEffect(() => { if (typeof conversationEndRef.current?.scrollIntoView === 'function') { @@ -354,7 +355,9 @@ export function ImageCanvas() { aria-label="刷新设计项目" className="h-9 w-9 rounded-full" onClick={() => { - void Promise.all([refreshWorkspace(), refreshTasks()]).catch(() => undefined); + void refreshWorkspace() + .then(() => refreshTasks()) + .catch(() => undefined); }} > diff --git a/src/stores/image-workspace.ts b/src/stores/image-workspace.ts index 8f622fa..278ef5f 100644 --- a/src/stores/image-workspace.ts +++ b/src/stores/image-workspace.ts @@ -6,6 +6,7 @@ import { fetchImageWorkspaceProject, fetchImageWorkspaceTasks, ImageWorkspaceApiError, + openImageWorkspaceTaskEvents, renameImageWorkspaceProject, sendImageWorkspaceMessage, } from '@/lib/image-workspace'; @@ -13,6 +14,8 @@ import { useAuthStore } from '@/stores/auth'; import { IMAGE_WORKSPACE_UNAVAILABLE_CODE, type DesignGenerationTask, + type DesignGenerationTasksSnapshotEvent, + type DesignGenerationTaskUpdatedEvent, type DesignWorkspace, type DesignWorkspaceBootstrap, type DesignWorkspaceSummary, @@ -26,12 +29,15 @@ export type ImageWorkspaceLoadStatus = | 'error' | 'auth-required'; +export type ImageWorkspaceTaskStreamState = 'idle' | 'connecting' | 'connected' | 'degraded'; + type ImageWorkspaceState = { status: ImageWorkspaceLoadStatus; bootstrap: DesignWorkspaceBootstrap | null; activeWorkspaceId: string | null; workspace: DesignWorkspace | null; tasks: DesignGenerationTask[]; + taskStreamState: ImageWorkspaceTaskStreamState; error: string | null; load: () => Promise; createProject: (title: string) => Promise; @@ -39,17 +45,88 @@ type ImageWorkspaceState = { selectProject: (workspaceId: string) => Promise; refreshWorkspace: () => Promise; refreshTasks: () => Promise; + connectTaskStream: () => void; + disconnectTaskStream: () => void; sendMessage: (message: string) => Promise; confirmGeneration: (quoteId: string) => Promise; reset: () => void; }; let inFlightLoad: Promise | null = null; -const GENERATION_TASK_SYNC_ATTEMPTS = 20; -const GENERATION_TASK_SYNC_DELAY_MS = 500; +const TASK_FALLBACK_POLL_INTERVAL_MS = 15_000; +let activeTaskEventSource: EventSource | null = null; +let activeTaskEventWorkspaceId: string | null = null; +let taskStreamGeneration = 0; +let taskFallbackTimer: ReturnType | null = null; +const taskEventRevisions = new Map(); -function wait(milliseconds: number): Promise { - return new Promise((resolve) => setTimeout(resolve, milliseconds)); +function taskRevisionKey(workspaceId: string, taskId: string): string { + return `${workspaceId}:${taskId}`; +} + +function stopFallbackPolling(): void { + if (taskFallbackTimer !== null) clearInterval(taskFallbackTimer); + taskFallbackTimer = null; +} + +function closeTaskEventSource(): void { + taskStreamGeneration += 1; + stopFallbackPolling(); + if (activeTaskEventSource) { + activeTaskEventSource.onopen = null; + activeTaskEventSource.onerror = null; + activeTaskEventSource.close(); + } + activeTaskEventSource = null; + activeTaskEventWorkspaceId = null; + taskEventRevisions.clear(); +} + +function parseTaskUpdatedEvent(event: Event): DesignGenerationTaskUpdatedEvent | null { + const data = (event as MessageEvent).data; + if (typeof data !== 'string') return null; + try { + const payload = JSON.parse(data) as Partial; + if (payload.type !== 'design.generation_task.updated' + || typeof payload.id !== 'string' + || typeof payload.workspaceId !== 'string' + || !Number.isInteger(payload.workspaceViewRevision) + || !payload.generationTask + || typeof payload.generationTask.taskId !== 'string' + || payload.generationTask.workspaceId !== payload.workspaceId) { + return null; + } + return payload as DesignGenerationTaskUpdatedEvent; + } catch { + return null; + } +} + +function parseTasksSnapshotEvent(event: Event): DesignGenerationTasksSnapshotEvent | null { + const data = (event as MessageEvent).data; + if (typeof data !== 'string') return null; + try { + const payload = JSON.parse(data) as Partial; + if (payload.type !== 'design.generation_tasks.snapshot' + || typeof payload.id !== 'string' + || typeof payload.workspaceId !== 'string' + || !Number.isInteger(payload.workspaceViewRevision) + || !Array.isArray(payload.generationTasks) + || !payload.generationTasks.every((task) => ( + typeof task?.taskId === 'string' && task.workspaceId === payload.workspaceId + ))) { + return null; + } + return payload as DesignGenerationTasksSnapshotEvent; + } catch { + return null; + } +} + +function sortTasks(tasks: DesignGenerationTask[]): DesignGenerationTask[] { + return [...tasks].sort((left, right) => ( + right.createdAt.localeCompare(left.createdAt) || right.taskId.localeCompare(left.taskId) + )); } function unavailable(error: unknown): boolean { @@ -86,9 +163,147 @@ function upsertSummary( } export const useImageWorkspaceStore = create((set, get) => { + const stopTaskStream = () => { + closeTaskEventSource(); + set({ taskStreamState: 'idle' }); + }; + + const startFallbackPolling = () => { + if (taskFallbackTimer !== null) return; + taskFallbackTimer = setInterval(() => { + if (get().taskStreamState !== 'degraded' || !get().activeWorkspaceId) return; + void get().refreshWorkspace() + .catch(() => null) + .then(() => get().refreshTasks().catch(() => [])); + }, TASK_FALLBACK_POLL_INTERVAL_MS); + }; + + const startTaskStream = (workspaceId: string) => { + if (activeTaskEventWorkspaceId === workspaceId + && get().taskStreamState !== 'idle') return; + closeTaskEventSource(); + const generation = taskStreamGeneration; + activeTaskEventWorkspaceId = workspaceId; + set({ taskStreamState: 'connecting' }); + void openImageWorkspaceTaskEvents(workspaceId).then((source) => { + if (generation !== taskStreamGeneration || get().activeWorkspaceId !== workspaceId) { + source.close(); + return; + } + activeTaskEventSource = source; + source.addEventListener('design.generation_tasks.snapshot', (event) => { + const snapshot = parseTasksSnapshotEvent(event); + if (!snapshot || snapshot.workspaceId !== activeTaskEventWorkspaceId) return; + set((state) => { + if (state.activeWorkspaceId !== snapshot.workspaceId) return state; + const tasksById = new Map(state.tasks.map((task) => [task.taskId, task])); + for (const task of snapshot.generationTasks) { + const key = taskRevisionKey(snapshot.workspaceId, task.taskId); + const knownRevision = taskEventRevisions.get(key) ?? 0; + const existing = tasksById.get(task.taskId); + if (snapshot.workspaceViewRevision >= knownRevision + && (!existing || existing.updatedAt <= task.updatedAt)) { + tasksById.set(task.taskId, task); + } + taskEventRevisions.set( + key, + Math.max(knownRevision, snapshot.workspaceViewRevision), + ); + } + const workspace = state.workspace?.workspaceId === snapshot.workspaceId + ? { + ...state.workspace, + viewRevision: Math.max( + state.workspace.viewRevision, + snapshot.workspaceViewRevision, + ), + } + : state.workspace; + const bootstrap = state.bootstrap + ? { + ...state.bootstrap, + workspaces: state.bootstrap.workspaces.map((candidate) => ( + candidate.workspaceId === snapshot.workspaceId + ? { + ...candidate, + viewRevision: Math.max( + candidate.viewRevision, + snapshot.workspaceViewRevision, + ), + } + : candidate + )), + } + : null; + return { tasks: sortTasks([...tasksById.values()]), workspace, bootstrap }; + }); + }); + source.addEventListener('design.generation_task.updated', (event) => { + const update = parseTaskUpdatedEvent(event); + if (!update || update.workspaceId !== activeTaskEventWorkspaceId) return; + set((state) => { + if (state.activeWorkspaceId !== update.workspaceId) return state; + const key = taskRevisionKey(update.workspaceId, update.generationTask.taskId); + const knownRevision = taskEventRevisions.get(key) ?? 0; + if (update.workspaceViewRevision <= knownRevision) return state; + const existing = state.tasks.find( + (candidate) => candidate.taskId === update.generationTask.taskId, + ); + if (existing && existing.updatedAt > update.generationTask.updatedAt) return state; + taskEventRevisions.set(key, update.workspaceViewRevision); + const tasks = sortTasks([ + update.generationTask, + ...state.tasks.filter((candidate) => candidate.taskId !== update.generationTask.taskId), + ]); + const workspace = state.workspace?.workspaceId === update.workspaceId + ? { + ...state.workspace, + viewRevision: Math.max( + state.workspace.viewRevision, + update.workspaceViewRevision, + ), + } + : state.workspace; + const bootstrap = state.bootstrap + ? { + ...state.bootstrap, + workspaces: state.bootstrap.workspaces.map((candidate) => ( + candidate.workspaceId === update.workspaceId + ? { + ...candidate, + viewRevision: Math.max( + candidate.viewRevision, + update.workspaceViewRevision, + ), + } + : candidate + )), + } + : null; + return { tasks, workspace, bootstrap }; + }); + }); + source.onopen = () => { + if (activeTaskEventSource !== source) return; + stopFallbackPolling(); + set({ taskStreamState: 'connected' }); + }; + source.onerror = () => { + if (activeTaskEventSource !== source) return; + set({ taskStreamState: 'degraded' }); + startFallbackPolling(); + }; + }).catch(() => { + if (generation !== taskStreamGeneration || get().activeWorkspaceId !== workspaceId) return; + set({ taskStreamState: 'degraded' }); + startFallbackPolling(); + }); + }; + const handleRequestError = (error: unknown): string => { const message = messageOf(error); if (authenticationRequired(error)) { + closeTaskEventSource(); useAuthStore.getState().invalidateSession(); set({ status: 'auth-required', @@ -96,6 +311,7 @@ export const useImageWorkspaceStore = create((set, get) => activeWorkspaceId: null, workspace: null, tasks: [], + taskStreamState: 'idle', error: message, }); } @@ -113,34 +329,23 @@ export const useImageWorkspaceStore = create((set, get) => return workspace; }; - const waitForGenerationTask = async ( - workspaceId: string, - quoteId: string, - ): Promise => { - for (let attempt = 0; attempt < GENERATION_TASK_SYNC_ATTEMPTS; attempt += 1) { - if (get().activeWorkspaceId !== workspaceId) return []; - const tasks = await get().refreshTasks().catch(() => []); - if (tasks.some((task) => task.quoteId === quoteId)) return tasks; - if (attempt < GENERATION_TASK_SYNC_ATTEMPTS - 1) { - await wait(GENERATION_TASK_SYNC_DELAY_MS); - } - } - throw new Error('生成请求已确认,但任务列表尚未同步,请点击刷新重试'); - }; - const loadWorkspace = async (workspaceId: string): Promise => { - const [workspace, tasks] = await Promise.all([ - fetchImageWorkspaceProject(workspaceId), - fetchImageWorkspaceTasks(workspaceId), - ]); + const workspace = await fetchImageWorkspaceProject(workspaceId); + if (get().activeWorkspaceId !== workspaceId) return; + const tasks = await fetchImageWorkspaceTasks(workspaceId); + if (get().activeWorkspaceId !== workspaceId) return; set((state) => ({ status: 'ready', bootstrap: upsertSummary(state.bootstrap, workspace), activeWorkspaceId: workspaceId, workspace, - tasks, + tasks: sortTasks(tasks), error: null, })); + startTaskStream(workspaceId); + for (const task of tasks) { + taskEventRevisions.set(taskRevisionKey(workspaceId, task.taskId), workspace.viewRevision); + } }; const recoverRevisionConflict = async (error: unknown): Promise => { @@ -162,6 +367,7 @@ export const useImageWorkspaceStore = create((set, get) => activeWorkspaceId: null, workspace: null, tasks: [], + taskStreamState: 'idle', error: null, load: () => { @@ -176,12 +382,14 @@ export const useImageWorkspaceStore = create((set, get) => ) ? currentId : bootstrap.workspaces[0]?.workspaceId ?? null; + closeTaskEventSource(); set({ status: 'ready', bootstrap, activeWorkspaceId, workspace: null, tasks: [], + taskStreamState: 'idle', error: null, }); if (activeWorkspaceId) await loadWorkspace(activeWorkspaceId); @@ -189,12 +397,14 @@ export const useImageWorkspaceStore = create((set, get) => } catch (error) { const message = handleRequestError(error); if (!authenticationRequired(error)) { + closeTaskEventSource(); set({ status: unavailable(error) ? 'unavailable' : 'error', bootstrap: null, activeWorkspaceId: null, workspace: null, tasks: [], + taskStreamState: 'idle', error: message, }); } @@ -210,6 +420,7 @@ export const useImageWorkspaceStore = create((set, get) => try { const workspace = applyWorkspace(await createImageWorkspaceProject(title)); set({ tasks: [] }); + startTaskStream(workspace.workspaceId); return workspace; } catch (error) { set({ error: handleRequestError(error) }); @@ -228,10 +439,12 @@ export const useImageWorkspaceStore = create((set, get) => selectProject: async (workspaceId) => { if (!get().bootstrap?.workspaces.some((item) => item.workspaceId === workspaceId)) return; + stopTaskStream(); set({ activeWorkspaceId: workspaceId, workspace: null, tasks: [], + taskStreamState: 'idle', error: null, }); try { @@ -245,7 +458,14 @@ export const useImageWorkspaceStore = create((set, get) => const workspaceId = get().activeWorkspaceId; if (!workspaceId) return null; try { - return applyWorkspace(await fetchImageWorkspaceProject(workspaceId)); + const refreshed = await fetchImageWorkspaceProject(workspaceId); + if (get().activeWorkspaceId !== workspaceId) return null; + const current = get().workspace; + if (current?.workspaceId === workspaceId + && current.viewRevision > refreshed.viewRevision) { + return current; + } + return applyWorkspace(refreshed); } catch (error) { set({ error: handleRequestError(error) }); throw error; @@ -258,9 +478,24 @@ export const useImageWorkspaceStore = create((set, get) => set({ tasks: [] }); return []; } + const requestedViewRevision = get().workspace?.workspaceId === workspaceId + ? get().workspace?.viewRevision ?? 0 + : 0; try { const tasks = await fetchImageWorkspaceTasks(workspaceId); - set({ tasks }); + const currentWorkspace = get().workspace; + if (get().activeWorkspaceId === workspaceId + && currentWorkspace?.workspaceId === workspaceId + && currentWorkspace.viewRevision === requestedViewRevision) { + set({ tasks: sortTasks(tasks) }); + for (const task of tasks) { + const key = taskRevisionKey(workspaceId, task.taskId); + taskEventRevisions.set( + key, + Math.max(taskEventRevisions.get(key) ?? 0, requestedViewRevision), + ); + } + } return tasks; } catch (error) { set({ error: handleRequestError(error) }); @@ -268,6 +503,13 @@ export const useImageWorkspaceStore = create((set, get) => } }, + connectTaskStream: () => { + const workspaceId = get().activeWorkspaceId; + if (workspaceId) startTaskStream(workspaceId); + }, + + disconnectTaskStream: stopTaskStream, + sendMessage: async (message) => { const workspace = get().workspace; if (!workspace) throw new Error('请先选择设计项目'); @@ -294,7 +536,7 @@ export const useImageWorkspaceStore = create((set, get) => workspace.turnRevision, quoteId, )); - await waitForGenerationTask(workspace.workspaceId, quoteId); + await get().refreshTasks().catch(() => []); return updated; } catch (error) { set({ error: handleRequestError(error) }); @@ -304,12 +546,14 @@ export const useImageWorkspaceStore = create((set, get) => reset: () => { inFlightLoad = null; + closeTaskEventSource(); set({ status: 'idle', bootstrap: null, activeWorkspaceId: null, workspace: null, tasks: [], + taskStreamState: 'idle', error: null, }); }, diff --git a/tests/unit/auth-routes.test.ts b/tests/unit/auth-routes.test.ts index 316042c..3f2c9a3 100644 --- a/tests/unit/auth-routes.test.ts +++ b/tests/unit/auth-routes.test.ts @@ -272,6 +272,7 @@ describe('auth host api routes', () => { ); vi.stubGlobal('fetch', fetchMock); const stop = vi.fn(async () => undefined); + const closeEventSessions = vi.fn(async () => undefined); const response = createResponse(); const handled = await handleAuthRoutes( @@ -280,6 +281,7 @@ describe('auth host api routes', () => { new URL('http://127.0.0.1:13210/api/auth/logout'), { opencodeManager: { stop }, + imageWorkspace: { closeEventSessions }, } as never, ); @@ -287,6 +289,7 @@ describe('auth host api routes', () => { expect(response.statusCode).toBe(200); expect(response.json()).toEqual({ success: true }); expect(stop).toHaveBeenCalledOnce(); + expect(closeEventSessions).toHaveBeenCalledOnce(); expect(providerServiceMock.deleteAccountApiKey).toHaveBeenCalledWith('niancode-user-models'); }); diff --git a/tests/unit/image-canvas-page.test.tsx b/tests/unit/image-canvas-page.test.tsx index 0d0c44a..630edf3 100644 --- a/tests/unit/image-canvas-page.test.tsx +++ b/tests/unit/image-canvas-page.test.tsx @@ -1,4 +1,4 @@ -import { fireEvent, render, screen, waitFor } from '@testing-library/react'; +import { act, fireEvent, render, screen, waitFor } from '@testing-library/react'; import { MemoryRouter } from 'react-router-dom'; import { beforeEach, describe, expect, it, vi } from 'vitest'; import { ImageWorkspaceApiError } from '@/lib/image-workspace'; @@ -7,6 +7,7 @@ import { useAuthStore } from '@/stores/auth'; import { useImageWorkspaceStore } from '@/stores/image-workspace'; import type { DesignGenerationTask, + DesignGenerationTaskUpdatedEvent, DesignWorkspace, DesignWorkspaceBootstrap, } from '../../shared/image-workspace'; @@ -16,8 +17,29 @@ const fetchImageWorkspaceProjectMock = vi.hoisted(() => vi.fn()); const fetchImageWorkspaceTasksMock = vi.hoisted(() => vi.fn()); const sendImageWorkspaceMessageMock = vi.hoisted(() => vi.fn()); const confirmImageWorkspaceGenerationMock = vi.hoisted(() => vi.fn()); +const openImageWorkspaceTaskEventsMock = vi.hoisted(() => vi.fn()); const resolveImageWorkspaceAssetUrlMock = vi.hoisted(() => vi.fn()); +type EventListener = (event: MessageEvent) => void; + +class MockEventSource { + onopen: ((event: Event) => void) | null = null; + onerror: ((event: Event) => void) | null = null; + readonly close = vi.fn(); + private readonly listeners = new Map>(); + + addEventListener(type: string, listener: EventListener): void { + const listeners = this.listeners.get(type) ?? new Set(); + listeners.add(listener); + this.listeners.set(type, listeners); + } + + emit(type: string, payload: unknown): void { + const event = { data: JSON.stringify(payload) } as MessageEvent; + for (const listener of this.listeners.get(type) ?? []) listener(event); + } +} + vi.mock('@/lib/image-workspace', async (importOriginal) => { const actual = await importOriginal(); return { @@ -29,6 +51,9 @@ vi.mock('@/lib/image-workspace', async (importOriginal) => { confirmImageWorkspaceGeneration: (...args: unknown[]) => ( confirmImageWorkspaceGenerationMock(...args) ), + openImageWorkspaceTaskEvents: (...args: unknown[]) => ( + openImageWorkspaceTaskEventsMock(...args) + ), resolveImageWorkspaceAssetUrl: (...args: unknown[]) => ( resolveImageWorkspaceAssetUrlMock(...args) ), @@ -122,6 +147,8 @@ const taskFixture: DesignGenerationTask = { }; describe('ImageCanvas Workspace-first design experience', () => { + let taskEventSource: MockEventSource; + beforeEach(() => { vi.clearAllMocks(); useImageWorkspaceStore.getState().reset(); @@ -134,6 +161,10 @@ describe('ImageCanvas Workspace-first design experience', () => { turnRevision: 2, phase: 'shaping', }); + taskEventSource = new MockEventSource(); + openImageWorkspaceTaskEventsMock.mockResolvedValue( + taskEventSource as unknown as EventSource, + ); resolveImageWorkspaceAssetUrlMock.mockResolvedValue( 'http://127.0.0.1:13210/content?token=host', ); @@ -330,27 +361,33 @@ describe('ImageCanvas Workspace-first design experience', () => { await waitFor(() => expect(fetchImageWorkspaceTasksMock).toHaveBeenCalledTimes(2)); }); - it('keeps checking after Quote confirmation until the new task becomes visible', async () => { + it('renders a new task pushed by the design event stream without repeated polling', async () => { const queuedTask = { ...taskFixture, taskId: 'task-delayed', status: 'queued' as const, resultAssets: [], }; - fetchImageWorkspaceTasksMock - .mockResolvedValueOnce([]) - .mockResolvedValueOnce([]) - .mockResolvedValueOnce([queuedTask]); + fetchImageWorkspaceTasksMock.mockResolvedValue([]); render(); await screen.findByTestId('design-quote-quote-one'); + await waitFor(() => expect(openImageWorkspaceTaskEventsMock) + .toHaveBeenCalledWith('workspace-cloud')); + await waitFor(() => expect(taskEventSource.onopen).not.toBeNull()); + act(() => { + taskEventSource.onopen?.(new Event('open')); + taskEventSource.emit('design.generation_task.updated', { + id: 'session-one:2', + type: 'design.generation_task.updated', + workspaceId: 'workspace-cloud', + workspaceViewRevision: 2, + generationTask: queuedTask, + } satisfies DesignGenerationTaskUpdatedEvent); + }); - fireEvent.click(screen.getByRole('button', { name: '确认并开始生成' })); - - await waitFor(() => expect(confirmImageWorkspaceGenerationMock) - .toHaveBeenCalledWith('workspace-cloud', 1, 'quote-one')); - await waitFor(() => expect(fetchImageWorkspaceTasksMock).toHaveBeenCalledTimes(3)); expect(await screen.findByTestId('design-task-task-delayed')) .toBeInTheDocument(); + expect(fetchImageWorkspaceTasksMock).toHaveBeenCalledOnce(); }); }); diff --git a/tests/unit/image-workspace-api.test.ts b/tests/unit/image-workspace-api.test.ts index f224f64..434c23d 100644 --- a/tests/unit/image-workspace-api.test.ts +++ b/tests/unit/image-workspace-api.test.ts @@ -1,6 +1,7 @@ import { beforeEach, describe, expect, it, vi } from 'vitest'; import { AppError } from '@/lib/error-model'; import { + createHostEventSource, ensureHostApiToken, getHostApiBase, hostApiFetch, @@ -10,17 +11,20 @@ import { createImageWorkspaceProject, fetchImageWorkspace, ImageWorkspaceApiError, + openImageWorkspaceTaskEvents, resolveImageWorkspaceAssetUrl, sendImageWorkspaceMessage, } from '@/lib/image-workspace'; vi.mock('@/lib/host-api', () => ({ hostApiFetch: vi.fn(), + createHostEventSource: vi.fn(), ensureHostApiToken: vi.fn(), getHostApiBase: vi.fn(), })); const hostApiFetchMock = vi.mocked(hostApiFetch); +const createHostEventSourceMock = vi.mocked(createHostEventSource); const ensureHostApiTokenMock = vi.mocked(ensureHostApiToken); const getHostApiBaseMock = vi.mocked(getHostApiBase); @@ -40,6 +44,8 @@ describe('AI design renderer API boundary', () => { hostApiFetchMock.mockResolvedValue({ success: true, status: 200, data: bootstrap }); ensureHostApiTokenMock.mockReset(); ensureHostApiTokenMock.mockResolvedValue('host-token'); + createHostEventSourceMock.mockReset(); + createHostEventSourceMock.mockReturnValue({ close: vi.fn() } as unknown as EventSource); getHostApiBaseMock.mockReset(); getHostApiBaseMock.mockReturnValue('http://127.0.0.1:13210'); }); @@ -50,6 +56,17 @@ describe('AI design renderer API boundary', () => { expect(hostApiFetchMock).toHaveBeenCalledWith('/api/works/image-workspace', {}); }); + it('opens task events only after the local Host API token is ready', async () => { + const source = await openImageWorkspaceTaskEvents('workspace/one'); + + expect(source).toBe(createHostEventSourceMock.mock.results[0].value); + expect(createHostEventSourceMock).toHaveBeenCalledWith( + '/api/works/image-workspace/workspaces/workspace%2Fone/events', + ); + expect(ensureHostApiTokenMock.mock.invocationCallOrder[0]) + .toBeLessThan(createHostEventSourceMock.mock.invocationCallOrder[0]); + }); + it('creates an idempotent Workspace with a trimmed user-visible title', async () => { await createImageWorkspaceProject(' 角色设计 ', 'workspace-client-1'); diff --git a/tests/unit/image-workspace-route.test.ts b/tests/unit/image-workspace-route.test.ts index e97041b..bfe072d 100644 --- a/tests/unit/image-workspace-route.test.ts +++ b/tests/unit/image-workspace-route.test.ts @@ -215,4 +215,60 @@ describe('AI design Main route boundary', () => { expect(response.headers.get('content-range')).toBe('bytes 0-6/100'); expect(Buffer.concat(response.chunks).toString()).toBe('partial'); }); + + it('relays normalized task events over local SSE and forwards the opaque resume cursor', async () => { + const close = vi.fn(); + const openWorkspaceEvents = vi.fn().mockResolvedValue({ + events: (async function* () { + yield { + id: 'session-one:7', + type: 'design.generation_task.updated' as const, + workspaceId: 'workspace/one', + workspaceViewRevision: 8, + generationTask: { + taskId: 'task-live', + workspaceId: 'workspace/one', + medium: 'image' as const, + status: 'succeeded' as const, + briefVersion: 1, + briefSummary: '海洋公益海报', + quoteId: 'quote-live', + quotedDesignPoints: 1, + failureCode: null, + resultAssets: [], + createdAt: '2026-08-02T10:00:00Z', + updatedAt: '2026-08-02T10:02:00Z', + }, + }; + })(), + close, + }); + const request = createRequest('GET'); + request.headers = { 'last-event-id': 'session-one:6' }; + const response = new MediaResponse(); + + const handled = await handleImageWorkspaceRoutes( + request, + response as unknown as ServerResponse, + new URL( + 'http://127.0.0.1/api/works/image-workspace/workspaces/workspace%2Fone/events', + ), + { imageWorkspace: { openWorkspaceEvents } } as unknown as HostApiContext, + ); + + expect(handled).toBe(true); + expect(openWorkspaceEvents).toHaveBeenCalledWith({ + workspaceId: 'workspace/one', + afterEventId: 'session-one:6', + }); + expect(response.statusCode).toBe(200); + expect(response.headers.get('content-type')).toBe('text/event-stream; charset=utf-8'); + expect(response.headers.get('x-accel-buffering')).toBe('no'); + expect(Buffer.concat(response.chunks).toString()).toContain([ + 'id: session-one:7', + 'event: design.generation_task.updated', + 'data: {"id":"session-one:7"', + ].join('\n')); + expect(close).toHaveBeenCalledOnce(); + }); }); diff --git a/tests/unit/image-workspace-store.test.ts b/tests/unit/image-workspace-store.test.ts new file mode 100644 index 0000000..9fb75a5 --- /dev/null +++ b/tests/unit/image-workspace-store.test.ts @@ -0,0 +1,261 @@ +import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'; +import { useImageWorkspaceStore } from '@/stores/image-workspace'; +import type { + DesignGenerationTask, + DesignGenerationTasksSnapshotEvent, + DesignGenerationTaskUpdatedEvent, + DesignWorkspace, + DesignWorkspaceBootstrap, +} from '../../shared/image-workspace'; + +const fetchImageWorkspaceMock = vi.hoisted(() => vi.fn()); +const fetchImageWorkspaceProjectMock = vi.hoisted(() => vi.fn()); +const fetchImageWorkspaceTasksMock = vi.hoisted(() => vi.fn()); +const openImageWorkspaceTaskEventsMock = vi.hoisted(() => vi.fn()); + +vi.mock('@/lib/image-workspace', async (importOriginal) => { + const actual = await importOriginal(); + return { + ...actual, + fetchImageWorkspace: (...args: unknown[]) => fetchImageWorkspaceMock(...args), + fetchImageWorkspaceProject: (...args: unknown[]) => fetchImageWorkspaceProjectMock(...args), + fetchImageWorkspaceTasks: (...args: unknown[]) => fetchImageWorkspaceTasksMock(...args), + openImageWorkspaceTaskEvents: (...args: unknown[]) => openImageWorkspaceTaskEventsMock(...args), + }; +}); + +type EventListener = (event: MessageEvent) => void; + +class MockEventSource { + onopen: ((event: Event) => void) | null = null; + onerror: ((event: Event) => void) | null = null; + readonly close = vi.fn(); + private readonly listeners = new Map>(); + + addEventListener(type: string, listener: EventListener): void { + const listeners = this.listeners.get(type) ?? new Set(); + listeners.add(listener); + this.listeners.set(type, listeners); + } + + emit(type: string, payload: unknown): void { + const event = { data: JSON.stringify(payload) } as MessageEvent; + for (const listener of this.listeners.get(type) ?? []) listener(event); + } +} + +function deferred(): { promise: Promise; resolve(value: T): void } { + let resolve!: (value: T) => void; + const promise = new Promise((accept) => { + resolve = accept; + }); + return { promise, resolve }; +} + +const task: DesignGenerationTask = { + taskId: 'task-one', + workspaceId: 'workspace-one', + medium: 'image', + status: 'queued', + briefVersion: 1, + briefSummary: '海洋公益海报', + quoteId: 'quote-one', + quotedDesignPoints: 1, + failureCode: null, + resultAssets: [], + createdAt: '2026-08-02T10:00:00Z', + updatedAt: '2026-08-02T10:00:00Z', +}; + +function workspace(workspaceId = 'workspace-one', viewRevision = 1): DesignWorkspace { + return { + workspaceId, + title: workspaceId, + turnRevision: 1, + viewRevision, + phase: 'shaping', + brief: { + version: 1, + status: 'ready', + medium: 'image', + summary: '海洋公益海报', + ready: true, + missingDecision: null, + }, + updatedAt: '2026-08-02T10:00:00Z', + messages: [], + }; +} + +function bootstrap(workspaceIds = ['workspace-one']): DesignWorkspaceBootstrap { + return { + capabilities: { conversation: true, generation: true, image: true, video: true }, + workspaces: workspaceIds.map((workspaceId) => { + const { messages: _messages, ...summary } = workspace(workspaceId); + return summary; + }), + }; +} + +function taskEvent( + workspaceViewRevision: number, + status: DesignGenerationTask['status'], +): DesignGenerationTaskUpdatedEvent { + return { + id: `session-one:${workspaceViewRevision}`, + type: 'design.generation_task.updated', + workspaceId: 'workspace-one', + workspaceViewRevision, + generationTask: { + ...task, + status, + updatedAt: `2026-08-02T10:0${workspaceViewRevision}:00Z`, + }, + }; +} + +function taskSnapshotEvent( + workspaceViewRevision: number, + status: DesignGenerationTask['status'], +): DesignGenerationTasksSnapshotEvent { + return { + id: `session-one:${workspaceViewRevision}`, + type: 'design.generation_tasks.snapshot', + workspaceId: 'workspace-one', + workspaceViewRevision, + generationTasks: [{ + ...task, + status, + updatedAt: `2026-08-02T10:0${workspaceViewRevision}:00Z`, + }], + }; +} + +describe('AI design task event store', () => { + beforeEach(() => { + vi.useRealTimers(); + vi.clearAllMocks(); + useImageWorkspaceStore.getState().reset(); + fetchImageWorkspaceMock.mockResolvedValue(bootstrap()); + fetchImageWorkspaceProjectMock.mockResolvedValue(workspace()); + fetchImageWorkspaceTasksMock.mockResolvedValue([task]); + }); + + afterEach(() => { + useImageWorkspaceStore.getState().reset(); + vi.useRealTimers(); + }); + + it('keeps a revision-5 snapshot over revision 3, then upserts a newer task by task id', async () => { + const source = new MockEventSource(); + openImageWorkspaceTaskEventsMock.mockResolvedValue(source as unknown as EventSource); + fetchImageWorkspaceProjectMock.mockResolvedValue(workspace('workspace-one', 5)); + + await useImageWorkspaceStore.getState().load(); + await vi.waitFor(() => expect(openImageWorkspaceTaskEventsMock) + .toHaveBeenCalledWith('workspace-one')); + source.onopen?.(new Event('open')); + source.emit('design.generation_task.updated', taskEvent(3, 'failed')); + expect(useImageWorkspaceStore.getState().tasks[0].status).toBe('queued'); + source.emit('design.generation_tasks.snapshot', taskSnapshotEvent(6, 'running')); + source.emit('design.generation_task.updated', taskEvent(5, 'failed')); + + expect(useImageWorkspaceStore.getState()).toMatchObject({ + taskStreamState: 'connected', + workspace: { viewRevision: 6 }, + tasks: [{ taskId: 'task-one', status: 'running' }], + }); + }); + + it('closes the previous stream when switching workspaces and on reset', async () => { + const first = new MockEventSource(); + const second = new MockEventSource(); + openImageWorkspaceTaskEventsMock + .mockResolvedValueOnce(first as unknown as EventSource) + .mockResolvedValueOnce(second as unknown as EventSource); + fetchImageWorkspaceMock.mockResolvedValue(bootstrap(['workspace-one', 'workspace-two'])); + fetchImageWorkspaceProjectMock.mockImplementation((workspaceId: string) => ( + Promise.resolve(workspace(workspaceId)) + )); + + await useImageWorkspaceStore.getState().load(); + await vi.waitFor(() => expect(openImageWorkspaceTaskEventsMock).toHaveBeenCalledTimes(1)); + await useImageWorkspaceStore.getState().selectProject('workspace-two'); + await vi.waitFor(() => expect(openImageWorkspaceTaskEventsMock).toHaveBeenCalledTimes(2)); + + expect(first.close).toHaveBeenCalledOnce(); + useImageWorkspaceStore.getState().reset(); + expect(second.close).toHaveBeenCalledOnce(); + expect(useImageWorkspaceStore.getState().taskStreamState).toBe('idle'); + }); + + it('does not let a slow previous Workspace selection overwrite the latest one', async () => { + const rootSource = new MockEventSource(); + const latestSource = new MockEventSource(); + const slowWorkspace = deferred(); + openImageWorkspaceTaskEventsMock + .mockResolvedValueOnce(rootSource as unknown as EventSource) + .mockResolvedValueOnce(latestSource as unknown as EventSource); + fetchImageWorkspaceMock.mockResolvedValue(bootstrap([ + 'workspace-root', + 'workspace-one', + 'workspace-two', + ])); + fetchImageWorkspaceProjectMock.mockImplementation((workspaceId: string) => ( + workspaceId === 'workspace-one' + ? slowWorkspace.promise + : Promise.resolve(workspace(workspaceId)) + )); + fetchImageWorkspaceTasksMock.mockResolvedValue([]); + + await useImageWorkspaceStore.getState().load(); + await vi.waitFor(() => expect(openImageWorkspaceTaskEventsMock) + .toHaveBeenCalledWith('workspace-root')); + const staleSelection = useImageWorkspaceStore.getState().selectProject('workspace-one'); + await vi.waitFor(() => expect(fetchImageWorkspaceProjectMock) + .toHaveBeenCalledWith('workspace-one')); + await useImageWorkspaceStore.getState().selectProject('workspace-two'); + slowWorkspace.resolve(workspace('workspace-one')); + await staleSelection; + + expect(useImageWorkspaceStore.getState().activeWorkspaceId).toBe('workspace-two'); + expect(useImageWorkspaceStore.getState().workspace?.workspaceId).toBe('workspace-two'); + expect(openImageWorkspaceTaskEventsMock).not.toHaveBeenCalledWith('workspace-one'); + }); + + it('advances each task watermark after a degraded REST snapshot', async () => { + const source = new MockEventSource(); + openImageWorkspaceTaskEventsMock.mockResolvedValue(source as unknown as EventSource); + + await useImageWorkspaceStore.getState().load(); + await vi.waitFor(() => expect(openImageWorkspaceTaskEventsMock).toHaveBeenCalledOnce()); + fetchImageWorkspaceProjectMock.mockResolvedValue(workspace('workspace-one', 8)); + fetchImageWorkspaceTasksMock.mockResolvedValue([{ + ...task, + status: 'succeeded', + }]); + await useImageWorkspaceStore.getState().refreshWorkspace(); + await useImageWorkspaceStore.getState().refreshTasks(); + source.emit('design.generation_task.updated', taskEvent(7, 'failed')); + + expect(useImageWorkspaceStore.getState().tasks[0].status).toBe('succeeded'); + }); + + it('polls slowly only while the event stream is degraded', async () => { + vi.useFakeTimers(); + const source = new MockEventSource(); + openImageWorkspaceTaskEventsMock.mockResolvedValue(source as unknown as EventSource); + + await useImageWorkspaceStore.getState().load(); + await vi.waitFor(() => expect(source.onerror).not.toBeNull()); + source.onerror?.(new Event('error')); + await vi.advanceTimersByTimeAsync(14_999); + expect(fetchImageWorkspaceTasksMock).toHaveBeenCalledTimes(1); + await vi.advanceTimersByTimeAsync(1); + expect(fetchImageWorkspaceTasksMock).toHaveBeenCalledTimes(2); + + source.onopen?.(new Event('open')); + await vi.advanceTimersByTimeAsync(15_000); + expect(fetchImageWorkspaceTasksMock).toHaveBeenCalledTimes(2); + }); +}); diff --git a/tests/unit/works-square-design-workspace.test.ts b/tests/unit/works-square-design-workspace.test.ts index 6b23e52..a1f41dc 100644 --- a/tests/unit/works-square-design-workspace.test.ts +++ b/tests/unit/works-square-design-workspace.test.ts @@ -49,6 +49,68 @@ function jsonResponse(payload: unknown, status = 200): Response { }); } +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(); @@ -193,4 +255,356 @@ describe('Works Square AI design adapter', () => { }), ); }); + + 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: { + workspace_id: 'workspace-one', + view_revision: 2, + }, + generation_tasks: [snapshotTask], + }, + }; + 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: 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, + generationTasks: [expect.objectContaining({ + taskId: 'task-snapshot', + medium: 'image', + status: 'running', + })], + }, + { + 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}$/); + }); });