screenshot-to-code 变体系统非阻塞改造方案：从阻塞式并行生成到逐变体实时交付

本文基于仓库根目录的设计文档 plan.md（Variant System Transformation Plan）展开，完整讲解 screenshot-to-code 多模型并行生成变体（Variant）系统的架构诊断：现有"全量等待"式阻塞流程的成因、非阻塞化改造的 WebSocket 协议设计与前后端实现方案、分阶段实施计划、风险与收益度量。读完本文，你能理解该项目如何通过 variantComplete 逐变体完成信号、变体级状态机与混合状态策略，把"用户必须等最慢的模型"改造为"第一个完成的变体即可交互"，并掌握对应源码中的落地位置与验证路径。

一、问题诊断：现有变体系统的阻塞式生成

1.1 现状：NUM_VARIANTS 个变体并行生成，但全量等待

screenshot-to-code 的核心体验是：一次生成请求会并行调用多个不同的 AI 模型，产出多个候选实现（变体）供用户对比。plan.md 对该系统现状的诊断如下：

1. 阻塞式生成（Blocking Generation）：当用户创建或更新代码时，系统会并行生成 NUM_VARIANTS（当时为 2）个变体，每个变体使用不同的 AI 模型；
2. 全有或全无（All-or-Nothing）：WebSocket 连接保持打开，直到所有变体全部生成完毕；
3. 生成期间无法交互：用户必须等全部变体完成之后，才能选择变体、发起更新或开始新一轮生成。

用一张流程线图概括就是：

User Request → Generate All Variants → Wait for All → Close WS → Show Results → Allow Interaction

这里的"变体数量"是系统级配置。在 backend/config.py 中可以看到当前取值已经演进：

NUM_VARIANTS = 4
NUM_VARIANTS_VIDEO = 2

即常规生成为 4 个变体，视频模式为 2 个；而"update"（基于已有代码的修改）流程则在后端被固定为 2 个变体以压低延迟与成本——这一点可以从 backend/routes/generate_code.py 中 ModelSelectionStage.select_models 的 num_variants = 2 if generation_type == "update" else NUM_VARIANTS 得到印证。plan.md 写作时的 "2 variants" 是当时的配置值，架构语义不变：变体数量由配置驱动，全流程自动适配。

变体与模型的配对采用"模型轮转"（cycling）策略：根据可用的 API key 组合出一组模型列表，再按 models[i % len(models)] 循环填充到每个变体。当前实现位于 backend/routes/generate_code.py 的 ModelSelectionStage._get_variant_models：三把 key 齐备时使用 ALL_KEYS_MODELS_DEFAULT 等预定义组合，单 key 时只用该 provider 的模型；若 models = [A, B] 且变体数为 5，结果就是 [A, B, A, B, A]。这一机制在 design-docs/variant-system.md 中也有完整描述。

1.2 关键组件定位：阻塞发生在哪一层

plan.md 对阻塞成因做了精确定位，我们逐层用当前源码核对：

后端 backend/routes/generate_code.py（plan.md 写作时的行号对应当时版本，这里给出当前代码中对应的结构）：

• 为所有变体创建并行任务，并用 asyncio.gather() 等待全部任务完成——对应 AgenticGenerationStage.process_variants：

# backend/routes/generate_code.py（当前实现）
results = await asyncio.gather(*tasks, return_exceptions=True)
variant_completions: Dict[int, str] = {}
for index, result in enumerate(results):
    if isinstance(result, BaseException):
        print(f"Variant {index + 1} failed: {result}")
        continue
    if result:
        variant_completions[index] = result

• 每条消息携带 variantIndex 字段路由到正确的变体——这一设计在 plan.md 提出之前就存在，是变体系统的基础设施；
• WebSocket 在整轮生成结束后才关闭——对应 WebSocketSetupMiddleware 中 finally: await context.ws_comm.close() 的兜底逻辑（backend/routes/generate_code.py）。

前端：

• frontend/src/generateCode.ts：WebSocket 客户端。注意其 close 事件处理（frontend/src/generateCode.ts）——只有当连接以正常状态码（1000）关闭时才回调 onComplete()，这正是"全量等待"的前端侧表现：onComplete 是全局完成信号，而非逐变体信号；
• frontend/src/App.tsx：生成期间把全局应用状态置为 CODING，阻塞交互 UI；
• frontend/src/components/variants/Variants.tsx：只有在生成全部完成后才向用户展示变体列表。

这三点合起来构成了 plan.md 描述的"Current Flow"：用户在 CODING 状态下被完全锁死，直到所有变体完成、WebSocket 关闭、前端收到 onComplete。

二、目标架构：非阻塞的逐变体交付

plan.md 提出的新流程是：

User Request → Start Generation → Show First Complete → User Can Interact → Cancel Others if Needed

