# AI Chat Builder Implementation Plan > **For Claude:** REQUIRED SUB-SKILL: Use superpowers:executing-plans to implement this plan task-by-task. **Goal:** Build a conversational AI flow builder where the AI conducts a multi-phase interview as a senior MSP engineer, progressively building a TreeStructure through natural dialogue. **Architecture:** New chat-based builder coexists alongside the existing wizard builder. Backend adds a new model, service, and endpoint module in `core/` and `api/endpoints/`. Frontend adds a split-panel page with chat left / tree preview right, a new Zustand store, and new components in `components/ai-chat/`. Reuses existing AI provider abstraction, quota system, and tree validation. **Tech Stack:** FastAPI, SQLAlchemy 2.0 (async), Alembic, Pydantic v2 (backend); React 19, Zustand, TypeScript, Tailwind CSS (frontend). AI via existing `AIProvider` abstraction (Gemini/Anthropic). **Design doc:** `docs/plans/2026-02-27-ai-chat-builder-design.md` --- ## Phase 1: Backend Foundation ### Task 1: Database Model — `AIChatSession` **Files:** - Create: `backend/app/models/ai_chat_session.py` - Modify: `backend/app/models/__init__.py` (lines 29-30, 68-69) **Step 1: Create the model file** ```python # backend/app/models/ai_chat_session.py """AI Chat Builder session tracking. Stores conversational flow builder state across the multi-phase interview. Sessions expire after 24 hours. """ import uuid from datetime import datetime, timezone from typing import Optional, Any from sqlalchemy import String, DateTime, ForeignKey, Integer from sqlalchemy.orm import Mapped, mapped_column from sqlalchemy.dialects.postgresql import UUID, JSONB from app.core.database import Base class AIChatSession(Base): __tablename__ = "ai_chat_sessions" id: Mapped[uuid.UUID] = mapped_column( UUID(as_uuid=True), primary_key=True, default=uuid.uuid4 ) user_id: Mapped[uuid.UUID] = mapped_column( UUID(as_uuid=True), ForeignKey("users.id", ondelete="CASCADE"), nullable=False, index=True, ) account_id: Mapped[uuid.UUID] = mapped_column( UUID(as_uuid=True), ForeignKey("accounts.id", ondelete="CASCADE"), nullable=False, index=True, ) status: Mapped[str] = mapped_column( String(20), nullable=False, default="active", comment="active | completed | abandoned", ) current_phase: Mapped[str] = mapped_column( String(20), nullable=False, default="scoping", comment="scoping | discovery | enrichment | review | generation", ) flow_type: Mapped[str] = mapped_column( String(20), nullable=False, comment="troubleshooting | procedural", ) conversation_history: Mapped[list[dict[str, Any]]] = mapped_column( JSONB, nullable=False, default=list ) working_tree: Mapped[Optional[dict[str, Any]]] = mapped_column( JSONB, nullable=True ) tree_metadata: Mapped[dict[str, Any]] = mapped_column( JSONB, nullable=False, default=dict ) provider_used: Mapped[Optional[str]] = mapped_column( String(20), nullable=True ) message_count: Mapped[int] = mapped_column( Integer, nullable=False, default=0 ) total_input_tokens: Mapped[int] = mapped_column( Integer, nullable=False, default=0 ) total_output_tokens: Mapped[int] = mapped_column( Integer, nullable=False, default=0 ) generated_tree_id: Mapped[Optional[uuid.UUID]] = mapped_column( UUID(as_uuid=True), ForeignKey("trees.id", ondelete="SET NULL"), nullable=True, ) expires_at: Mapped[datetime] = mapped_column( DateTime(timezone=True), nullable=False ) created_at: Mapped[datetime] = mapped_column( DateTime(timezone=True), default=lambda: datetime.now(timezone.utc) ) updated_at: Mapped[datetime] = mapped_column( DateTime(timezone=True), default=lambda: datetime.now(timezone.utc), onupdate=lambda: datetime.now(timezone.utc), ) ``` **Step 2: Register the model in `__init__.py`** Add import after line 30 (`from .ai_usage import AIUsage`): ```python from .ai_chat_session import AIChatSession ``` Add to `__all__` after line 69 (`"AIUsage"`): ```python "AIChatSession", ``` **Step 3: Create Alembic migration** Run: `cd backend && alembic revision -m "add ai_chat_sessions table"` Then write the migration manually (do NOT use `--autogenerate` — see CLAUDE.md lesson #17): ```python """add ai_chat_sessions table""" from alembic import op import sqlalchemy as sa from sqlalchemy.dialects.postgresql import UUID, JSONB def upgrade() -> None: op.create_table( "ai_chat_sessions", sa.Column("id", UUID(as_uuid=True), primary_key=True), sa.Column("user_id", UUID(as_uuid=True), sa.ForeignKey("users.id", ondelete="CASCADE"), nullable=False, index=True), sa.Column("account_id", UUID(as_uuid=True), sa.ForeignKey("accounts.id", ondelete="CASCADE"), nullable=False, index=True), sa.Column("status", sa.String(20), nullable=False, server_default="active"), sa.Column("current_phase", sa.String(20), nullable=False, server_default="scoping"), sa.Column("flow_type", sa.String(20), nullable=False), sa.Column("conversation_history", JSONB, nullable=False, server_default="[]"), sa.Column("working_tree", JSONB, nullable=True), sa.Column("tree_metadata", JSONB, nullable=False, server_default="{}"), sa.Column("provider_used", sa.String(20), nullable=True), sa.Column("message_count", sa.Integer, nullable=False, server_default="0"), sa.Column("total_input_tokens", sa.Integer, nullable=False, server_default="0"), sa.Column("total_output_tokens", sa.Integer, nullable=False, server_default="0"), sa.Column("generated_tree_id", UUID(as_uuid=True), sa.ForeignKey("trees.id", ondelete="SET NULL"), nullable=True), sa.Column("expires_at", sa.DateTime(timezone=True), nullable=False), sa.Column("created_at", sa.DateTime(timezone=True), server_default=sa.func.now()), sa.Column("updated_at", sa.DateTime(timezone=True), server_default=sa.func.now()), ) def downgrade() -> None: op.drop_table("ai_chat_sessions") ``` **Step 4: Run migration** Run: `cd backend && alembic upgrade head` Expected: Migration applied successfully, `ai_chat_sessions` table created. **Step 5: Commit** ```bash git add backend/app/models/ai_chat_session.py backend/app/models/__init__.py backend/alembic/versions/*ai_chat_sessions* git commit -m "feat: add ai_chat_sessions database model and migration" ``` --- ### Task 2: Pydantic Schemas **Files:** - Create: `backend/app/schemas/ai_chat.py` **Step 1: Create the schemas file** ```python # backend/app/schemas/ai_chat.py """Pydantic schemas for the AI Chat Builder.""" from typing import Any, Literal, Optional from uuid import UUID from pydantic import BaseModel, Field # ── Requests ── class AIChatStartRequest(BaseModel): """Start a new chat builder session.""" flow_type: Literal["troubleshooting", "procedural"] = Field( ..., description="Type of flow to build" ) class AIChatMessageRequest(BaseModel): """Send a user message in a chat session.""" content: str = Field(..., min_length=1, max_length=5000) class AIChatImportRequest(BaseModel): """Import generated tree with optional metadata overrides.""" name: Optional[str] = Field(None, min_length=1, max_length=255) description: Optional[str] = Field(None, max_length=2000) category_id: Optional[UUID] = None tags: list[str] = Field(default_factory=list) # ── Responses ── class AIChatStartResponse(BaseModel): """Response after creating a chat session.""" session_id: UUID greeting: str current_phase: str class AIChatMessageResponse(BaseModel): """Response after sending a message.""" content: str current_phase: str working_tree: Optional[dict[str, Any]] = None tree_metadata: Optional[dict[str, Any]] = None class AIChatSessionResponse(BaseModel): """Full session state for resume.""" session_id: UUID status: str current_phase: str flow_type: str conversation_history: list[dict[str, Any]] working_tree: Optional[dict[str, Any]] = None tree_metadata: Optional[dict[str, Any]] = None message_count: int generated_tree: Optional[dict[str, Any]] = None class AIChatGenerateResponse(BaseModel): """Response with the final generated tree.""" tree_structure: dict[str, Any] tree_metadata: dict[str, Any] status: str class AIChatImportResponse(BaseModel): """Response after importing tree to editor.""" tree_id: UUID tree_type: str ``` **Step 2: Commit** ```bash git add backend/app/schemas/ai_chat.py git commit -m "feat: add Pydantic schemas for AI chat builder" ``` --- ### Task 3: Chat Service — System Prompt & Conversation Loop **Files:** - Create: `backend/app/core/ai_chat_service.py` This is the largest file. It contains the system prompt, response parsing, and the three main functions. **Step 1: Create the service file** ```python # backend/app/core/ai_chat_service.py """AI Chat Builder service. Manages the conversational flow builder: system prompt construction, message exchange with AI provider, and response parsing (extracting tree updates, phase transitions, and metadata from structured markers). """ import json import logging import re import uuid from datetime import datetime, timezone, timedelta from typing import Any, Optional from sqlalchemy import select from sqlalchemy.ext.asyncio import AsyncSession from app.core.ai_provider import get_ai_provider from app.core.ai_tree_validator import validate_generated_tree from app.core.config import settings from app.models.ai_chat_session import AIChatSession logger = logging.getLogger(__name__) # ── Cost estimation ── COST_PER_INPUT_TOKEN = 1.0 / 1_000_000 COST_PER_OUTPUT_TOKEN = 5.0 / 1_000_000 # ── Max messages per session ── MAX_MESSAGES_FREE = 10 MAX_MESSAGES_PAID = 25 # ── System Prompt ── ROLE_PERSONA = """You are a senior IT engineer embedded in ResolutionFlow, a troubleshooting platform for MSP (Managed Service Provider) engineers. You have 15+ years of hands-on experience across Windows Server, Active Directory, Entra ID/Azure AD, Microsoft 365, networking (DNS, DHCP, routing, VPN, firewalls), virtualization (Hyper-V, VMware), security, backup/DR, and cloud infrastructure. Your job is to help engineers build troubleshooting decision trees by interviewing them about a problem space. You are NOT a generic assistant. You are a colleague who has seen these issues hundreds of times and knows the optimal diagnostic order. CRITICAL BEHAVIORS: - Act as a senior engineer, not a chatbot. Use your domain knowledge to SUGGEST diagnostic steps, not just record what the user says. - When the user describes a problem area, demonstrate understanding by naming specific sub-categories, common causes, and relevant tools. - Challenge assumptions constructively: "Before we go down that path, have you considered checking X first? In my experience, that resolves 60% of these cases." - Capture SPECIFIC commands with exact syntax. Not "check the service" but "Get-Service ADSync | Select-Object Status, StartType". - Include expected outcomes for every action: what does success look like? - Surface edge cases proactively: "What about multi-forest environments?" or "Does this change if they have conditional access policies?" - Explain WHY the diagnostic order matters: "We check connectivity before auth because a network issue masquerades as an auth failure." - Ask ONE focused question at a time. Do not overwhelm with multiple questions. - Use plain, collegial language. Sound like a colleague, not a form.""" SCHEMA_CONTEXT = """ TREESTRUCTURE SCHEMA — This is what you are building: The tree is a recursive JSON structure. Each node has a "type" field: 1. decision — A diagnostic question with branching options Required: id (string), type ("decision"), question (string), options (array), children (array) Optional: help_text (string) Each option: { id (string), label (string), next_node_id (string — must match a child's id) } 2. action — A step the engineer performs Required: id (string), type ("action"), title (string), description (string) Optional: commands (string array — exact CLI/PowerShell syntax), expected_outcome (string), help_text (string), next_node_id (string — ID of the next node to navigate to) 3. solution — A resolution endpoint Required: id (string), type ("solution"), title (string), description (string) Optional: resolution_steps (string array) STRUCTURAL RULES: - Root node MUST be type "decision" - Decision nodes contain their children in the "children" array - Each decision option's next_node_id must reference a child node's id - Action nodes use next_node_id to chain to the next step (NOT children) - Solution nodes are terminal — no next_node_id or children - All IDs must be unique strings (use descriptive slugs like "check-service-status") """ INTERVIEW_PROTOCOL = """ INTERVIEW PHASES — Follow this progression: PHASE 1 - SCOPING (current_phase: scoping): Ask broad questions to understand the problem domain and scope: - What type of issue is this flow for? - Who is the target audience? (Tier 1 help desk, Tier 2, Tier 3?) - What environment assumptions? (On-prem, hybrid, specific vendors?) Demonstrate domain expertise immediately. If the user says "Azure AD Sync failures," show understanding: "Are you primarily seeing password hash sync issues, object attribute sync failures, or full directory sync errors?" DO NOT emit [TREE_UPDATE] during scoping. You are still understanding the problem. PHASE 2 - DISCOVERY (current_phase: discovery): Work through the troubleshooting logic branch by branch: - Establish the first diagnostic question (the root decision node) - For each branch, ask what the engineer would check next - Suggest checks the user might not have considered - Capture specific commands, tools, and procedures EMIT [TREE_UPDATE] ONLY when you and the user have agreed on a concrete node — a decision with clear options, or an action with a specific command. If you are asking a question, you are NOT updating the tree. PHASE 3 - ENRICHMENT (current_phase: enrichment): Circle back to enrich existing nodes: - Add exact PowerShell/CLI commands with syntax - Add help text with relevant documentation links - Add expected outcomes for action nodes - Suggest edge cases needing additional branches EMIT [TREE_UPDATE] when enriching existing nodes or adding edge case branches. PHASE 4 - REVIEW (current_phase: review): Present a summary: - Total node count by type - Text outline of the flow structure - Flag any areas of uncertainty - Offer chance to add/remove/modify branches EMIT [TREE_UPDATE] only if the user requests structural changes. TRANSITION between phases by emitting [PHASE:phase_name] when the conversation naturally moves to the next stage. You decide when enough information has been gathered for each phase. """ RESPONSE_FORMAT = """ RESPONSE FORMAT: Your response is natural conversational text. When the tree structure changes, include structured markers that will be parsed by the system (the user will NOT see these markers): 1. Tree update (only when structure changes — see phase rules above): [TREE_UPDATE] {...valid TreeStructure JSON...} [/TREE_UPDATE] 2. Phase transition (when moving to next phase): [PHASE:discovery] 3. Metadata capture (when you learn the flow's name, description, or tags): [METADATA] {"name": "...", "description": "...", "tags": ["..."]} [/METADATA] IMPORTANT: - Include [TREE_UPDATE] sparingly. Only when concrete nodes are established or modified. - The tree update should be the COMPLETE working tree, not a diff. - Always include conversational text OUTSIDE the markers — never respond with only markers. """ def _build_system_prompt(flow_type: str) -> str: """Assemble the full system prompt for the chat builder.""" flow_context = ( "The user wants to build a TROUBLESHOOTING flow — a diagnostic decision tree " "that guides engineers through symptom identification, diagnostic checks, and " "resolution steps." if flow_type == "troubleshooting" else "The user wants to build a PROCEDURAL flow — a step-by-step process guide " "with phases, checklists, and verification steps." ) return f"{ROLE_PERSONA}\n\n{flow_context}\n\n{SCHEMA_CONTEXT}\n\n{INTERVIEW_PROTOCOL}\n\n{RESPONSE_FORMAT}" def _strip_markdown_fences(text: str) -> str: """Strip markdown code fences if the model wrapped its JSON response.""" text = text.strip() match = re.match(r"^```(?:json)?\s*([\s\S]*?)```$", text) if match: return match.group(1).strip() return text def _parse_ai_response(raw_response: str) -> dict[str, Any]: """Parse structured markers from AI response. Returns dict with: - content: str (conversational text with markers stripped) - tree_update: dict | None (parsed TreeStructure JSON) - phase: str | None (new phase name) - metadata: dict | None (name, description, tags) """ result: dict[str, Any] = { "content": raw_response, "tree_update": None, "phase": None, "metadata": None, } # Extract [TREE_UPDATE]...[/TREE_UPDATE] tree_match = re.search( r"\[TREE_UPDATE\]\s*([\s\S]*?)\s*\[/TREE_UPDATE\]", raw_response ) if tree_match: try: raw_json = _strip_markdown_fences(tree_match.group(1)) result["tree_update"] = json.loads(raw_json) except (json.JSONDecodeError, ValueError) as e: logger.warning("Failed to parse tree update JSON: %s", e) result["content"] = raw_response[: tree_match.start()] + raw_response[tree_match.end() :] # Extract [PHASE:name] phase_match = re.search(r"\[PHASE:(\w+)\]", result["content"]) if phase_match: result["phase"] = phase_match.group(1) result["content"] = result["content"][: phase_match.start()] + result["content"][phase_match.end() :] # Extract [METADATA]...[/METADATA] meta_match = re.search( r"\[METADATA\]\s*([\s\S]*?)\s*\[/METADATA\]", result["content"] ) if meta_match: try: raw_json = _strip_markdown_fences(meta_match.group(1)) result["metadata"] = json.loads(raw_json) except (json.JSONDecodeError, ValueError) as e: logger.warning("Failed to parse metadata JSON: %s", e) result["content"] = result["content"][: meta_match.start()] + result["content"][meta_match.end() :] # Clean up extra whitespace from marker removal result["content"] = re.sub(r"\n{3,}", "\n\n", result["content"]).strip() return result # ── Main Service Functions ── async def start_chat_session( flow_type: str, user_id: uuid.UUID, account_id: uuid.UUID, db: AsyncSession, ) -> tuple[AIChatSession, str]: """Create a chat session and return the AI's opening greeting. Returns (session, greeting_text). """ session = AIChatSession( user_id=user_id, account_id=account_id, flow_type=flow_type, expires_at=datetime.now(timezone.utc) + timedelta(hours=settings.AI_CONVERSATION_TTL_HOURS), ) db.add(session) await db.flush() # Build system prompt and get opening message system_prompt = _build_system_prompt(flow_type) primer = ( f"I want to build a {flow_type} flow. Help me get started." ) provider = get_ai_provider() provider_name = settings.AI_PROVIDER messages = [{"role": "user", "content": primer}] response_text, input_tokens, output_tokens = await provider.generate_json( system_prompt=system_prompt, messages=messages, max_tokens=1500, ) # Parse response (greeting shouldn't have tree updates, but handle gracefully) parsed = _parse_ai_response(response_text) # Store conversation history now_iso = datetime.now(timezone.utc).isoformat() session.conversation_history = [ {"role": "user", "content": primer, "timestamp": now_iso, "hidden": True}, {"role": "assistant", "content": parsed["content"], "timestamp": now_iso}, ] session.provider_used = provider_name session.message_count = 1 session.total_input_tokens = input_tokens session.total_output_tokens = output_tokens if parsed["metadata"]: session.tree_metadata = parsed["metadata"] return session, parsed["content"] async def send_message( session: AIChatSession, user_message: str, db: AsyncSession, ) -> tuple[str, Optional[dict], Optional[str], Optional[dict]]: """Send a user message and get AI response. Returns (ai_content, working_tree_update, new_phase, metadata_update). """ system_prompt = _build_system_prompt(session.flow_type) # Build messages array from conversation history now_iso = datetime.now(timezone.utc).isoformat() history = list(session.conversation_history) history.append({"role": "user", "content": user_message, "timestamp": now_iso}) # Convert to provider format (just role + content) provider_messages = [ {"role": msg["role"], "content": msg["content"]} for msg in history ] provider = get_ai_provider() response_text, input_tokens, output_tokens = await provider.generate_json( system_prompt=system_prompt, messages=provider_messages, max_tokens=2000, ) parsed = _parse_ai_response(response_text) # Validate tree update if present tree_update = parsed["tree_update"] if tree_update: errors = validate_generated_tree(tree_update) if errors: logger.warning("AI tree update failed validation: %s", errors) tree_update = None # Silently discard invalid updates # Update session state history.append({"role": "assistant", "content": parsed["content"], "timestamp": now_iso}) session.conversation_history = history session.message_count = session.message_count + 1 session.total_input_tokens = session.total_input_tokens + input_tokens session.total_output_tokens = session.total_output_tokens + output_tokens if tree_update: session.working_tree = tree_update if parsed["phase"]: valid_phases = {"scoping", "discovery", "enrichment", "review", "generation"} if parsed["phase"] in valid_phases: session.current_phase = parsed["phase"] if parsed["metadata"]: merged = dict(session.tree_metadata) merged.update(parsed["metadata"]) session.tree_metadata = merged session.updated_at = datetime.now(timezone.utc) return parsed["content"], tree_update, parsed["phase"], parsed["metadata"] async def generate_final_tree( session: AIChatSession, db: AsyncSession, ) -> tuple[dict[str, Any], dict[str, Any]]: """Generate the final validated TreeStructure from the conversation. Returns (tree_structure, metadata). Raises ValueError if generation fails after retry. """ system_prompt = _build_system_prompt(session.flow_type) # Build generation prompt from full conversation provider_messages = [ {"role": msg["role"], "content": msg["content"]} for msg in session.conversation_history ] generation_instruction = """Based on our entire conversation, generate the COMPLETE and FINAL TreeStructure JSON for this flow. Requirements: - Include ALL branches, steps, and solutions we discussed - Use descriptive node IDs (slugs, not UUIDs) - Root node must be type "decision" - Every decision option must have a valid next_node_id pointing to a child - Every action node should have commands with exact syntax where discussed - Every action node should have expected_outcome where discussed - Solution nodes should have resolution_steps - Respond with ONLY the JSON — no conversational text, no markdown fences Also provide metadata as a separate JSON object after the tree: [METADATA] {"name": "...", "description": "...", "tags": ["..."]} [/METADATA]""" provider_messages.append({"role": "user", "content": generation_instruction}) provider = get_ai_provider() for attempt in range(2): # One try + one retry response_text, input_tokens, output_tokens = await provider.generate_json( system_prompt=system_prompt, messages=provider_messages, max_tokens=8000, ) session.total_input_tokens = session.total_input_tokens + input_tokens session.total_output_tokens = session.total_output_tokens + output_tokens # Extract metadata first parsed = _parse_ai_response(response_text) metadata = parsed["metadata"] or dict(session.tree_metadata) # Parse tree JSON — could be in tree_update marker or raw tree = parsed["tree_update"] if not tree: try: raw = _strip_markdown_fences(parsed["content"]) tree = json.loads(raw) except (json.JSONDecodeError, ValueError): pass if not tree: if attempt == 0: provider_messages.append({"role": "assistant", "content": response_text}) provider_messages.append({ "role": "user", "content": "That response was not valid JSON. Please respond with ONLY the TreeStructure JSON object, starting with { and ending with }. No markdown fences, no explanatory text.", }) continue raise ValueError("AI failed to produce valid JSON after retry") errors = validate_generated_tree(tree) if errors: if attempt == 0: provider_messages.append({"role": "assistant", "content": response_text}) correction = ( f"The tree has validation errors: {'; '.join(errors)}. " "Please fix these issues and respond with the corrected JSON only." ) provider_messages.append({"role": "user", "content": correction}) continue raise ValueError(f"Generated tree failed validation: {'; '.join(errors)}") # Success session.working_tree = tree session.tree_metadata = metadata session.current_phase = "generation" session.updated_at = datetime.now(timezone.utc) return tree, metadata raise ValueError("AI failed to generate a valid tree") async def get_chat_session( session_id: uuid.UUID, user_id: uuid.UUID, db: AsyncSession, ) -> AIChatSession: """Get a chat session, validating ownership and expiry. Raises HTTPException on not found, forbidden, or expired. """ from fastapi import HTTPException, status result = await db.execute( select(AIChatSession).where(AIChatSession.id == session_id) ) session = result.scalar_one_or_none() if not session: raise HTTPException(status_code=404, detail="Chat session not found") if session.user_id != user_id: raise HTTPException(status_code=403, detail="Access denied") if session.expires_at < datetime.now(timezone.utc): session.status = "abandoned" await db.flush() raise HTTPException(status_code=410, detail="Chat session has expired") return session ``` **Step 2: Commit** ```bash git add backend/app/core/ai_chat_service.py git commit -m "feat: add AI chat builder service with system prompt and conversation loop" ``` --- ### Task 4: API Endpoints **Files:** - Create: `backend/app/api/endpoints/ai_chat.py` - Modify: `backend/app/api/router.py` (lines 9, 40) **Step 1: Create the endpoints file** ```python # backend/app/api/endpoints/ai_chat.py """AI Chat Builder endpoints. Conversational flow builder: POST /ai/chat/sessions — Start session, get AI greeting POST /ai/chat/sessions/{id}/messages — Send message, get AI response GET /ai/chat/sessions/{id} — Get session state (for resume) POST /ai/chat/sessions/{id}/generate — Generate final TreeStructure POST /ai/chat/sessions/{id}/import — Create Tree from generated structure DELETE /ai/chat/sessions/{id} — Abandon session """ import logging from typing import Annotated from uuid import UUID from fastapi import APIRouter, Depends, HTTPException, Request, status from sqlalchemy.ext.asyncio import AsyncSession from app.core.rate_limit import limiter from app.api.deps import get_current_active_user, get_db, require_engineer_or_admin from app.core.config import settings from app.core.ai_chat_service import ( start_chat_session, send_message, generate_final_tree, get_chat_session, MAX_MESSAGES_FREE, MAX_MESSAGES_PAID, ) from app.core.ai_quota_service import check_ai_quota, record_ai_usage, get_user_plan from app.models.user import User from app.models.tree import Tree from app.schemas.ai_chat import ( AIChatStartRequest, AIChatStartResponse, AIChatMessageRequest, AIChatMessageResponse, AIChatSessionResponse, AIChatGenerateResponse, AIChatImportRequest, AIChatImportResponse, ) logger = logging.getLogger(__name__) router = APIRouter(prefix="/ai/chat", tags=["ai-chat-builder"]) def _require_ai_enabled() -> None: if not settings.ai_enabled: raise HTTPException( status_code=status.HTTP_503_SERVICE_UNAVAILABLE, detail="AI is not configured. Set GOOGLE_AI_API_KEY or ANTHROPIC_API_KEY.", ) @router.post("/sessions", response_model=AIChatStartResponse, status_code=201) @limiter.limit("10/minute") async def create_session( request: Request, data: AIChatStartRequest, current_user: Annotated[User, Depends(get_current_active_user)], db: Annotated[AsyncSession, Depends(get_db)], _: None = Depends(require_engineer_or_admin), ): """Start a new AI chat builder session.""" _require_ai_enabled() # Check quota allowed, quota_status = await check_ai_quota( user_id=current_user.id, account_id=current_user.account_id, db=db, billing_anchor=current_user.ai_billing_cycle_anchor_at, is_super_admin=current_user.is_super_admin, ) if not allowed: reset_key = ( "daily_reset_at" if quota_status.get("deny_reason") == "daily" else "monthly_reset_at" ) raise HTTPException( status_code=status.HTTP_429_TOO_MANY_REQUESTS, detail={ "message": f"AI build limit exceeded ({quota_status['deny_reason']})", "reset_at": quota_status.get(reset_key), "quota": quota_status, }, ) plan = await get_user_plan(current_user.account_id, db) try: session, greeting = await start_chat_session( flow_type=data.flow_type, user_id=current_user.id, account_id=current_user.account_id, db=db, ) except Exception as e: logger.exception("AI chat session start failed: %s", e) raise HTTPException( status_code=status.HTTP_502_BAD_GATEWAY, detail=f"AI provider error ({type(e).__name__}). Please try again.", ) # Record usage for the greeting await record_ai_usage( user_id=current_user.id, account_id=current_user.account_id, conversation_id=session.id, generation_type="chat_message", tier=plan, input_tokens=session.total_input_tokens, output_tokens=session.total_output_tokens, estimated_cost=( session.total_input_tokens * 1.0 / 1_000_000 + session.total_output_tokens * 5.0 / 1_000_000 ), succeeded=True, counts_toward_quota=False, error_code=None, extra_data={"phase": "scoping"}, db=db, ) await db.commit() return AIChatStartResponse( session_id=session.id, greeting=greeting, current_phase=session.current_phase, ) @router.post("/sessions/{session_id}/messages", response_model=AIChatMessageResponse) @limiter.limit("10/minute") async def post_message( request: Request, session_id: UUID, data: AIChatMessageRequest, current_user: Annotated[User, Depends(get_current_active_user)], db: Annotated[AsyncSession, Depends(get_db)], _: None = Depends(require_engineer_or_admin), ): """Send a user message and get AI response.""" _require_ai_enabled() session = await get_chat_session(session_id, current_user.id, db) if session.status != "active": raise HTTPException( status_code=status.HTTP_400_BAD_REQUEST, detail=f"Session is {session.status}, cannot send messages", ) # Per-session message limit plan = await get_user_plan(current_user.account_id, db) max_messages = MAX_MESSAGES_PAID if plan != "free" else MAX_MESSAGES_FREE if current_user.is_super_admin: max_messages = 999 if session.message_count >= max_messages: raise HTTPException( status_code=status.HTTP_429_TOO_MANY_REQUESTS, detail=f"Maximum messages per session reached ({max_messages}). Generate your tree or start a new session.", ) prev_input = session.total_input_tokens prev_output = session.total_output_tokens try: ai_content, tree_update, new_phase, metadata = await send_message( session, data.content, db ) except Exception as e: logger.exception("AI chat message failed: %s", e) await record_ai_usage( user_id=current_user.id, account_id=current_user.account_id, conversation_id=session.id, generation_type="chat_message", tier=plan, input_tokens=0, output_tokens=0, estimated_cost=0, succeeded=False, counts_toward_quota=False, error_code=type(e).__name__, extra_data=None, db=db, ) await db.commit() raise HTTPException( status_code=status.HTTP_502_BAD_GATEWAY, detail=f"AI provider error ({type(e).__name__}). Please try again.", ) # Record usage input_delta = session.total_input_tokens - prev_input output_delta = session.total_output_tokens - prev_output await record_ai_usage( user_id=current_user.id, account_id=current_user.account_id, conversation_id=session.id, generation_type="chat_message", tier=plan, input_tokens=input_delta, output_tokens=output_delta, estimated_cost=( input_delta * 1.0 / 1_000_000 + output_delta * 5.0 / 1_000_000 ), succeeded=True, counts_toward_quota=False, error_code=None, extra_data={"phase": session.current_phase}, db=db, ) await db.commit() return AIChatMessageResponse( content=ai_content, current_phase=session.current_phase, working_tree=session.working_tree, tree_metadata=session.tree_metadata if session.tree_metadata else None, ) @router.get("/sessions/{session_id}", response_model=AIChatSessionResponse) async def get_session( session_id: UUID, current_user: Annotated[User, Depends(get_current_active_user)], db: Annotated[AsyncSession, Depends(get_db)], ): """Get full session state for resume after page reload.""" session = await get_chat_session(session_id, current_user.id, db) # Filter out hidden messages (like the primer) visible_history = [ msg for msg in session.conversation_history if not msg.get("hidden") ] return AIChatSessionResponse( session_id=session.id, status=session.status, current_phase=session.current_phase, flow_type=session.flow_type, conversation_history=visible_history, working_tree=session.working_tree, tree_metadata=session.tree_metadata if session.tree_metadata else None, message_count=session.message_count, generated_tree=session.working_tree if session.status == "completed" else None, ) @router.post("/sessions/{session_id}/generate", response_model=AIChatGenerateResponse) @limiter.limit("10/minute") async def generate_tree( request: Request, session_id: UUID, current_user: Annotated[User, Depends(get_current_active_user)], db: Annotated[AsyncSession, Depends(get_db)], _: None = Depends(require_engineer_or_admin), ): """Generate final TreeStructure JSON from conversation.""" _require_ai_enabled() session = await get_chat_session(session_id, current_user.id, db) # If already generated, return cached result if session.status == "completed" and session.working_tree: return AIChatGenerateResponse( tree_structure=session.working_tree, tree_metadata=session.tree_metadata, status="completed", ) if session.status != "active": raise HTTPException( status_code=status.HTTP_400_BAD_REQUEST, detail=f"Session is {session.status}, cannot generate", ) plan = await get_user_plan(current_user.account_id, db) prev_input = session.total_input_tokens prev_output = session.total_output_tokens try: tree_structure, metadata = await generate_final_tree(session, db) except ValueError as e: await record_ai_usage( user_id=current_user.id, account_id=current_user.account_id, conversation_id=session.id, generation_type="chat_generate", tier=plan, input_tokens=session.total_input_tokens - prev_input, output_tokens=session.total_output_tokens - prev_output, estimated_cost=0, succeeded=False, counts_toward_quota=False, error_code="invalid_output", extra_data={"error": str(e)}, db=db, ) await db.commit() raise HTTPException( status_code=status.HTTP_422_UNPROCESSABLE_ENTITY, detail=f"Tree generation failed: {e}", ) except Exception as e: logger.exception("AI chat generate failed: %s", e) input_delta = session.total_input_tokens - prev_input output_delta = session.total_output_tokens - prev_output await record_ai_usage( user_id=current_user.id, account_id=current_user.account_id, conversation_id=session.id, generation_type="chat_generate", tier=plan, input_tokens=input_delta, output_tokens=output_delta, estimated_cost=0, succeeded=False, counts_toward_quota=False, error_code=type(e).__name__, extra_data={"error": str(e)}, db=db, ) await db.commit() # Distinguish timeout from other errors error_name = type(e).__name__ if "timeout" in error_name.lower() or "Timeout" in str(e): raise HTTPException( status_code=status.HTTP_504_GATEWAY_TIMEOUT, detail="Tree generation timed out. Please try again.", ) raise HTTPException( status_code=status.HTTP_502_BAD_GATEWAY, detail=f"AI provider error ({error_name}). Please try again.", ) # Record successful quota-consuming usage input_delta = session.total_input_tokens - prev_input output_delta = session.total_output_tokens - prev_output await record_ai_usage( user_id=current_user.id, account_id=current_user.account_id, conversation_id=session.id, generation_type="chat_generate", tier=plan, input_tokens=input_delta, output_tokens=output_delta, estimated_cost=( input_delta * 1.0 / 1_000_000 + output_delta * 5.0 / 1_000_000 ), succeeded=True, counts_toward_quota=True, error_code=None, extra_data=None, db=db, ) session.status = "completed" await db.commit() return AIChatGenerateResponse( tree_structure=tree_structure, tree_metadata=metadata, status="completed", ) @router.post("/sessions/{session_id}/import", response_model=AIChatImportResponse) @limiter.limit("10/minute") async def import_tree( request: Request, session_id: UUID, data: AIChatImportRequest, current_user: Annotated[User, Depends(get_current_active_user)], db: Annotated[AsyncSession, Depends(get_db)], _: None = Depends(require_engineer_or_admin), ): """Create a Tree record from the generated tree structure.""" session = await get_chat_session(session_id, current_user.id, db) if session.status != "completed" or not session.working_tree: raise HTTPException( status_code=status.HTTP_400_BAD_REQUEST, detail="Session must be completed with a generated tree before importing", ) # Already imported? if session.generated_tree_id: return AIChatImportResponse( tree_id=session.generated_tree_id, tree_type=session.flow_type, ) metadata = session.tree_metadata or {} tree = Tree( name=data.name or metadata.get("name", "AI-Generated Flow"), description=data.description or metadata.get("description", ""), tree_type=session.flow_type, tree_structure=session.working_tree, author_id=current_user.id, account_id=current_user.account_id, category_id=data.category_id, is_public=False, ) db.add(tree) await db.flush() session.generated_tree_id = tree.id await db.commit() return AIChatImportResponse( tree_id=tree.id, tree_type=session.flow_type, ) @router.delete("/sessions/{session_id}", status_code=204) async def abandon_session( session_id: UUID, current_user: Annotated[User, Depends(get_current_active_user)], db: Annotated[AsyncSession, Depends(get_db)], ): """Abandon a chat session.""" session = await get_chat_session(session_id, current_user.id, db) session.status = "abandoned" await db.commit() ``` **Step 2: Register in router** In `backend/app/api/router.py`, add after line 9 (`from app.api.endpoints import ai_fix`): ```python from app.api.endpoints import ai_chat ``` Add after line 40 (`api_router.include_router(ai_fix.router)`): ```python api_router.include_router(ai_chat.router) ``` **Step 3: Update quota service daily limit** In `backend/app/core/ai_quota_service.py`, change line 118: ```python # Before: AIUsage.generation_type.in_(["scaffold", "branch_detail"]), # After: AIUsage.generation_type.in_(["scaffold", "branch_detail", "chat_message", "chat_generate"]), ``` **Step 4: Commit** ```bash git add backend/app/api/endpoints/ai_chat.py backend/app/api/router.py backend/app/core/ai_quota_service.py git commit -m "feat: add AI chat builder endpoints and update quota service" ``` --- ### Task 5: Backend Tests **Files:** - Create: `backend/tests/test_ai_chat.py` **Step 1: Write integration tests** ```python # backend/tests/test_ai_chat.py """Integration tests for AI Chat Builder endpoints. These tests mock the AI provider to avoid real API calls. """ import pytest from unittest.mock import AsyncMock, patch pytestmark = pytest.mark.asyncio @pytest.fixture def mock_ai_provider(): """Mock AI provider that returns realistic responses.""" provider = AsyncMock() provider.generate_json = AsyncMock(return_value=( "Great question! Let's build a troubleshooting flow for DNS resolution issues. " "To start, I need to understand the scope.\n\n" "Who is the target audience for this flow? Are we targeting:\n" "- Tier 1 help desk (basic checks only)\n" "- Tier 2 desktop support (intermediate diagnostics)\n" "- Tier 3 systems engineers (deep DNS troubleshooting)\n\n" "[PHASE:scoping]", 500, # input tokens 200, # output tokens )) return provider async def test_create_chat_session(client, auth_headers, mock_ai_provider): """POST /ai/chat/sessions creates a session and returns AI greeting.""" with patch("app.core.ai_chat_service.get_ai_provider", return_value=mock_ai_provider): resp = await client.post( "/api/v1/ai/chat/sessions", json={"flow_type": "troubleshooting"}, headers=auth_headers, ) assert resp.status_code == 201 data = resp.json() assert "session_id" in data assert "greeting" in data assert data["current_phase"] == "scoping" assert len(data["greeting"]) > 0 async def test_send_message(client, auth_headers, mock_ai_provider): """POST /ai/chat/sessions/{id}/messages returns AI response.""" # Create session first with patch("app.core.ai_chat_service.get_ai_provider", return_value=mock_ai_provider): create_resp = await client.post( "/api/v1/ai/chat/sessions", json={"flow_type": "troubleshooting"}, headers=auth_headers, ) session_id = create_resp.json()["session_id"] # Mock response with tree update mock_ai_provider.generate_json = AsyncMock(return_value=( "Good, targeting Tier 2 support. Let's start with the first diagnostic question.\n\n" "The root question should be: 'What DNS symptom is the user experiencing?'\n\n" "[TREE_UPDATE]\n" '{"id": "root", "type": "decision", "question": "What DNS symptom is the user experiencing?", ' '"options": [{"id": "opt-1", "label": "Cannot resolve any domains", "next_node_id": "check-dns-service"}], ' '"children": [{"id": "check-dns-service", "type": "action", "title": "Check DNS Client Service", ' '"description": "Verify the DNS Client service is running", ' '"commands": ["Get-Service Dnscache | Select-Object Status, StartType"], ' '"expected_outcome": "Service should show Running status"}]}\n' "[/TREE_UPDATE]\n\n" "[PHASE:discovery]", 800, 400, )) with patch("app.core.ai_chat_service.get_ai_provider", return_value=mock_ai_provider): resp = await client.post( f"/api/v1/ai/chat/sessions/{session_id}/messages", json={"content": "This is for Tier 2 support, hybrid environment with on-prem AD."}, headers=auth_headers, ) assert resp.status_code == 200 data = resp.json() assert "content" in data assert data["current_phase"] == "discovery" assert data["working_tree"] is not None assert data["working_tree"]["type"] == "decision" # Markers should be stripped from content assert "[TREE_UPDATE]" not in data["content"] assert "[PHASE:" not in data["content"] async def test_get_session(client, auth_headers, mock_ai_provider): """GET /ai/chat/sessions/{id} returns full session state.""" with patch("app.core.ai_chat_service.get_ai_provider", return_value=mock_ai_provider): create_resp = await client.post( "/api/v1/ai/chat/sessions", json={"flow_type": "troubleshooting"}, headers=auth_headers, ) session_id = create_resp.json()["session_id"] resp = await client.get( f"/api/v1/ai/chat/sessions/{session_id}", headers=auth_headers, ) assert resp.status_code == 200 data = resp.json() assert data["session_id"] == session_id assert data["status"] == "active" assert data["flow_type"] == "troubleshooting" # Hidden primer message should be filtered out assert all(msg.get("role") == "assistant" or not msg.get("hidden") for msg in data["conversation_history"]) async def test_abandon_session(client, auth_headers, mock_ai_provider): """DELETE /ai/chat/sessions/{id} sets status to abandoned.""" with patch("app.core.ai_chat_service.get_ai_provider", return_value=mock_ai_provider): create_resp = await client.post( "/api/v1/ai/chat/sessions", json={"flow_type": "troubleshooting"}, headers=auth_headers, ) session_id = create_resp.json()["session_id"] resp = await client.delete( f"/api/v1/ai/chat/sessions/{session_id}", headers=auth_headers, ) assert resp.status_code == 204 # Verify session is abandoned get_resp = await client.get( f"/api/v1/ai/chat/sessions/{session_id}", headers=auth_headers, ) assert get_resp.json()["status"] == "abandoned" async def test_message_limit_enforced(client, auth_headers, mock_ai_provider): """Per-session message limit should return 429.""" with patch("app.core.ai_chat_service.get_ai_provider", return_value=mock_ai_provider): create_resp = await client.post( "/api/v1/ai/chat/sessions", json={"flow_type": "troubleshooting"}, headers=auth_headers, ) session_id = create_resp.json()["session_id"] # Patch the session to have max messages already with patch("app.api.endpoints.ai_chat.MAX_MESSAGES_FREE", 1): with patch("app.core.ai_chat_service.get_ai_provider", return_value=mock_ai_provider): resp = await client.post( f"/api/v1/ai/chat/sessions/{session_id}/messages", json={"content": "test message"}, headers=auth_headers, ) assert resp.status_code == 429 async def test_session_not_found(client, auth_headers): """Accessing nonexistent session returns 404.""" import uuid fake_id = str(uuid.uuid4()) resp = await client.get( f"/api/v1/ai/chat/sessions/{fake_id}", headers=auth_headers, ) assert resp.status_code == 404 async def test_ai_disabled_returns_503(client, auth_headers): """When AI is not configured, endpoints return 503.""" with patch("app.api.endpoints.ai_chat.settings") as mock_settings: mock_settings.ai_enabled = False resp = await client.post( "/api/v1/ai/chat/sessions", json={"flow_type": "troubleshooting"}, headers=auth_headers, ) assert resp.status_code == 503 ``` **Step 2: Run tests** Run: `cd backend && python -m pytest tests/test_ai_chat.py -v --override-ini="addopts="` Expected: All tests pass. If fixture issues, check `conftest.py` for available fixtures (CLAUDE.md lesson #21). **Step 3: Commit** ```bash git add backend/tests/test_ai_chat.py git commit -m "test: add integration tests for AI chat builder endpoints" ``` --- ## Phase 2: Frontend Chat UI ### Task 6: TypeScript Types **Files:** - Create: `frontend/src/types/ai-chat.ts` - Modify: `frontend/src/types/index.ts` (after line 54) **Step 1: Create the types file** ```typescript // frontend/src/types/ai-chat.ts export type InterviewPhase = 'scoping' | 'discovery' | 'enrichment' | 'review' | 'generation' export interface ChatMessage { role: 'user' | 'assistant' content: string timestamp: string } export interface AIChatStartResponse { session_id: string greeting: string current_phase: InterviewPhase } export interface AIChatMessageResponse { content: string current_phase: InterviewPhase working_tree: Record | null tree_metadata: Record | null } export interface AIChatSessionResponse { session_id: string status: 'active' | 'completed' | 'abandoned' current_phase: InterviewPhase flow_type: 'troubleshooting' | 'procedural' conversation_history: ChatMessage[] working_tree: Record | null tree_metadata: Record | null message_count: number generated_tree: Record | null } export interface AIChatGenerateResponse { tree_structure: Record tree_metadata: Record status: string } export interface AIChatImportResponse { tree_id: string tree_type: string } ``` **Step 2: Export from types/index.ts** Add after line 54 (the `ai-fix` export block): ```typescript export type { InterviewPhase, ChatMessage, AIChatStartResponse, AIChatMessageResponse, AIChatSessionResponse, AIChatGenerateResponse, AIChatImportResponse, } from './ai-chat' ``` **Step 3: Commit** ```bash git add frontend/src/types/ai-chat.ts frontend/src/types/index.ts git commit -m "feat: add TypeScript types for AI chat builder" ``` --- ### Task 7: API Client Module **Files:** - Create: `frontend/src/api/aiChat.ts` - Modify: `frontend/src/api/index.ts` (after line 19) **Step 1: Create the API client** ```typescript // frontend/src/api/aiChat.ts import { apiClient } from './client' import type { AIChatStartResponse, AIChatMessageResponse, AIChatSessionResponse, AIChatGenerateResponse, AIChatImportResponse, } from '@/types' export const aiChatApi = { startSession: async (flowType: 'troubleshooting' | 'procedural'): Promise => { const { data } = await apiClient.post('/ai/chat/sessions', { flow_type: flowType }) return data }, sendMessage: async (sessionId: string, content: string): Promise => { const { data } = await apiClient.post(`/ai/chat/sessions/${sessionId}/messages`, { content }) return data }, getSession: async (sessionId: string): Promise => { const { data } = await apiClient.get(`/ai/chat/sessions/${sessionId}`) return data }, generateTree: async (sessionId: string): Promise => { const { data } = await apiClient.post(`/ai/chat/sessions/${sessionId}/generate`) return data }, importTree: async ( sessionId: string, params?: { name?: string; description?: string; category_id?: string; tags?: string[] } ): Promise => { const { data } = await apiClient.post(`/ai/chat/sessions/${sessionId}/import`, params || {}) return data }, abandonSession: async (sessionId: string): Promise => { await apiClient.delete(`/ai/chat/sessions/${sessionId}`) }, } export default aiChatApi ``` **Step 2: Export from api/index.ts** Add after line 19 (`export { default as aiBuilderApi } from './aiBuilder'`): ```typescript export { default as aiChatApi } from './aiChat' ``` **Step 3: Commit** ```bash git add frontend/src/api/aiChat.ts frontend/src/api/index.ts git commit -m "feat: add API client for AI chat builder" ``` --- ### Task 8: Zustand Store **Files:** - Create: `frontend/src/store/aiChatStore.ts` **Step 1: Create the store** ```typescript // frontend/src/store/aiChatStore.ts import { create } from 'zustand' import { aiChatApi } from '@/api/aiChat' import type { ChatMessage, InterviewPhase, TreeStructure, } from '@/types' interface TreeMetadata { name?: string description?: string tags?: string[] category_id?: string } interface AIChatState { // Session sessionId: string | null status: 'idle' | 'active' | 'completed' | 'abandoned' currentPhase: InterviewPhase flowType: 'troubleshooting' | 'procedural' | null // Conversation messages: ChatMessage[] isResponding: boolean // Progressive tree workingTree: TreeStructure | null treeMetadata: TreeMetadata | null // Final generation generatedTree: TreeStructure | null isGenerating: boolean importedTreeId: string | null // Error error: string | null // Actions startSession: (flowType: 'troubleshooting' | 'procedural') => Promise sendMessage: (content: string) => Promise generateTree: () => Promise importToEditor: (params?: { name?: string; description?: string; category_id?: string; tags?: string[] }) => Promise abandonSession: () => Promise resumeSession: (sessionId: string) => Promise reset: () => void } const initialState = { sessionId: null, status: 'idle' as const, currentPhase: 'scoping' as InterviewPhase, flowType: null, messages: [], isResponding: false, workingTree: null, treeMetadata: null, generatedTree: null, isGenerating: false, importedTreeId: null, error: null, } export const useAIChatStore = create((set, get) => ({ ...initialState, startSession: async (flowType) => { set({ ...initialState, status: 'active', flowType, isResponding: true, error: null }) try { const response = await aiChatApi.startSession(flowType) set({ sessionId: response.session_id, currentPhase: response.current_phase, messages: [{ role: 'assistant', content: response.greeting, timestamp: new Date().toISOString(), }], isResponding: false, }) } catch (e: unknown) { const message = e instanceof Error ? e.message : 'Failed to start session' set({ error: message, isResponding: false, status: 'idle' }) } }, sendMessage: async (content) => { const { sessionId, messages } = get() if (!sessionId) return const userMessage: ChatMessage = { role: 'user', content, timestamp: new Date().toISOString(), } set({ messages: [...messages, userMessage], isResponding: true, error: null, }) try { const response = await aiChatApi.sendMessage(sessionId, content) const aiMessage: ChatMessage = { role: 'assistant', content: response.content, timestamp: new Date().toISOString(), } set((state) => ({ messages: [...state.messages, aiMessage], currentPhase: response.current_phase, workingTree: response.working_tree as TreeStructure | null ?? state.workingTree, treeMetadata: response.tree_metadata as TreeMetadata | null ?? state.treeMetadata, isResponding: false, })) } catch (e: unknown) { const message = e instanceof Error ? e.message : 'Failed to send message' set({ error: message, isResponding: false }) } }, generateTree: async () => { const { sessionId } = get() if (!sessionId) return set({ isGenerating: true, error: null }) try { const response = await aiChatApi.generateTree(sessionId) set({ generatedTree: response.tree_structure as TreeStructure, workingTree: response.tree_structure as TreeStructure, treeMetadata: response.tree_metadata as TreeMetadata, status: 'completed', isGenerating: false, }) } catch (e: unknown) { const message = e instanceof Error ? e.message : 'Failed to generate tree' set({ error: message, isGenerating: false }) } }, importToEditor: async (params) => { const { sessionId } = get() if (!sessionId) throw new Error('No active session') const response = await aiChatApi.importTree(sessionId, params) set({ importedTreeId: response.tree_id }) return response.tree_id }, abandonSession: async () => { const { sessionId } = get() if (!sessionId) return try { await aiChatApi.abandonSession(sessionId) } catch { // Best effort — session may have already expired } set({ ...initialState }) }, resumeSession: async (sessionId) => { set({ isResponding: true, error: null }) try { const session = await aiChatApi.getSession(sessionId) set({ sessionId: session.session_id, status: session.status === 'active' ? 'active' : session.status as 'completed' | 'abandoned', currentPhase: session.current_phase, flowType: session.flow_type, messages: session.conversation_history as ChatMessage[], workingTree: session.working_tree as TreeStructure | null, treeMetadata: session.tree_metadata as TreeMetadata | null, generatedTree: session.generated_tree as TreeStructure | null, isResponding: false, }) } catch (e: unknown) { const message = e instanceof Error ? e.message : 'Failed to resume session' set({ error: message, isResponding: false }) } }, reset: () => set({ ...initialState }), })) ``` **Step 2: Commit** ```bash git add frontend/src/store/aiChatStore.ts git commit -m "feat: add Zustand store for AI chat builder" ``` --- ### Task 9: Chat Components **Files:** - Create: `frontend/src/components/ai-chat/ChatMessage.tsx` - Create: `frontend/src/components/ai-chat/ChatInput.tsx` - Create: `frontend/src/components/ai-chat/ChatPanel.tsx` - Create: `frontend/src/components/ai-chat/PhaseIndicator.tsx` - Create: `frontend/src/components/ai-chat/ChatToolbar.tsx` - Create: `frontend/src/components/ai-chat/EmptyPreview.tsx` - Create: `frontend/src/components/ai-chat/StaticTreePreview.tsx` This task creates all 7 components. Each is small and focused. **Step 1: Create `ChatMessage.tsx`** ```typescript // frontend/src/components/ai-chat/ChatMessage.tsx import { Bot, User } from 'lucide-react' import { MarkdownContent } from '@/components/ui/MarkdownContent' import { cn } from '@/lib/utils' import type { ChatMessage as ChatMessageType } from '@/types' interface ChatMessageProps { message: ChatMessageType } export function ChatMessage({ message }: ChatMessageProps) { const isAI = message.role === 'assistant' return (
{isAI && (
)}
{isAI ? ( ) : (

{message.content}

)}
{!isAI && (
)}
) } ``` **Step 2: Create `ChatInput.tsx`** ```typescript // frontend/src/components/ai-chat/ChatInput.tsx import { useState, useRef, useCallback } from 'react' import { Send } from 'lucide-react' import { cn } from '@/lib/utils' interface ChatInputProps { onSend: (content: string) => void disabled?: boolean placeholder?: string } export function ChatInput({ onSend, disabled, placeholder = 'Type a message...' }: ChatInputProps) { const [value, setValue] = useState('') const textareaRef = useRef(null) const handleSend = useCallback(() => { const trimmed = value.trim() if (!trimmed || disabled) return onSend(trimmed) setValue('') if (textareaRef.current) { textareaRef.current.style.height = 'auto' } }, [value, disabled, onSend]) const handleKeyDown = (e: React.KeyboardEvent) => { if (e.key === 'Enter' && !e.shiftKey) { e.preventDefault() handleSend() } } const handleInput = () => { if (textareaRef.current) { textareaRef.current.style.height = 'auto' textareaRef.current.style.height = Math.min(textareaRef.current.scrollHeight, 160) + 'px' } } return (