Files
resolutionflow/docs/plans/2026-02-27-ai-chat-builder-implementation.md
chihlasm 09c5d60067 docs: add AI Chat Builder implementation plan
14 tasks across 4 phases: backend foundation (model, schemas, service,
endpoints, tests), frontend chat UI (types, API client, store,
components, page), tree preview integration, and polish (session
resume, responsive layout).

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
2026-02-27 03:25:17 -05:00

84 KiB

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

# 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):

from .ai_chat_session import AIChatSession

Add to __all__ after line 69 ("AIUsage"):

    "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):

"""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

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

# 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

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

# 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

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

# 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):

from app.api.endpoints import ai_chat

Add after line 40 (api_router.include_router(ai_fix.router)):

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:

# Before:
            AIUsage.generation_type.in_(["scaffold", "branch_detail"]),
# After:
            AIUsage.generation_type.in_(["scaffold", "branch_detail", "chat_message", "chat_generate"]),

Step 4: Commit

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

# 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

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

// 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<string, unknown> | null
  tree_metadata: Record<string, unknown> | null
}

export interface AIChatSessionResponse {
  session_id: string
  status: 'active' | 'completed' | 'abandoned'
  current_phase: InterviewPhase
  flow_type: 'troubleshooting' | 'procedural'
  conversation_history: ChatMessage[]
  working_tree: Record<string, unknown> | null
  tree_metadata: Record<string, unknown> | null
  message_count: number
  generated_tree: Record<string, unknown> | null
}

export interface AIChatGenerateResponse {
  tree_structure: Record<string, unknown>
  tree_metadata: Record<string, unknown>
  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):


export type {
  InterviewPhase,
  ChatMessage,
  AIChatStartResponse,
  AIChatMessageResponse,
  AIChatSessionResponse,
  AIChatGenerateResponse,
  AIChatImportResponse,
} from './ai-chat'

Step 3: Commit

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

// 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<AIChatStartResponse> => {
    const { data } = await apiClient.post('/ai/chat/sessions', { flow_type: flowType })
    return data
  },

  sendMessage: async (sessionId: string, content: string): Promise<AIChatMessageResponse> => {
    const { data } = await apiClient.post(`/ai/chat/sessions/${sessionId}/messages`, { content })
    return data
  },

  getSession: async (sessionId: string): Promise<AIChatSessionResponse> => {
    const { data } = await apiClient.get(`/ai/chat/sessions/${sessionId}`)
    return data
  },

  generateTree: async (sessionId: string): Promise<AIChatGenerateResponse> => {
    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<AIChatImportResponse> => {
    const { data } = await apiClient.post(`/ai/chat/sessions/${sessionId}/import`, params || {})
    return data
  },

  abandonSession: async (sessionId: string): Promise<void> => {
    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'):

export { default as aiChatApi } from './aiChat'

Step 3: Commit

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

// 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<void>
  sendMessage: (content: string) => Promise<void>
  generateTree: () => Promise<void>
  importToEditor: (params?: { name?: string; description?: string; category_id?: string; tags?: string[] }) => Promise<string>
  abandonSession: () => Promise<void>
  resumeSession: (sessionId: string) => Promise<void>
  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<AIChatState>((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

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

// 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 (
    <div className={cn('flex gap-3', isAI ? 'items-start' : 'items-start justify-end')}>
      {isAI && (
        <div className="flex h-8 w-8 shrink-0 items-center justify-center rounded-lg bg-primary/10">
          <Bot className="h-4 w-4 text-primary" />
        </div>
      )}
      <div
        className={cn(
          'max-w-[85%] rounded-xl px-4 py-3',
          isAI
            ? 'bg-card border border-border'
            : 'bg-primary/10 border border-primary/20'
        )}
      >
        {isAI ? (
          <MarkdownContent content={message.content} className="text-sm" />
        ) : (
          <p className="text-sm text-foreground whitespace-pre-wrap">{message.content}</p>
        )}
      </div>
      {!isAI && (
        <div className="flex h-8 w-8 shrink-0 items-center justify-center rounded-lg bg-accent">
          <User className="h-4 w-4 text-foreground" />
        </div>
      )}
    </div>
  )
}

Step 2: Create ChatInput.tsx

// 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<HTMLTextAreaElement>(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 (
    <div className="flex items-end gap-2 border-t border-border bg-card px-4 py-3">
      <textarea
        ref={textareaRef}
        value={value}
        onChange={(e) => setValue(e.target.value)}
        onKeyDown={handleKeyDown}
        onInput={handleInput}
        placeholder={placeholder}
        disabled={disabled}
        rows={1}
        className={cn(
          'flex-1 resize-none rounded-lg border border-border bg-background px-3 py-2',
          'text-sm text-foreground placeholder:text-muted-foreground',
          'focus:border-primary focus:ring-1 focus:ring-primary/20 focus:outline-none',
          'disabled:opacity-50 disabled:cursor-not-allowed',
          'max-h-40'
        )}
      />
      <button
        onClick={handleSend}
        disabled={disabled || !value.trim()}
        className={cn(
          'flex h-10 w-10 shrink-0 items-center justify-center rounded-lg',
          'bg-gradient-brand text-white shadow-lg shadow-primary/20',
          'hover:opacity-90 transition-opacity',
          'disabled:opacity-50 disabled:cursor-not-allowed'
        )}
      >
        <Send className="h-4 w-4" />
      </button>
    </div>
  )
}

Step 3: Create ChatPanel.tsx

// frontend/src/components/ai-chat/ChatPanel.tsx
import { useEffect, useRef } from 'react'
import { ChatMessage } from './ChatMessage'
import { ChatInput } from './ChatInput'
import { Spinner } from '@/components/common/Spinner'
import type { ChatMessage as ChatMessageType } from '@/types'

interface ChatPanelProps {
  messages: ChatMessageType[]
  isResponding: boolean
  onSendMessage: (content: string) => void
  disabled?: boolean
}

export function ChatPanel({ messages, isResponding, onSendMessage, disabled }: ChatPanelProps) {
  const scrollRef = useRef<HTMLDivElement>(null)

  // Auto-scroll to bottom on new messages
  useEffect(() => {
    if (scrollRef.current) {
      scrollRef.current.scrollTop = scrollRef.current.scrollHeight
    }
  }, [messages, isResponding])

  return (
    <div className="flex h-full flex-col">
      {/* Messages */}
      <div ref={scrollRef} className="flex-1 overflow-y-auto px-4 py-4 space-y-4">
        {messages.map((msg, i) => (
          <ChatMessage key={i} message={msg} />
        ))}
        {isResponding && (
          <div className="flex items-center gap-2 text-sm text-muted-foreground">
            <Spinner size="sm" />
            <span>Thinking...</span>
          </div>
        )}
      </div>

      {/* Input */}
      <ChatInput
        onSend={onSendMessage}
        disabled={disabled || isResponding}
        placeholder={isResponding ? 'Waiting for response...' : 'Type a message...'}
      />
    </div>
  )
}

Step 4: Create PhaseIndicator.tsx

// frontend/src/components/ai-chat/PhaseIndicator.tsx
import { cn } from '@/lib/utils'
import type { InterviewPhase } from '@/types'

const PHASES: { key: InterviewPhase; label: string }[] = [
  { key: 'scoping', label: 'Scoping' },
  { key: 'discovery', label: 'Discovery' },
  { key: 'enrichment', label: 'Enrichment' },
  { key: 'review', label: 'Review' },
  { key: 'generation', label: 'Generate' },
]

interface PhaseIndicatorProps {
  currentPhase: InterviewPhase
}

export function PhaseIndicator({ currentPhase }: PhaseIndicatorProps) {
  const currentIndex = PHASES.findIndex((p) => p.key === currentPhase)

  return (
    <div className="flex items-center gap-1">
      {PHASES.map((phase, i) => {
        const isActive = phase.key === currentPhase
        const isCompleted = i < currentIndex

        return (
          <div key={phase.key} className="flex items-center">
            {i > 0 && (
              <div
                className={cn(
                  'mx-1 h-px w-4',
                  isCompleted ? 'bg-primary' : 'bg-border'
                )}
              />
            )}
            <span
              className={cn(
                'font-label text-[0.6875rem] uppercase tracking-wide px-2 py-0.5 rounded',
                isActive && 'text-primary bg-primary/10 font-medium',
                isCompleted && 'text-primary',
                !isActive && !isCompleted && 'text-muted-foreground'
              )}
            >
              {phase.label}
            </span>
          </div>
        )
      })}
    </div>
  )
}

Step 5: Create ChatToolbar.tsx

// frontend/src/components/ai-chat/ChatToolbar.tsx
import { Sparkles, Download, RotateCcw, ArrowRight } from 'lucide-react'
import { PhaseIndicator } from './PhaseIndicator'
import { cn } from '@/lib/utils'
import type { InterviewPhase } from '@/types'

interface ChatToolbarProps {
  currentPhase: InterviewPhase
  status: 'idle' | 'active' | 'completed' | 'abandoned'
  isGenerating: boolean
  hasGeneratedTree: boolean
  onGenerate: () => void
  onImport: () => void
  onReset: () => void
}

export function ChatToolbar({
  currentPhase,
  status,
  isGenerating,
  hasGeneratedTree,
  onGenerate,
  onImport,
  onReset,
}: ChatToolbarProps) {
  return (
    <div className="flex items-center justify-between border-b border-border bg-card px-4 py-2">
      <div className="flex items-center gap-3">
        <div className="flex items-center gap-2 text-sm font-medium text-foreground">
          <Sparkles className="h-4 w-4 text-primary" />
          Build with AI
        </div>
        <PhaseIndicator currentPhase={currentPhase} />
      </div>

      <div className="flex items-center gap-2">
        {status === 'active' && !hasGeneratedTree && (
          <button
            onClick={onGenerate}
            disabled={isGenerating}
            className={cn(
              'flex items-center gap-1.5 rounded-lg px-3 py-1.5 text-sm font-medium',
              'bg-gradient-brand text-white shadow-lg shadow-primary/20',
              'hover:opacity-90 transition-opacity',
              'disabled:opacity-50 disabled:cursor-not-allowed'
            )}
          >
            <Download className="h-3.5 w-3.5" />
            {isGenerating ? 'Generating...' : 'Generate Tree'}
          </button>
        )}

        {hasGeneratedTree && (
          <button
            onClick={onImport}
            className={cn(
              'flex items-center gap-1.5 rounded-lg px-3 py-1.5 text-sm font-medium',
              'bg-gradient-brand text-white shadow-lg shadow-primary/20',
              'hover:opacity-90 transition-opacity'
            )}
          >
            <ArrowRight className="h-3.5 w-3.5" />
            Import to Editor
          </button>
        )}

        <button
          onClick={onReset}
          className="flex items-center gap-1.5 rounded-lg border border-border px-3 py-1.5 text-sm text-muted-foreground hover:bg-accent hover:text-foreground transition-colors"
        >
          <RotateCcw className="h-3.5 w-3.5" />
          Start Over
        </button>
      </div>
    </div>
  )
}

Step 6: Create EmptyPreview.tsx

// frontend/src/components/ai-chat/EmptyPreview.tsx
import { TreeDeciduous } from 'lucide-react'

export function EmptyPreview() {
  return (
    <div className="flex h-full flex-col items-center justify-center p-8 text-center">
      <TreeDeciduous className="h-12 w-12 text-muted-foreground/30 mb-4" />
      <h3 className="text-sm font-medium text-muted-foreground mb-1">Flow Preview</h3>
      <p className="text-xs text-muted-foreground/70 max-w-48">
        Your flow will appear here as you describe it to the AI
      </p>
    </div>
  )
}

Step 7: Create StaticTreePreview.tsx

This wraps TreePreviewNode without depending on useTreeEditorStore.

// frontend/src/components/ai-chat/StaticTreePreview.tsx
import { useState, useMemo, useCallback } from 'react'
import { TreePreviewNode } from '@/components/tree-preview/TreePreviewNode'
import type { SharedLinksMap } from '@/components/tree-preview/TreePreviewPanel'
import type { TreeStructure } from '@/types'

interface StaticTreePreviewProps {
  tree: TreeStructure
  name?: string
}

function findNodeInTree(nodeId: string, tree: TreeStructure): TreeStructure | null {
  if (tree.id === nodeId) return tree
  if (tree.children) {
    for (const child of tree.children) {
      const found = findNodeInTree(nodeId, child)
      if (found) return found
    }
  }
  return null
}

function buildSharedLinksMap(node: TreeStructure, map: SharedLinksMap = new Map()): SharedLinksMap {
  const nodeLabel = node.type === 'decision' ? node.question : node.title
  if (node.type === 'decision' && node.options) {
    for (const opt of node.options) {
      if (opt.next_node_id) {
        const existing = map.get(opt.next_node_id) || []
        existing.push({ id: node.id, label: nodeLabel || 'Untitled' })
        map.set(opt.next_node_id, existing)
      }
    }
  }
  if (node.type === 'action' && node.next_node_id) {
    const existing = map.get(node.next_node_id) || []
    existing.push({ id: node.id, label: nodeLabel || 'Untitled' })
    map.set(node.next_node_id, existing)
  }
  if (node.children) {
    for (const child of node.children) {
      buildSharedLinksMap(child, map)
    }
  }
  return map
}

export function StaticTreePreview({ tree, name }: StaticTreePreviewProps) {
  const [selectedNodeId, setSelectedNodeId] = useState<string | null>(null)

  const findNode = useCallback(
    (nodeId: string) => findNodeInTree(nodeId, tree),
    [tree]
  )

  const sharedLinksMap = useMemo(() => buildSharedLinksMap(tree), [tree])

  return (
    <div className="flex h-full flex-col">
      <div className="border-b border-border px-4 py-2">
        <h3 className="text-sm font-semibold text-foreground">
          Preview: {name || 'Untitled Flow'}
        </h3>
        <p className="text-xs text-muted-foreground">
          Click a node to select
        </p>
      </div>
      <div className="flex-1 overflow-auto p-4">
        <div className="inline-block min-w-full">
          <TreePreviewNode
            node={tree}
            selectedNodeId={selectedNodeId}
            onSelect={setSelectedNodeId}
            depth={0}
            findNode={findNode}
            sharedLinksMap={sharedLinksMap}
          />
        </div>
      </div>
    </div>
  )
}

Step 8: Commit

git add frontend/src/components/ai-chat/
git commit -m "feat: add AI chat builder components — ChatPanel, ChatInput, ChatMessage, PhaseIndicator, ChatToolbar, EmptyPreview, StaticTreePreview"

Task 10: Main Page & Routing

Files:

  • Create: frontend/src/pages/AIChatBuilderPage.tsx
  • Modify: frontend/src/router.tsx (add lazy import + route)
  • Modify: frontend/src/pages/TreeLibraryPage.tsx (add "Build with AI" button)

Step 1: Create the page

// frontend/src/pages/AIChatBuilderPage.tsx
import { useCallback, useEffect } from 'react'
import { useNavigate, useSearchParams } from 'react-router-dom'
import { useAIChatStore } from '@/store/aiChatStore'
import { ChatPanel } from '@/components/ai-chat/ChatPanel'
import { ChatToolbar } from '@/components/ai-chat/ChatToolbar'
import { EmptyPreview } from '@/components/ai-chat/EmptyPreview'
import { StaticTreePreview } from '@/components/ai-chat/StaticTreePreview'
import { Spinner } from '@/components/common/Spinner'
import { getTreeEditorPath } from '@/lib/routing'
import { toast } from '@/lib/toast'
import type { TreeStructure } from '@/types'

export function AIChatBuilderPage() {
  const navigate = useNavigate()
  const [searchParams] = useSearchParams()
  const flowType = searchParams.get('type') === 'procedural' ? 'procedural' : 'troubleshooting'

  const {
    sessionId,
    status,
    currentPhase,
    messages,
    isResponding,
    workingTree,
    treeMetadata,
    generatedTree,
    isGenerating,
    error,
    startSession,
    sendMessage,
    generateTree,
    importToEditor,
    abandonSession,
    reset,
  } = useAIChatStore()

  // Start session on mount if no active session
  useEffect(() => {
    if (!sessionId && status === 'idle') {
      startSession(flowType as 'troubleshooting' | 'procedural')
    }
  }, [sessionId, status, flowType, startSession])

  const handleSendMessage = useCallback(
    (content: string) => {
      sendMessage(content)
    },
    [sendMessage]
  )

  const handleGenerate = useCallback(() => {
    generateTree()
  }, [generateTree])

  const handleImport = useCallback(async () => {
    try {
      const treeId = await importToEditor({
        name: treeMetadata?.name,
        description: treeMetadata?.description,
        tags: treeMetadata?.tags,
      })
      const path = getTreeEditorPath(treeId, flowType)
      navigate(path)
      toast.success('Flow imported to editor')
    } catch {
      toast.error('Failed to import flow')
    }
  }, [importToEditor, treeMetadata, flowType, navigate])

  const handleReset = useCallback(() => {
    abandonSession()
    startSession(flowType as 'troubleshooting' | 'procedural')
  }, [abandonSession, startSession, flowType])

  // Show error toast
  useEffect(() => {
    if (error) {
      toast.error(error)
    }
  }, [error])

  if (status === 'idle' && !sessionId) {
    return (
      <div className="flex h-full items-center justify-center">
        <Spinner />
      </div>
    )
  }

  const previewTree = (generatedTree || workingTree) as TreeStructure | null

  return (
    <div className="flex h-full flex-col">
      <ChatToolbar
        currentPhase={currentPhase}
        status={status}
        isGenerating={isGenerating}
        hasGeneratedTree={!!generatedTree}
        onGenerate={handleGenerate}
        onImport={handleImport}
        onReset={handleReset}
      />

      <div className="flex flex-1 overflow-hidden">
        {/* Left panel: Chat (60%) */}
        <div className="flex w-3/5 flex-col border-r border-border">
          <ChatPanel
            messages={messages}
            isResponding={isResponding}
            onSendMessage={handleSendMessage}
            disabled={status !== 'active'}
          />
        </div>

        {/* Right panel: Tree preview (40%) */}
        <div className="w-2/5 overflow-hidden bg-background">
          {previewTree ? (
            <StaticTreePreview
              tree={previewTree}
              name={treeMetadata?.name}
            />
          ) : (
            <EmptyPreview />
          )}
        </div>
      </div>
    </div>
  )
}

export default AIChatBuilderPage

Step 2: Add route to router.tsx

Add lazy import after line 35 (const StepLibraryPage = ...):

const AIChatBuilderPage = lazy(() => import('@/pages/AIChatBuilderPage'))

Add route after the step-library route block (after line 255):

      {
        path: 'ai/chat',
        element: (
          <Suspense fallback={<PageLoader />}>
            <AIChatBuilderPage />
          </Suspense>
        ),
      },

Step 3: Add "Build with AI" navigation to TreeLibraryPage.tsx

Find the section near line 560 where {showAIBuilder && ...} renders the modal. We need to add a navigation button that goes to /ai/chat instead. Look for the area where CreateFlowDropdown is rendered and add a "Build with AI" button next to it.

The exact location depends on the template markup — find the Create button area and add:

{canCreateTrees && aiEnabled && (
  <button
    onClick={() => navigate('/ai/chat')}
    className="flex items-center gap-2 rounded-lg bg-gradient-brand px-4 py-2 text-sm font-medium text-white shadow-lg shadow-primary/20 hover:opacity-90 transition-opacity"
  >
    <Sparkles className="h-4 w-4" />
    Build with AI
  </button>
)}

Add Sparkles to the Lucide import at line 3. Add useNavigate if not already imported.

Step 4: Run frontend build to verify

Run: cd frontend && npm run build

Expected: Build succeeds with no TypeScript errors.

Step 5: Commit

git add frontend/src/pages/AIChatBuilderPage.tsx frontend/src/router.tsx frontend/src/pages/TreeLibraryPage.tsx
git commit -m "feat: add AI chat builder page with routing and library entry point"

Phase 3: Tree Preview & Import Integration

Task 11: End-to-End Wiring & Polish

Files:

  • May need minor adjustments based on Phase 2 testing

Step 1: Manual testing checklist

Start the dev servers and test:

  1. Navigate to Tree Library → click "Build with AI" → lands on /ai/chat
  2. Session starts, AI greeting appears in chat panel
  3. Type a message → AI responds (loading spinner while waiting)
  4. After discovery phase: working tree appears in right panel
  5. Click "Generate Tree" → final tree generated (loading state, then preview updates)
  6. Click "Import to Editor" → navigates to Tree Editor with the tree loaded
  7. Click "Start Over" → resets and starts fresh session
  8. Refresh page mid-conversation → session can be resumed (stretch goal — requires storing sessionId in URL params or localStorage)

Step 2: Fix any issues found during testing

Common things to fix:

  • TreePreviewNode prop mismatches (check the findNode return type)
  • Markdown rendering issues in AI responses
  • Scroll behavior (auto-scroll to bottom on new messages)
  • Error states (network errors, AI timeouts)

Step 3: Commit any fixes

git add -A
git commit -m "fix: polish AI chat builder based on manual testing"

Phase 4: Polish

Task 12: Session Resume via URL

Files:

  • Modify: frontend/src/pages/AIChatBuilderPage.tsx

Step 1: Store sessionId in URL search params

After a session starts, update the URL: ?session=<id>. On mount, check for ?session= param and call resumeSession() instead of startSession().

// In AIChatBuilderPage.tsx, update the mount effect:
useEffect(() => {
  const resumeId = searchParams.get('session')
  if (resumeId && !sessionId) {
    resumeSession(resumeId)
  } else if (!sessionId && status === 'idle') {
    startSession(flowType as 'troubleshooting' | 'procedural')
  }
}, []) // eslint-disable-line react-hooks/exhaustive-deps

// After session starts, update URL:
useEffect(() => {
  if (sessionId) {
    const params = new URLSearchParams(searchParams)
    params.set('session', sessionId)
    navigate(`/ai/chat?${params.toString()}`, { replace: true })
  }
}, [sessionId]) // eslint-disable-line react-hooks/exhaustive-deps

Step 2: Commit

git add frontend/src/pages/AIChatBuilderPage.tsx
git commit -m "feat: add session resume via URL parameter for AI chat builder"

Task 13: Responsive Layout

Files:

  • Modify: frontend/src/pages/AIChatBuilderPage.tsx

Step 1: Stack panels vertically on narrow screens

Replace the fixed w-3/5 / w-2/5 split with responsive classes:

<div className="flex flex-1 overflow-hidden flex-col lg:flex-row">
  {/* Left panel: Chat */}
  <div className="flex flex-1 lg:w-3/5 flex-col border-b lg:border-b-0 lg:border-r border-border min-h-0">
    <ChatPanel ... />
  </div>

  {/* Right panel: Tree preview — collapsible on mobile */}
  <div className="hidden lg:block lg:w-2/5 overflow-hidden bg-background">
    {previewTree ? <StaticTreePreview ... /> : <EmptyPreview />}
  </div>
</div>

On mobile, the preview is hidden (the chat is the primary interface). Users can generate and import without seeing the preview.

Step 2: Run build

Run: cd frontend && npm run build

Step 3: Commit

git add frontend/src/pages/AIChatBuilderPage.tsx
git commit -m "feat: add responsive layout for AI chat builder"

Task 14: Run Backend Tests & Verify

Step 1: Run full backend test suite

Run: cd backend && python -m pytest --override-ini="addopts=" -v

Expected: All existing tests still pass + new AI chat tests pass.

Step 2: Run frontend build

Run: cd frontend && npm run build

Expected: Clean build with no errors.

Step 3: Final commit if needed

Fix any issues and commit.


Summary of All Files

New Files (15)

# File Task
1 backend/app/models/ai_chat_session.py Task 1
2 backend/alembic/versions/XXX_add_ai_chat_sessions.py Task 1
3 backend/app/schemas/ai_chat.py Task 2
4 backend/app/core/ai_chat_service.py Task 3
5 backend/app/api/endpoints/ai_chat.py Task 4
6 backend/tests/test_ai_chat.py Task 5
7 frontend/src/types/ai-chat.ts Task 6
8 frontend/src/api/aiChat.ts Task 7
9 frontend/src/store/aiChatStore.ts Task 8
10 frontend/src/components/ai-chat/ChatMessage.tsx Task 9
11 frontend/src/components/ai-chat/ChatInput.tsx Task 9
12 frontend/src/components/ai-chat/ChatPanel.tsx Task 9
13 frontend/src/components/ai-chat/PhaseIndicator.tsx Task 9
14 frontend/src/components/ai-chat/ChatToolbar.tsx Task 9
15 frontend/src/components/ai-chat/EmptyPreview.tsx Task 9
16 frontend/src/components/ai-chat/StaticTreePreview.tsx Task 9
17 frontend/src/pages/AIChatBuilderPage.tsx Task 10

Modified Files (6)

# File Change Task
1 backend/app/models/__init__.py Import AIChatSession Task 1
2 backend/app/api/router.py Include ai_chat.router Task 4
3 backend/app/core/ai_quota_service.py Add chat types to daily limit Task 4
4 frontend/src/types/index.ts Export ai-chat types Task 6
5 frontend/src/api/index.ts Export aiChat module Task 7
6 frontend/src/router.tsx Add /ai/chat route Task 10
7 frontend/src/pages/TreeLibraryPage.tsx Add "Build with AI" button Task 10