核心思想：把"整批完成"这一粗粒度信号，细化为"单个变体完成"的细粒度信号，让第一个跑完的模型立刻交付给用户，其余变体继续后台生成、可按需取消。

2.1 WebSocket 协议增强：新增 variantComplete 消息类型

plan.md 对后端协议提出的改动：

• 新增消息类型 "variantComplete"，用于通知单个变体已完成；
• 第一个变体完成后保持 WebSocket 打开；
• 支持在生成过程中取消指定变体；
• 为每个变体追踪状态：pending、generating、complete、failed、cancelled。

对照当前仓库，这一部分已经落地：backend/routes/generate_code.py 中的 MessageType 类型定义包含 variantComplete 与 variantError：

# WebSocket message types
MessageType = Literal[
    "chunk",
    "status",
    "setCode",
    "error",
    "variantComplete",
    "variantError",
    "variantCount",
    "variantModels",
    "thinking",
    "assistant",
    "toolStart",
    "toolResult",
]

且实际比 plan 更丰富：variantCount 让前端动态获知本轮变体数量（由 StatusBroadcastMiddleware 发出，同时先广播每路 "Generating code..." 状态），variantModels 在调试开关打开时下发每路使用的模型名。前端侧的完整类型定义见 frontend/src/generateCode.ts：

type WebSocketResponse = {
  type:
    | "chunk"
    | "status"
    | "setCode"
    | "error"
    | "variantComplete"
    | "variantError"
    | "variantCount"
    | "variantModels"
    | "thinking"
    | "assistant"
    | "toolStart"
    | "toolResult";
  value?: string;
  data?: any;
  eventId?: string;
  variantIndex: number;
};

每条消息都带 variantIndex，前端据此把 chunk、状态、完成信号分别归入对应变体。backend/tests/test_websocket_communicator.py 对发送与断连处理做了测试覆盖，可作为协议行为的验证入口。

2.2 生成逻辑改造：从"等全部"到"完成即推送"

plan.md 给出的改造示意——从"等待 asyncio.gather 全部返回"改为逐个 await 任务、完成一个就发一个信号：

# Instead of:
completions = await asyncio.gather(*tasks, return_exceptions=True)

# Use:
async def process_variants():
    for index, task in enumerate(tasks):
        try:
            completion = await task
            await send_message("variantComplete", index)
            # Allow frontend to interact immediately
        except Exception as e:
            await send_message("variantFailed", index)

需要说明的是：plan 中的这段代码是设计示意（串行逐路 await 仅为表达"每路完成立即通知"的意图），实际落地采用了等价但更清晰的写法——任务仍然并行创建，但每路变体在自己完成的瞬间独立推送 setCode + variantComplete，而不是由聚合层统一推送。当前 AgenticGenerationStage._run_variant 的收尾逻辑：

completion = await runner.run(model, prompt_messages)
if completion:
    await self.send_message("setCode", completion, index, None, None)
await self.send_message(
    "variantComplete",
    "Variant generation complete",
    index,
    None,
    None,
)

也就是说：第 1 路模型（通常是最快的）先返回，前端立即收到它的完整代码并完成信号；第 3、4 路仍在生成。asyncio.gather(*tasks, return_exceptions=True) 依然保留在 process_variants 中，但它的角色从"交付闸门"退化为"结果汇总"——最终用于判断是否所有变体都失败（全失败则 throw_error 并关闭连接）。每路变体的异常处理粒度也按计划细化：OpenAI 鉴权失败、模型不存在、配额超限分别给出针对性文案并通过 variantError 下发（backend/routes/generate_code.py），失败只影响对应的那一路，不拖累其他变体。

此外，整个 WebSocket 请求处理被重构为管道（Pipeline）模式：WebSocketSetupMiddleware → ParameterExtractionMiddleware → StatusBroadcastMiddleware → PromptCreationMiddleware → CodeGenerationMiddleware → PostProcessingMiddleware 依次装配（stream_code）。WebSocketCommunicator.send_message 内部维护 is_closed 标志并对客户端断连做静默降级（backend/routes/generate_code.py），保证了"连接保持打开、消息持续推送"的非阻塞模型在异常场景下依然稳健。

2.3 前端状态管理：变体级状态追踪

plan.md 提出在 Zustand store（project-store.ts）中加入变体级状态：

interface VariantState {
  code: string;
  status: 'pending' | 'generating' | 'complete' | 'failed' | 'cancelled';
  generationTime?: number;
}

落地实现采用了相近的建模：frontend/src/store/project-store.ts 中每路变体携带 code 与 status（默认值 "generating"），并暴露 updateVariantStatus(commitHash, variantIndex, status) 动作供 WebSocket 回调逐路更新。变体状态集合最终收敛为：

