From cbd211df4bb1d82beb4d45c973827efc2245c69d Mon Sep 17 00:00:00 2001 From: Kayshen-X Date: Sat, 21 Feb 2026 16:34:33 +0800 Subject: [PATCH] feat(ai): enhance chat functionality and orchestration capabilities - Introduce context optimization to manage chat history size and prevent unbounded growth. - Implement complexity assessment for design prompts to determine orchestration needs. - Add orchestrator functionality for parallel design generation, improving performance and responsiveness. - Update chat message parsing to include explicit status handling for orchestrator steps. - Refactor design generator to support orchestration and streamline context building. --- src/components/panels/ai-chat-panel.tsx | 5 +- src/components/panels/chat-message.tsx | 24 +- src/services/ai/ai-service.ts | 14 +- src/services/ai/ai-types.ts | 52 +++ src/services/ai/complexity-classifier.ts | 90 +++++ src/services/ai/context-optimizer.ts | 56 +++ src/services/ai/design-generator.ts | 24 +- src/services/ai/orchestrator-prompts.ts | 20 + src/services/ai/orchestrator.ts | 476 +++++++++++++++++++++++ 9 files changed, 750 insertions(+), 11 deletions(-) create mode 100644 src/services/ai/complexity-classifier.ts create mode 100644 src/services/ai/context-optimizer.ts create mode 100644 src/services/ai/orchestrator-prompts.ts create mode 100644 src/services/ai/orchestrator.ts diff --git a/src/components/panels/ai-chat-panel.tsx b/src/components/panels/ai-chat-panel.tsx index 341b47301..3c56baa8b 100644 --- a/src/components/panels/ai-chat-panel.tsx +++ b/src/components/panels/ai-chat-panel.tsx @@ -19,6 +19,7 @@ import { extractAndApplyDesign, extractAndApplyDesignModification } from '@/services/ai/design-generator' +import { trimChatHistory } from '@/services/ai/context-optimizer' import type { ChatMessage as ChatMessageType } from '@/services/ai/ai-types' import type { AIProviderType } from '@/types/agent-settings' import ClaudeLogo from '@/components/icons/claude-logo' @@ -227,12 +228,14 @@ function useChatHandlers() { } else { // --- CHAT MODE --- chatHistory.push({ role: 'user', content: fullUserMessage }) + // Trim history to prevent unbounded context growth + const trimmedHistory = trimChatHistory(chatHistory) // Resolve which provider the currently selected model belongs to const currentProvider = useAIStore.getState().modelGroups.find((g) => g.models.some((m) => m.value === model), )?.provider let chatThinking = false - for await (const chunk of streamChat(CHAT_SYSTEM_PROMPT, chatHistory, model, undefined, currentProvider)) { + for await (const chunk of streamChat(CHAT_SYSTEM_PROMPT, trimmedHistory, model, undefined, currentProvider)) { if (chunk.type === 'thinking') { // Show a brief thinking indicator so the UI isn't stuck on empty if (!chatThinking && !accumulated) { diff --git a/src/components/panels/chat-message.tsx b/src/components/panels/chat-message.tsx index dd85cf18c..415bb895d 100644 --- a/src/components/panels/chat-message.tsx +++ b/src/components/panels/chat-message.tsx @@ -49,16 +49,22 @@ function stripToolCallXml(text: string): string { export interface ParsedStep { title: string content: string + /** Explicit status from orchestrator steps (undefined for normal steps) */ + status?: 'pending' | 'streaming' | 'done' | 'error' } export function parseStepBlocks(text: string, isStreaming?: boolean): ParsedStep[] { - const stepRegex = /]*title="([^"]+)")?[^>]*>([\s\S]*?)<\/step>/gi + const stepRegex = /]*)>([\s\S]*?)<\/step>/gi const parsed: ParsedStep[] = [] let match: RegExpExecArray | null while ((match = stepRegex.exec(text)) !== null) { + const attrs = match[1] + const titleMatch = attrs.match(/title="([^"]+)"/) + const statusMatch = attrs.match(/status="([^"]+)"/) parsed.push({ - title: (match[1] ?? 'Processing').trim() || 'Processing', + title: (titleMatch?.[1] ?? 'Processing').trim() || 'Processing', + status: (statusMatch?.[1] as ParsedStep['status']) ?? undefined, content: (match[2] ?? '').trim(), }) } @@ -68,9 +74,11 @@ export function parseStepBlocks(text: string, isStreaming?: boolean): ParsedStep if (isStreaming && lastOpen > lastClose) { const partial = text.slice(lastOpen) const titleMatch = partial.match(/title="([^"]+)"/i) + const statusMatch = partial.match(/status="([^"]+)"/i) const contentStart = partial.indexOf('>') parsed.push({ title: (titleMatch?.[1] ?? 'Design').trim() || 'Design', + status: (statusMatch?.[1] as ParsedStep['status']) ?? undefined, content: contentStart >= 0 ? partial @@ -149,7 +157,17 @@ export function buildPipelineProgress( return steps.map((s) => ({ label: s.title, done: true, active: false })) } - // Map each step to done/active/pending based on completed JSON blocks. + // If steps have explicit status (orchestrator mode), use that directly + const hasExplicitStatus = steps.some((s) => s.status !== undefined) + if (hasExplicitStatus) { + return steps.map((s) => ({ + label: s.title, + done: s.status === 'done', + active: s.status === 'streaming', + })) + } + + // Fallback: Map each step to done/active/pending based on completed JSON blocks. // Step[i] is done when jsonBlockCount > i. // The step at jsonBlockCount is active (currently being generated). return steps.map((s, index) => { diff --git a/src/services/ai/ai-service.ts b/src/services/ai/ai-service.ts index 6a3c6cd1d..e64e0dfe6 100644 --- a/src/services/ai/ai-service.ts +++ b/src/services/ai/ai-service.ts @@ -7,6 +7,12 @@ const DEFAULT_STREAM_NO_TEXT_TIMEOUT_MS = 75_000 interface StreamChatOptions { hardTimeoutMs?: number noTextTimeoutMs?: number + /** + * Whether thinking events should reset the no-text timeout. + * Default: true (backward compatible). Set to false for fast calls + * where thinking should NOT prevent the no-text timeout from firing. + */ + thinkingResetsTimeout?: boolean } /** @@ -22,6 +28,7 @@ export async function* streamChat( ): AsyncGenerator { const hardTimeoutMs = Math.max(10_000, options?.hardTimeoutMs ?? DEFAULT_STREAM_HARD_TIMEOUT_MS) const noTextTimeoutMs = Math.max(10_000, options?.noTextTimeoutMs ?? DEFAULT_STREAM_NO_TEXT_TIMEOUT_MS) + const thinkingResetsTimeout = options?.thinkingResetsTimeout ?? true const controller = new AbortController() let abortReason: 'hard_timeout' | 'no_text_timeout' | null = null @@ -113,8 +120,11 @@ export async function* streamChat( continue } - // Any non-empty content (text or thinking) counts as activity - if ((chunk.type === 'text' || chunk.type === 'thinking') && chunk.content.trim().length > 0) { + // Any non-empty text counts as activity; thinking only resets + // the timeout when thinkingResetsTimeout is true (default). + if (chunk.type === 'text' && chunk.content.trim().length > 0) { + resetActivityTimeout() + } else if (chunk.type === 'thinking' && chunk.content.trim().length > 0 && thinkingResetsTimeout) { resetActivityTimeout() } diff --git a/src/services/ai/ai-types.ts b/src/services/ai/ai-types.ts index 37ae8212b..9aa0ac6e6 100644 --- a/src/services/ai/ai-types.ts +++ b/src/services/ai/ai-types.ts @@ -27,3 +27,55 @@ export interface AIStreamChunk { type: 'text' | 'thinking' | 'done' | 'error' | 'ping' content: string } + +// --------------------------------------------------------------------------- +// Orchestrator types — used for parallel sub-agent design generation +// --------------------------------------------------------------------------- + +/** A sub-task produced by the orchestrator (lightweight — only spatial info) */ +export interface SubTask { + /** Unique ID for this sub-task, e.g. "sidebar", "header" */ + id: string + /** Human-readable label for progress UI */ + label: string + /** Spatial region allocated on the canvas */ + region: { width: number; height: number } + /** ID prefix for all nodes generated by this sub-agent (assigned at runtime) */ + idPrefix: string + /** Parent frame ID to insert into (assigned at runtime) */ + parentFrameId: string | null +} + +/** Plan produced by the orchestrator (lightweight — only structure) */ +export interface OrchestratorPlan { + rootFrame: { + id: string + name: string + width: number + height: number + layout?: 'none' | 'vertical' | 'horizontal' + gap?: number + fill?: Array<{ type: string; color: string }> + } + subtasks: SubTask[] +} + +/** Progress state for the orchestrated generation */ +export interface OrchestrationProgress { + phase: 'planning' | 'generating' | 'merging' | 'done' | 'error' + subtasks: Array<{ + id: string + label: string + status: 'pending' | 'streaming' | 'done' | 'error' + nodeCount: number + }> + totalNodes: number +} + +/** Result from a single sub-agent */ +export interface SubAgentResult { + subtaskId: string + nodes: import('@/types/pen').PenNode[] + rawResponse: string + error?: string +} diff --git a/src/services/ai/complexity-classifier.ts b/src/services/ai/complexity-classifier.ts new file mode 100644 index 000000000..9f26e824f --- /dev/null +++ b/src/services/ai/complexity-classifier.ts @@ -0,0 +1,90 @@ +/** + * Heuristic complexity classifier for design prompts. + * Determines whether a prompt should be routed through + * the orchestrator (parallel sub-agents) or handled directly. + * + * No API calls — runs in <1ms. False positives are cheap + * (one extra planning call), false negatives are safe + * (fall back to existing single-call path). + */ + +// Structural keywords indicating distinct spatial sections +const SECTION_KEYWORDS = [ + 'sidebar', 'header', 'footer', 'nav', 'navigation', + 'hero', 'stats', 'statistics', 'chart', 'table', + 'form', 'modal', 'dialog', 'card section', 'cards section', + 'profile', 'settings', 'feed', 'content area', + 'grid', 'gallery', 'carousel', 'tabs', + 'search bar', 'filter', 'toolbar', 'breadcrumb', + 'notification', 'calendar', 'timeline', + // Chinese equivalents + '侧边栏', '头部', '底部', '导航', '英雄区', + '统计', '图表', '表格', '表单', '弹窗', + '卡片区', '个人资料', '设置', '时间线', +] + +// Full-page keywords that almost always produce complex output +const COMPLEX_PAGE_KEYWORDS = [ + 'dashboard', 'landing page', 'homepage', 'e-commerce', + 'admin panel', 'social media', 'email client', + 'settings page', 'analytics', 'crm', 'erp', + 'portfolio', 'blog', 'marketplace', 'checkout', + // Chinese + '仪表盘', '着陆页', '首页', '电商', '管理后台', + '社交', '邮件', '分析页', +] + +// Simple single-component keywords (suppress orchestration) +const SIMPLE_KEYWORDS = [ + 'button', 'input', 'avatar', 'badge', 'tooltip', + 'toggle', 'switch', 'checkbox', 'radio', + 'tag', 'chip', 'divider', 'spinner', 'icon', + '按钮', '输入框', '头像', '标签', '开关', +] + +/** Minimum sections required to trigger orchestrator */ +const COMPLEXITY_THRESHOLD = 3 + +export interface ComplexityAssessment { + isComplex: boolean + estimatedSections: number + reason: string +} + +export function assessComplexity(prompt: string): ComplexityAssessment { + const lower = prompt.toLowerCase() + + // Count distinct section keywords mentioned + const mentionedSections = SECTION_KEYWORDS.filter((kw) => lower.includes(kw)) + const isFullPage = COMPLEX_PAGE_KEYWORDS.some((kw) => lower.includes(kw)) + const isSimpleComponent = + SIMPLE_KEYWORDS.some((kw) => lower.includes(kw)) && mentionedSections.length <= 1 + + // Count conjunctions as section indicators + const conjunctions = (lower.match(/\band\b/g) || []).length + + (lower.match(/[,，、]/g) || []).length + + const estimatedSections = isFullPage + ? Math.max(mentionedSections.length, 4) + : mentionedSections.length + Math.min(Math.floor(conjunctions / 2), 2) + + if (isSimpleComponent) { + return { + isComplex: false, + estimatedSections: 1, + reason: 'Simple single component request', + } + } + + const isComplex = + estimatedSections >= COMPLEXITY_THRESHOLD || + (isFullPage && estimatedSections >= 2) + + return { + isComplex, + estimatedSections, + reason: isComplex + ? `Detected ${mentionedSections.length} sections (${mentionedSections.join(', ')})` + : `Below threshold: ${estimatedSections} sections`, + } +} diff --git a/src/services/ai/context-optimizer.ts b/src/services/ai/context-optimizer.ts new file mode 100644 index 000000000..300785497 --- /dev/null +++ b/src/services/ai/context-optimizer.ts @@ -0,0 +1,56 @@ +/** + * Context optimization utilities for AI chat. + * Prevents unbounded growth of chat history and context size. + */ + +const DEFAULT_MAX_MESSAGES = 10 +const DEFAULT_MAX_CHARS = 32_000 + +/** + * Sliding window for chat history. + * Keeps the most recent messages while respecting character limits. + * Always preserves the first user message for context continuity. + */ +export function trimChatHistory( + messages: T[], + maxMessages: number = DEFAULT_MAX_MESSAGES, + maxChars: number = DEFAULT_MAX_CHARS, +): T[] { + if (messages.length <= maxMessages) { + const totalChars = messages.reduce((sum, m) => sum + m.content.length, 0) + if (totalChars <= maxChars) return messages + } + + // Always keep the first user message for context continuity + const firstUser = messages.find((m) => m.role === 'user') + const recentMessages = messages.slice(-maxMessages) + + const window: T[] = [] + let charCount = 0 + + // Add first user message if it's not already in the recent window + if (firstUser && !recentMessages.includes(firstUser)) { + window.push(firstUser) + charCount += firstUser.content.length + } + + // Add recent messages, respecting char limit + for (const msg of recentMessages) { + const msgChars = msg.content.length + if (charCount + msgChars > maxChars) { + // Truncate this message to fit + const remaining = maxChars - charCount + if (remaining > 200) { + window.push({ + ...msg, + content: msg.content.slice(0, remaining) + '\n[...truncated...]', + } as T) + } + break + } + window.push(msg) + charCount += msgChars + } + + return window +} diff --git a/src/services/ai/design-generator.ts b/src/services/ai/design-generator.ts index b0f0c5b6e..1d13491a0 100644 --- a/src/services/ai/design-generator.ts +++ b/src/services/ai/design-generator.ts @@ -11,10 +11,13 @@ import { startNewAnimationBatch, resetAnimationState, } from './design-animation' +import { assessComplexity } from './complexity-classifier' +import { executeOrchestration } from './orchestrator' const DESIGN_STREAM_TIMEOUTS = { hardTimeoutMs: 180_000, noTextTimeoutMs: 60_000, + thinkingResetsTimeout: false, } // --------------------------------------------------------------------------- @@ -24,13 +27,13 @@ const DESIGN_STREAM_TIMEOUTS = { const generationRemappedIds = new Map() -function resetGenerationRemapping(): void { +export function resetGenerationRemapping(): void { generationRemappedIds.clear() } // Helper to find all complete JSON blocks in text -function extractJsonFromResponse(text: string): PenNode[] | null { +export function extractJsonFromResponse(text: string): PenNode[] | null { const parsedBlocks = extractAllJsonBlocks(text) .map((block) => tryParseNodes(block)) .filter(Boolean) as PenNode[][] @@ -97,7 +100,7 @@ function buildContextMessage(request: AIDesignRequest): string { } /** Build a concise summary of document variables for AI context. */ -function buildVariableContext( +export function buildVariableContext( variables?: Record, themes?: Record, ): string | null { @@ -159,7 +162,7 @@ interface StreamingNodeResult { * Uses brace-counting to detect complete objects before the block closes. * Each object is expected to have a `_parent` field for tree insertion. */ -function extractStreamingNodes( +export function extractStreamingNodes( text: string, processedOffset: number, ): { results: StreamingNodeResult[]; newOffset: number } { @@ -260,7 +263,7 @@ function parseJsonlToTree(text: string): PenNode[] | null { * Insert a single streaming node into the canvas with animation. * Handles root frame replacement and parent ID remapping. */ -function insertStreamingNode( +export function insertStreamingNode( node: PenNode, parentId: string | null, ): void { @@ -341,6 +344,17 @@ export async function generateDesign( animated?: boolean } ): Promise<{ nodes: PenNode[]; rawResponse: string }> { + // Route complex prompts through orchestrator for parallel generation + const { isComplex } = assessComplexity(request.prompt) + if (isComplex) { + try { + return await executeOrchestration(request, callbacks) + } catch (err) { + // Orchestrator failed — silently fall back to single-call generation + console.warn('Orchestrator failed, falling back to direct generation:', err) + } + } + const userMessage = buildContextMessage(request) let fullResponse = '' let streamingOffset = 0 // Tracks how far we've parsed in the streaming text diff --git a/src/services/ai/orchestrator-prompts.ts b/src/services/ai/orchestrator-prompts.ts new file mode 100644 index 000000000..d549cc5a5 --- /dev/null +++ b/src/services/ai/orchestrator-prompts.ts @@ -0,0 +1,20 @@ +/** + * Orchestrator prompt — ultra-lightweight, only splits into sections. + * No design details, no prompt rewriting. Just structure. + */ + +export const ORCHESTRATOR_PROMPT = `Split a UI request into 2-6 spatial sections. Output ONLY JSON, start with {. + +FORMAT: +{"rootFrame":{"id":"page","name":"Page","width":1200,"height":800,"layout":"vertical","fill":[{"type":"solid","color":"#16171B"}]},"subtasks":[{"id":"nav","label":"Navigation Bar","region":{"width":1200,"height":56}},{"id":"hero","label":"Hero Section","region":{"width":1200,"height":400}}]} + +RULES: +- Each subtask = one section with id, label, region (width+height). +- Regions tile to fill rootFrame. vertical = top-to-bottom. +- Mobile: 375x812. Desktop: 1200x800. +- NO explanation. NO markdown. JUST the JSON object.` + +export const ORCHESTRATOR_TIMEOUTS = { + hardTimeoutMs: 30_000, + noTextTimeoutMs: 20_000, +} diff --git a/src/services/ai/orchestrator.ts b/src/services/ai/orchestrator.ts new file mode 100644 index 000000000..375862b86 --- /dev/null +++ b/src/services/ai/orchestrator.ts @@ -0,0 +1,476 @@ +/** + * Orchestrator for parallel design generation. + * + * Flow: + * 1. Fast "architect" API call decomposes the prompt into spatial sub-tasks + * 2. Root frame is created on canvas + * 3. Multiple sub-agents execute in parallel, each streaming JSONL + * 4. Nodes are inserted to canvas in real-time with animation + * + * Falls back to single-call generation on any orchestrator failure. + */ + +import type { PenNode, FrameNode } from '@/types/pen' +import type { VariableDefinition } from '@/types/variables' +import type { + AIDesignRequest, + OrchestratorPlan, + OrchestrationProgress, + SubTask, + SubAgentResult, +} from './ai-types' +import { streamChat } from './ai-service' +import { DESIGN_GENERATOR_PROMPT } from './ai-prompts' +import { ORCHESTRATOR_PROMPT, ORCHESTRATOR_TIMEOUTS } from './orchestrator-prompts' +import { + extractStreamingNodes, + extractJsonFromResponse, + insertStreamingNode, + buildVariableContext, + resetGenerationRemapping, +} from './design-generator' +import { useDocumentStore } from '@/stores/document-store' +import { useHistoryStore } from '@/stores/history-store' +import { + resetAnimationState, + startNewAnimationBatch, +} from './design-animation' + +const SUB_AGENT_TIMEOUTS = { + hardTimeoutMs: 120_000, + noTextTimeoutMs: 45_000, + thinkingResetsTimeout: true, +} + +// --------------------------------------------------------------------------- +// Public API +// --------------------------------------------------------------------------- + +export async function executeOrchestration( + request: AIDesignRequest, + callbacks?: { + onApplyPartial?: (count: number) => void + onTextUpdate?: (text: string) => void + animated?: boolean + }, +): Promise<{ nodes: PenNode[]; rawResponse: string }> { + const animated = callbacks?.animated ?? false + + // -- Phase 1: Planning -- + callbacks?.onTextUpdate?.( + 'Analyzing design structure...', + ) + + const plan = await callOrchestrator(request.prompt, (thinking) => { + // Forward thinking progress to UI so user sees activity + const truncated = thinking.length > 200 + ? thinking.slice(-200) + '...' + : thinking + callbacks?.onTextUpdate?.( + `${truncated}`, + ) + }) + + // Assign ID prefixes + for (const st of plan.subtasks) { + st.idPrefix = st.id + st.parentFrameId = plan.rootFrame.id + } + + // Show planning done + all subtask steps as pending + emitProgress(plan, { + phase: 'generating', + subtasks: plan.subtasks.map((st) => ({ + id: st.id, label: st.label, status: 'pending' as const, nodeCount: 0, + })), + totalNodes: 0, + }, callbacks) + + // -- Phase 2: Setup canvas -- + resetGenerationRemapping() + + if (animated) { + resetAnimationState() + useHistoryStore.getState().startBatch(useDocumentStore.getState().document) + } + + // Insert root frame + const rootNode: FrameNode = { + id: plan.rootFrame.id, + type: 'frame', + name: plan.rootFrame.name, + x: 0, + y: 0, + width: plan.rootFrame.width, + height: plan.rootFrame.height, + layout: plan.rootFrame.layout ?? 'vertical', + gap: plan.rootFrame.gap ?? 0, + fill: (plan.rootFrame.fill as FrameNode['fill']) ?? [ + { type: 'solid', color: '#FFFFFF' }, + ], + children: [], + } + + insertStreamingNode(rootNode, null) + + // -- Phase 3: Parallel sub-agent execution -- + const progress: OrchestrationProgress = { + phase: 'generating', + subtasks: plan.subtasks.map((st) => ({ + id: st.id, + label: st.label, + status: 'pending' as const, + nodeCount: 0, + })), + totalNodes: 0, + } + + let results: SubAgentResult[] + try { + results = await executeSubAgentsSequentially( + plan, + request, + progress, + callbacks, + ) + } finally { + if (animated) { + useHistoryStore.getState().endBatch(useDocumentStore.getState().document) + } + } + + // -- Phase 4: Collect results -- + // Mark all completed subtasks as done in final progress + for (const entry of progress.subtasks) { + if (entry.status !== 'error') { + entry.status = 'done' + } + } + progress.phase = 'done' + emitProgress(plan, progress, callbacks) + + const allNodes: PenNode[] = [rootNode] + for (const r of results) { + allNodes.push(...r.nodes) + } + + // Build final rawResponse that includes step tags so the chat message + // shows the complete pipeline progress after streaming ends + const finalStepTags = buildFinalStepTags(plan, progress) + + return { nodes: allNodes, rawResponse: finalStepTags } +} + +// --------------------------------------------------------------------------- +// Orchestrator call — fast decomposition +// --------------------------------------------------------------------------- + +/** Max prompt length for the orchestrator — it only needs structure, not full detail */ +const MAX_ORCHESTRATOR_PROMPT_CHARS = 2000 + +async function callOrchestrator( + prompt: string, + onThinking?: (thinking: string) => void, +): Promise { + // Truncate long prompts — the orchestrator only needs high-level structure + const truncatedPrompt = prompt.length > MAX_ORCHESTRATOR_PROMPT_CHARS + ? prompt.slice(0, MAX_ORCHESTRATOR_PROMPT_CHARS) + '\n\n[... prompt truncated for planning — full details will be sent to each section agent]' + : prompt + + let rawResponse = '' + let thinkingContent = '' + + for await (const chunk of streamChat( + ORCHESTRATOR_PROMPT, + [{ role: 'user', content: truncatedPrompt }], + undefined, + { + ...ORCHESTRATOR_TIMEOUTS, + // Don't let thinking indefinitely extend the no-text timeout + thinkingResetsTimeout: false, + }, + )) { + if (chunk.type === 'text') { + rawResponse += chunk.content + } else if (chunk.type === 'thinking') { + thinkingContent += chunk.content + onThinking?.(thinkingContent) + } else if (chunk.type === 'error') { + throw new Error(`Orchestrator failed: ${chunk.content}`) + } + } + + const plan = parseOrchestratorResponse(rawResponse) + if (!plan) { + throw new Error('Failed to parse orchestrator plan') + } + + return plan +} + +function parseOrchestratorResponse(raw: string): OrchestratorPlan | null { + const trimmed = raw.trim() + + // Try direct parse + const plan = tryParsePlan(trimmed) + if (plan) return plan + + // Try extracting from code fences + const fenceMatch = trimmed.match(/```(?:json)?\s*\n?([\s\S]*?)\n?```/) + if (fenceMatch) { + const fenced = tryParsePlan(fenceMatch[1].trim()) + if (fenced) return fenced + } + + // Try extracting first { ... } block + const firstBrace = trimmed.indexOf('{') + const lastBrace = trimmed.lastIndexOf('}') + if (firstBrace >= 0 && lastBrace > firstBrace) { + const braced = tryParsePlan(trimmed.slice(firstBrace, lastBrace + 1)) + if (braced) return braced + } + + return null +} + +function tryParsePlan(text: string): OrchestratorPlan | null { + try { + const obj = JSON.parse(text) as Record + if (!obj.rootFrame || typeof obj.rootFrame !== 'object') return null + if (!Array.isArray(obj.subtasks) || obj.subtasks.length === 0) return null + + const rf = obj.rootFrame as Record + if (!rf.id || !rf.width || !rf.height) return null + + for (const st of obj.subtasks as Record[]) { + if (!st.id || !st.region) return null + } + + return obj as unknown as OrchestratorPlan + } catch { + return null + } +} + +// --------------------------------------------------------------------------- +// Parallel sub-agent execution +// --------------------------------------------------------------------------- + +async function executeSubAgentsSequentially( + plan: OrchestratorPlan, + request: AIDesignRequest, + progress: OrchestrationProgress, + callbacks?: { + onApplyPartial?: (count: number) => void + onTextUpdate?: (text: string) => void + animated?: boolean + }, +): Promise { + const results: SubAgentResult[] = [] + for (let i = 0; i < plan.subtasks.length; i++) { + const result = await executeSubAgent(plan.subtasks[i], plan, request, progress, i, callbacks) + results.push(result) + } + return results +} + +async function executeSubAgent( + subtask: SubTask, + plan: OrchestratorPlan, + request: AIDesignRequest, + progress: OrchestrationProgress, + index: number, + callbacks?: { + onApplyPartial?: (count: number) => void + onTextUpdate?: (text: string) => void + animated?: boolean + }, +): Promise { + const animated = callbacks?.animated ?? false + const progressEntry = progress.subtasks[index] + progressEntry.status = 'streaming' + emitProgress(plan, progress, callbacks) + + const userPrompt = buildSubAgentUserPrompt( + subtask, + request.prompt, + request.context?.variables, + request.context?.themes, + ) + + let rawResponse = '' + const nodes: PenNode[] = [] + let streamOffset = 0 + + try { + for await (const chunk of streamChat( + DESIGN_GENERATOR_PROMPT, + [{ role: 'user', content: userPrompt }], + undefined, + SUB_AGENT_TIMEOUTS, + )) { + if (chunk.type === 'text') { + rawResponse += chunk.content + + if (animated) { + const { results, newOffset } = extractStreamingNodes( + rawResponse, + streamOffset, + ) + if (results.length > 0) { + streamOffset = newOffset + startNewAnimationBatch() + + for (const { node, parentId } of results) { + // Enforce ID prefix + ensureIdPrefix(node, subtask.idPrefix) + if (parentId !== null) { + // Prefix the parent reference too + const prefixedParent = ensurePrefixStr( + parentId, + subtask.idPrefix, + ) + insertStreamingNode(node, prefixedParent) + } else { + // Sub-agent root → insert under the orchestrator root frame + insertStreamingNode(node, plan.rootFrame.id) + } + nodes.push(node) + progressEntry.nodeCount++ + progress.totalNodes++ + } + callbacks?.onApplyPartial?.(progress.totalNodes) + emitProgress(plan, progress, callbacks) + } + } + } else if (chunk.type === 'thinking') { + // Forward thinking progress so UI doesn't look stuck + emitProgress(plan, progress, callbacks) + } else if (chunk.type === 'error') { + progressEntry.status = 'error' + emitProgress(plan, progress, callbacks) + return { subtaskId: subtask.id, nodes, rawResponse, error: chunk.content } + } + } + + // Fallback: if streaming extraction found nothing, try batch extraction + if (nodes.length === 0 && rawResponse.trim().length > 0) { + const fallbackNodes = extractJsonFromResponse(rawResponse) + if (fallbackNodes && fallbackNodes.length > 0) { + startNewAnimationBatch() + for (const node of fallbackNodes) { + ensureIdPrefix(node, subtask.idPrefix) + insertStreamingNode(node, plan.rootFrame.id) + nodes.push(node) + progressEntry.nodeCount++ + progress.totalNodes++ + } + callbacks?.onApplyPartial?.(progress.totalNodes) + } + } + + progressEntry.status = 'done' + emitProgress(plan, progress, callbacks) + return { subtaskId: subtask.id, nodes, rawResponse } + } catch (err) { + const msg = err instanceof Error ? err.message : 'Unknown error' + progressEntry.status = 'error' + emitProgress(plan, progress, callbacks) + return { subtaskId: subtask.id, nodes, rawResponse, error: msg } + } +} + +// --------------------------------------------------------------------------- +// Sub-agent prompt builder +// --------------------------------------------------------------------------- + +function buildSubAgentUserPrompt( + subtask: SubTask, + originalPrompt: string, + variables?: Record, + themes?: Record, +): string { + const { region } = subtask + + // Extract a brief context line from the original prompt (first 200 chars) + const briefContext = originalPrompt.length > 200 + ? originalPrompt.slice(0, 200) + '...' + : originalPrompt + + let prompt = `Design: "${subtask.label}" +Context: ${briefContext} +Canvas: ${region.width}x${region.height}px +Root frame: id="${subtask.idPrefix}-root", width=${region.width}, height=${region.height} +All node IDs MUST start with "${subtask.idPrefix}-". +Generate ONLY this section, not the full page.` + + const varContext = buildVariableContext(variables, themes) + if (varContext) { + prompt += '\n\n' + varContext + } + + return prompt +} + +// --------------------------------------------------------------------------- +// ID namespace isolation +// --------------------------------------------------------------------------- + +function ensureIdPrefix(node: PenNode, prefix: string): void { + if (!node.id.startsWith(`${prefix}-`)) { + node.id = `${prefix}-${node.id}` + } +} + +function ensurePrefixStr(id: string, prefix: string): string { + if (id.startsWith(`${prefix}-`)) return id + return `${prefix}-${id}` +} + +// --------------------------------------------------------------------------- +// Progress emission — updates UI via tags +// --------------------------------------------------------------------------- + +function emitProgress( + plan: OrchestratorPlan, + progress: OrchestrationProgress, + callbacks?: { + onTextUpdate?: (text: string) => void + }, +): void { + if (!callbacks?.onTextUpdate) return + + // Always show "Planning layout" as done first + const planningStep = 'Analyzing design structure...' + + const subtaskSteps = plan.subtasks + .map((st, i) => { + const entry = progress.subtasks[i] + const status = entry.status === 'streaming' ? 'streaming' + : entry.status === 'done' ? 'done' + : entry.status === 'error' ? 'error' + : 'pending' + const nodeInfo = entry.nodeCount > 0 ? ` (${entry.nodeCount} elements)` : '' + return `` + }) + .join('\n') + + callbacks.onTextUpdate(`${planningStep}\n${subtaskSteps}`) +} + +/** Build step tags for the final rawResponse (shown in message after streaming ends) */ +function buildFinalStepTags( + plan: OrchestratorPlan, + progress: OrchestrationProgress, +): string { + const planningStep = 'Analyzing design structure...' + const subtaskSteps = plan.subtasks + .map((st, i) => { + const entry = progress.subtasks[i] + const status = entry.status + const nodeInfo = entry.nodeCount > 0 ? ` (${entry.nodeCount} elements)` : '' + return `` + }) + .join('\n') + return `${planningStep}\n${subtaskSteps}` +}