type VariantStatus = "generating" | "complete" | "cancelled";

即把 plan 中的 failed 语义并入 cancelled（失败的变体在 UI 上以"已取消"呈现），并配合 commit 结构管理——commit 即一次生成/修改的快照，内含 variants: Variant[] 与 selectedVariantIndex，历史通过 parentHash 链式回溯。这一整套 commit + 非阻塞变体的设计在 design-docs/commits-and-variants.md 中有权威描述。

UI 层面的关键决策是混合状态策略（Hybrid State）：全局 AppState 仍维持 INITIAL → CODING → CODE_READY 的粗粒度流转，但是否解锁更新交互由"全局就绪 或 当前选中变体已就绪"双条件决定：

// UI shows update interface when either:
const canUpdate =
  appState === AppState.CODE_READY ||           // All variants done
  isSelectedVariantComplete;                    // Selected variant done

这意味着用户不需要等待 CODE_READY，只要自己选中的那一路变体收到 variantComplete，就可以立即发起下一轮修改。plan.md 中"Enable 'Update' button when at least one variant is ready" 的诉求正是由此实现（对应 design-docs/commits-and-variants.md 中 Sidebar 的双条件渲染 appState === AppState.CODE_READY || isSelectedVariantComplete）。

2.4 变体列表的实时状态指示

plan.md 对视觉反馈的要求（生成中转圈、完成打勾、完成时轻微动效、耗时指示）在 frontend/src/components/variants/Variants.tsx 中逐条落地：

• 已完成的变体显示绿色指示点（variant.status === "complete" 时 statusColor = "bg-green-500"）；
• 生成中的变体显示 Spinner（variant.status === "generating" 分支）；
• 失败/取消的变体显示红色指示，用户仍可在其余变体间切换。

这对应 plan 中"UI Confusion"风险的缓解手段：清晰的状态指示让用户一眼分辨哪些结果可用、哪些还在路上。

2.5 WebSocket 客户端：不再依赖连接关闭

plan.md 要求前端"处理新的 variantComplete 消息类型、不等 WebSocket 关闭就放开交互、追踪仍在生成的变体"。当前 frontend/src/generateCode.ts 的回调接口完整实现了这一点——onVariantComplete、onVariantError、onVariantCount、onSetCode 等回调在 message 事件中逐条分发，而 onComplete 只保留给连接正常关闭这一时刻：

} else if (response.type === "variantComplete") {
  callbacks.onVariantComplete(response.variantIndex);
} else if (response.type === "variantError") {
  callbacks.onVariantError(response.variantIndex, response.value || "");
}

连接关闭时按关闭码区分语义：USER_CLOSE_WEB_SOCKET_CODE 视为用户主动取消、APP_ERROR_WEB_SOCKET_CODE 视为后端已知错误、其他非 1000 码视为连接故障，只有 1000 才触发 onComplete（frontend/src/generateCode.ts）。后端错误关闭码常量与前端保持同源（backend/ws/constants.py）。

2.6 更新流程中的变体取消

plan.md 定义的更新流程（Update Flow）是：用户带着未完成的变体发起新一轮修改时——取消剩余变体的生成、以 2 个变体开始新一轮生成、清掉旧的未完成变体。落地实现中：

• 新的生成请求会建立新的 WebSocket 连接，旧连接被前端关闭以释放资源，避免"僵尸变体"持续消耗模型配额；
• 旧 commit 中未完成的变体在状态上标记为 cancelled，保留在历史里但不参与当前会话；
• "update" 类型的生成在后端固定使用 2 个变体（如 1.2 节所述），天然降低了取消的成本。

需要谨慎说明的一点：design-docs/commits-and-variants.md 中展示了向前端→后端发送 cancel_variant 逐路取消消息的示意代码，但从当前 backend/routes/generate_code.py 看，服务端只在握手阶段 receive_params 一次，MessageType 枚举中也不存在 cancel_variant 类型——从源码结构看，逐路服务端取消协议并未随非阻塞改造一起落地，取消主要依赖"断开旧连接 + 状态置为 cancelled"的客户端策略。若未来要实现真正的中途中止（节省 token 花费），需要在协议中新增客户端控制消息并在 Agent 运行器中接入取消钩子。

三、用户体验改进：渐进式加载的完整闭环

plan.md 第 3 节列出的 UX 改进，与当前实现的对应关系如下：

plan.md 提出的改进	落地情况与验证位置
第一个变体就绪后立即展示	variantComplete → updateVariantStatus(..., "complete")，见 project-store.ts
未完成变体显示 skeleton/loading	Variants.tsx 中 generating 状态的 Spinner 分支
生成中可切换已完成的变体间	变体切换即时生效，切换目标仍在生成时显示 loading 直至完成
至少一个变体就绪即启用 Update	混合状态双条件 CODE_READY || isSelectedVariantComplete，见 design-docs/commits-and-variants.md
加载转圈 / 完成打勾 / 耗时指示	绿点/红点/Spinner 三态指示 + 生成耗时展示，见 Variants.tsx 与 frontend/src/components/agent/generation-time.ts

整体事件流（以 3 路变体为例）可以概括为：

Backend: Generate Variant 1 → "setCode" → "variantComplete"
Frontend: Update UI → 立即允许交互（若为选中变体）

Backend: Generate Variant 2 → "setCode" → "variantComplete"
Frontend: Update UI → 用户可切换到该变体

Backend: Generate Variant 3 → "variantError"
Frontend: 显示错误 → 该路标记为 cancelled，其余变体不受影响

这条链路上每一环都有可测试的行为边界：后端协议行为见 backend/tests/test_websocket_communicator.py 与 backend/tests/test_status_broadcast.py（variantCount/状态广播），变体选择逻辑见 backend/tests/test_model_selection.py，前端变体交互可用 frontend/jest.config.js 配置的测试环境对 project-store 行为做断言（frontend/src/store/project-store.test.ts）。

四、实施路线、风险与度量

4.1 四阶段实施计划

plan.md 给出的实施拆解，可作为理解本改造工作量分布的参考：

1. Phase 1: Backend Protocol（2–3 天）——改造 WebSocket 消息协议、实现逐变体完成追踪、加入取消机制；
2. Phase 2: Frontend State（2–3 天）——扩展 Zustand store 的变体状态、改造 WebSocket 客户端处理、调整 commit 结构；
3. Phase 3: UI Components（2–3 天）——Variants.tsx 渐进式展示、加载态与动画、Sidebar.tsx 即时交互；
4. Phase 4: Testing & Polish（2 天）——边界场景（全部变体失败等）、性能优化、用户测试。

从当前仓库状态看，Phase 1–3 的主体均已合入：协议消息类型、管道化生成流程、commit/变体状态管理与三态 UI 指示齐备；Phase 4 的"全部变体失败"边界在 CodeGenerationMiddleware 中以 len(variant_completions) == 0 → throw_error 收口。

4.2 风险与缓解

plan.md 识别的三大风险及其缓解策略值得保留为工程备忘：

1. 复杂度：状态更多、更难管理 → 谨慎的状态设计 + 完备测试（对应 VariantStatus 的最小状态集收敛与 commit 结构的规范化）；
2. 竞态条件：变体生成中用户发起更新 → 明确的取消逻辑与连接生命周期管理（新连接替换旧连接、旧连接关闭防泄漏）；
3. UI 困惑：用户可能不理解部分结果 → 清晰的状态指示（绿/红点、Spinner）。

4.3 收益与成功度量

plan.md 总结的收益：首次交互时间更快（用户可立即使用第一个变体）、体验更好（不必等慢模型拖累快模型）、迭代效率更高、灵活性（可随时放弃不想要的生成）。度量指标包括：

• 首次可交互时间下降约 50%（plan 的目标值，属于预期而非实测数据）；
• 生成速度的用户满意度提升；
• 生成过程中的流失率下降；
• 单会话内的迭代次数增加。

结合 design-docs/commits-and-variants.md 的补充总结，非阻塞架构的最终价值可归纳为五点：感知性能提升、多模型并行、灵活交互、资源效率（弃用未选变体）、以及优雅降级——部分变体失败时系统整体依然可用。

五、小结：从 plan 到落地的映射

回到 plan.md 本身，它回答了一个具体的架构问题：如何在保持多模型并行对比的前提下，把交互解锁时机从"整批完成"前移到"首个完成"。其方案三板斧——variantComplete 细粒度协议消息、变体级状态追踪、混合状态的双条件 UI 解锁——在 screenshot-to-code 当前代码中均可找到实现证据：

• 协议层：backend/routes/generate_code.py 的 MessageType 与 WebSocketCommunicator；
• 状态层：frontend/src/store/project-store.ts 的 VariantStatus / updateVariantStatus 与 commit 结构；
• 交互层：frontend/src/components/variants/Variants.tsx 的三态指示与 frontend/src/generateCode.ts 的逐路回调分发。

若你希望在本地体验并验证这套流程，可按 README.md 的后端/前端启动步骤配置至少一个模型 key（OpenAI、Anthropic 或 Gemini 三选一，key 越多每轮变体的模型组合越强），在 backend/.env 写入 key 后运行 poetry run uvicorn main:app --reload --port 7001 与 pnpm dev，即可在界面中观察多个变体逐一完成、随时切换并立即发起更新的非阻塞行为。深入阅读可沿 design-docs/variant-system.md、design-docs/commits-and-variants.md 与 backend/tests/test_model_selection.py 继续扩展。