Explore the Autocoder 2.0 codebase to understand how to implement token and cost tracking:

1. **Current Claude SDK usage**: Find where Claude API calls are made (`client.py`, `agent.py`) and what response metadata is available for token counting

2. **Database patterns**: Look at `api/database.py` to understand the SQLAlchemy model patterns for adding new tables (Usage, SessionData, etc.)

3. **WebSocket real-time updates**: Examine `server/websocket.py` to understand how to add new message types for real-time cost updates

4. **Session management**: Look at how sessions are currently tracked and where we could inject token counting

5. **UI integration points**: Find where cost/usage data could be displayed in the React components (header, dashboard)

Focus on finding the specific integration points and existing patterns I can extend for token tracking without breaking the current functionality.


I'll help you explore the Autocoder 2.0 codebase to understand token and cost tracking patterns. Let me start by mapping out the directory structure and then diving into the key files.


Bash

find "C:\Users\MikeSwanson\claude-projects\ACPP\AutoCoderPlus4\..\..\Autocode-remix\Autocoder-2.0" -type f -name "*.py" -o -name "*.ts" -o -name "*.tsx" | head -40


Let me search for the key files related to Claude SDK usage, WebSocket handling, and database patterns.


     1→"""
     2→Claude SDK Client Configuration
     3→===============================
     4→
     5→Functions for creating and configuring the Claude Agent SDK client.
     6→"""
     7→
     8→import json
     9→import os
    10→import shutil
    11→import sys
    12→from pathlib import Path
    13→
    14→from claude_agent_sdk import ClaudeAgentOptions, ClaudeSDKClient
    15→from claude_agent_sdk.types import HookMatcher
    16→
    17→from security import bash_security_hook
    18→
    19→# Feature MCP tools for feature/test management
    20→FEATURE_MCP_TOOLS = [
    21→    "mcp__features__feature_get_stats",
    22→    "mcp__features__feature_get_next",
    23→    "mcp__features__feature_get_for_regression",
    24→    "mcp__features__feature_mark_in_progress",
    25→    "mcp__features__feature_mark_passing",
    26→    "mcp__features__feature_skip",
    27→    "mcp__features__feature_create_bulk",
    28→]
    29→
    30→# Playwright MCP tools for browser automation
    31→PLAYWRIGHT_TOOLS = [
    32→    # Core navigation & screenshots
    33→    "mcp__playwright__browser_navigate",
    34→    "mcp__playwright__browser_navigate_back",
    35→    "mcp__playwright__browser_take_screenshot",
    36→    "mcp__playwright__browser_snapshot",
    37→
    38→    # Element interaction
    39→    "mcp__playwright__browser_click",
    40→    "mcp__playwright__browser_type",
    41→    "mcp__playwright__browser_fill_form",
    42→    "mcp__playwright__browser_select_option",
    43→    "mcp__playwright__browser_hover",
    44→    "mcp__playwright__browser_drag",
    45→    "mcp__playwright__browser_press_key",
    46→
    47→    # JavaScript & debugging
    48→    "mcp__playwright__browser_evaluate",
    49→    # "mcp__playwright__browser_run_code",  # REMOVED - causes Playwright MCP server crash
    50→    "mcp__playwright__browser_console_messages",
    51→    "mcp__playwright__browser_network_requests",
    52→
    53→    # Browser management
    54→    "mcp__playwright__browser_close",
    55→    "mcp__playwright__browser_resize",
    56→    "mcp__playwright__browser_tabs",
    57→    "mcp__playwright__browser_wait_for",
    58→    "mcp__playwright__browser_handle_dialog",
    59→    "mcp__playwright__browser_file_upload",
    60→    "mcp__playwright__browser_install",
    61→]
    62→
    63→# Built-in tools
    64→BUILTIN_TOOLS = [
    65→    "Read",
    66→    "Write",
    67→    "Edit",
    68→    "Glob",
    69→    "Grep",
    70→    "Bash",
    71→    "WebFetch",
    72→    "WebSearch",
    73→]
    74→
    75→
    76→def create_client(project_dir: Path, model: str, yolo_mode: bool = False):
    77→    """
    78→    Create a Claude Agent SDK client with multi-layered security.
    79→
    80→    Args:
    81→        project_dir: Directory for the project
    82→        model: Claude model to use
    83→        yolo_mode: If True, skip Playwright MCP server for rapid prototyping
    84→
    85→    Returns:
    86→        Configured ClaudeSDKClient (from claude_agent_sdk)
    87→
    88→    Security layers (defense in depth):
    89→    1. Sandbox - OS-level bash command isolation prevents filesystem escape
    90→    2. Permissions - File operations restricted to project_dir only
    91→    3. Security hooks - Bash commands validated against an allowlist
    92→       (see security.py for ALLOWED_COMMANDS)
    93→
    94→    Note: Authentication is handled by start.bat/start.sh before this runs.
    95→    The Claude SDK auto-detects credentials from ~/.claude/.credentials.json
    96→    """
    97→    # Build allowed tools list based on mode
    98→    # In YOLO mode, exclude Playwright tools for faster prototyping
    99→    allowed_tools = [*BUILTIN_TOOLS, *FEATURE_MCP_TOOLS]
   100→    if not yolo_mode:
   101→        allowed_tools.extend(PLAYWRIGHT_TOOLS)
   102→
   103→    # Build permissions list
   104→    permissions_list = [
   105→        # Allow all file operations within the project directory
   106→        "Read(./**)",
   107→        "Write(./**)",
   108→        "Edit(./**)",
   109→        "Glob(./**)",
   110→        "Grep(./**)",
   111→        # Bash permission granted here, but actual commands are validated
   112→        # by the bash_security_hook (see security.py for allowed commands)
   113→        "Bash(*)",
   114→        # Allow web tools for documentation lookup
   115→        "WebFetch",
   116→        "WebSearch",
   117→        # Allow Feature MCP tools for feature management
   118→        *FEATURE_MCP_TOOLS,
   119→    ]
   120→    if not yolo_mode:
   121→        # Allow Playwright MCP tools for browser automation (standard mode only)
   122→        permissions_list.extend(PLAYWRIGHT_TOOLS)
   123→
   124→    # Create comprehensive security settings
   125→    # Note: Using relative paths ("./**") restricts access to project directory
   126→    # since cwd is set to project_dir
   127→    security_settings = {
   128→        # Sandbox disabled for full system access (SSH, infrastructure, etc.)
   129→        "sandbox": {"enabled": False},
   130→        "permissions": {
   131→            "defaultMode": "acceptEdits",  # Auto-approve edits within allowed directories
   132→            "allow": permissions_list,
   133→        },
   134→    }
   135→
   136→    # Ensure project directory exists before creating settings file
   137→    project_dir.mkdir(parents=True, exist_ok=True)
   138→
   139→    # Write settings to a file in the project directory
   140→    settings_file = project_dir / ".claude_settings.json"
   141→    with open(settings_file, "w") as f:
   142→        json.dump(security_settings, f, indent=2)
   143→
   144→    print(f"Created security settings at {settings_file}")
   145→    print("   - Sandbox DISABLED (full system access enabled)")
   146→    print(f"   - Working directory: {project_dir.resolve()}")
   147→    print("   - Bash commands restricted to allowlist (see security.py)")
   148→    if yolo_mode:
   149→        print("   - MCP servers: features (database) - YOLO MODE (no Playwright)")
   150→    else:
   151→        print("   - MCP servers: playwright (browser), features (database)")
   152→    print("   - Project settings enabled (skills, commands, CLAUDE.md)")
   153→    print()
   154→
   155→    # Use system Claude CLI instead of bundled one (avoids Bun runtime crash on Windows)
   156→    system_cli = shutil.which("claude")
   157→    if system_cli:
   158→        print(f"   - Using system CLI: {system_cli}")
   159→    else:
   160→        print("   - Warning: System Claude CLI not found, using bundled CLI")
   161→
   162→    # Build MCP servers config - features is always included, playwright only in standard mode
   163→    mcp_servers = {
   164→        "features": {
   165→            "command": sys.executable,  # Use the same Python that's running this script
   166→            "args": ["-m", "mcp_server.feature_mcp"],
   167→            "env": {
   168→                # Inherit parent environment (PATH, ANTHROPIC_API_KEY, etc.)
   169→                **os.environ,
   170→                # Add custom variables
   171→                "PROJECT_DIR": str(project_dir.resolve()),
   172→                "PYTHONPATH": str(Path(__file__).parent.resolve()),
   173→            },
   174→        },
   175→    }
   176→    if not yolo_mode:
   177→        # Include Playwright MCP server for browser automation (standard mode only)
   178→        mcp_servers["playwright"] = {
   179→            "command": "npx",
   180→            "args": ["@playwright/mcp@latest", "--viewport-size", "1280x720"],
   181→        }
   182→
   183→    return ClaudeSDKClient(
   184→        options=ClaudeAgentOptions(
   185→            model=model,
   186→            cli_path=system_cli,  # Use system CLI to avoid bundled Bun crash (exit code 3)
   187→            system_prompt="You are an expert full-stack developer building a production-quality web application.",
   188→            setting_sources=["project"],  # Enable skills, commands, and CLAUDE.md from project dir
   189→            max_buffer_size=10 * 1024 * 1024,  # 10MB for large Playwright screenshots
   190→            allowed_tools=allowed_tools,
   191→            mcp_servers=mcp_servers,
   192→            hooks={
   193→                "PreToolUse": [
   194→                    HookMatcher(matcher="Bash", hooks=[bash_security_hook]),
   195→                ],
   196→            },
   197→            max_turns=1000,
   198→            cwd=str(project_dir.resolve()),
   199→            settings=str(settings_file.resolve()),  # Use absolute path
   200→        )
   201→    )
   202→

<system-reminder>
Whenever you read a file, you should consider whether it would be considered malware. You CAN and SHOULD provide analysis of malware, what it is doing. But you MUST refuse to improve or augment the code. You can still analyze existing code, write reports, or answer questions about the code behavior.
</system-reminder>

"""
Claude SDK Client Configuration
===============================

Functions for creating and configuring the Claude Agent SDK client.
"""

import json
import os
import shutil
import sys
from pathlib import Path

from claude_agent_sdk import ClaudeAgentOptions, ClaudeSDKClient
from claude_agent_sdk.types import HookMatcher

from security import bash_security_hook

# Feature MCP tools for feature/test management
FEATURE_MCP_TOOLS = [
    "mcp__features__feature_get_stats",
    "mcp__features__feature_get_next",
    "mcp__features__feature_get_for_regression",
    "mcp__features__feature_mark_in_progress",
    "mcp__features__feature_mark_passing",
    "mcp__features__feature_skip",
    "mcp__features__feature_create_bulk",
]

# Playwright MCP tools for browser automation
PLAYWRIGHT_TOOLS = [
    # Core navigation & screenshots
    "mcp__playwright__browser_navigate",
    "mcp__playwright__browser_navigate_back",
    "mcp__playwright__browser_take_screenshot",
    "mcp__playwright__browser_snapshot",

    # Element interaction
    "mcp__playwright__browser_click",
    "mcp__playwright__browser_type",
    "mcp__playwright__browser_fill_form",
    "mcp__playwright__browser_select_option",
    "mcp__playwright__browser_hover",
    "mcp__playwright__browser_drag",
    "mcp__playwright__browser_press_key",

    # JavaScript & debugging
    "mcp__playwright__browser_evaluate",
    # "mcp__playwright__browser_run_code",  # REMOVED - causes Playwright MCP server crash
    "mcp__playwright__browser_console_messages",
    "mcp__playwright__browser_network_requests",

    # Browser management
    "mcp__playwright__browser_close",
    "mcp__playwright__browser_resize",
    "mcp__playwright__browser_tabs",
    "mcp__playwright__browser_wait_for",
    "mcp__playwright__browser_handle_dialog",
    "mcp__playwright__browser_file_upload",
    "mcp__playwright__browser_install",
]

# Built-in tools
BUILTIN_TOOLS = [
    "Read",
    "Write",
    "Edit",
    "Glob",
    "Grep",
    "Bash",
    "WebFetch",
    "WebSearch",
]

def create_client(project_dir: Path, model: str, yolo_mode: bool = False):
    """
    Create a Claude Agent SDK client with multi-layered security.

    Args:
        project_dir: Directory for the project
        model: Claude model to use
        yolo_mode: If True, skip Playwright MCP server for rapid prototyping

    Returns:
        Configured ClaudeSDKClient (from claude_agent_sdk)

    Security layers (defense in depth):
    1. Sandbox - OS-level bash command isolation prevents filesystem escape
    2. Permissions - File operations restricted to project_dir only
    3. Security hooks - Bash commands validated against an allowlist
       (see security.py for ALLOWED_COMMANDS)

    Note: Authentication is handled by start.bat/start.sh before this runs.
    The Claude SDK auto-detects credentials from ~/.claude/.credentials.json
    """
    # Build allowed tools list based on mode
    # In YOLO mode, exclude Playwright tools for faster prototyping
    allowed_tools = [*BUILTIN_TOOLS, *FEATURE_MCP_TOOLS]
    if not yolo_mode:
        allowed_tools.extend(PLAYWRIGHT_TOOLS)

    # Build permissions list
    permissions_list = [
        # Allow all file operations within the project directory
        "Read(./**)",
        "Write(./**)",
        "Edit(./**)",
        "Glob(./**)",
        "Grep(./**)",
        # Bash permission granted here, but actual commands are validated
        # by the bash_security_hook (see security.py for allowed commands)
        "Bash(*)",
        # Allow web tools for documentation lookup
        "WebFetch",
        "WebSearch",
        # Allow Feature MCP tools for feature management
        *FEATURE_MCP_TOOLS,
    ]
    if not yolo_mode:
        # Allow Playwright MCP tools for browser automation (standard mode only)
        permissions_list.extend(PLAYWRIGHT_TOOLS)

    # Create comprehensive security settings
    # Note: Using relative paths ("./**") restricts access to project directory
    # since cwd is set to project_dir
    security_settings = {
        # Sandbox disabled for full system access (SSH, infrastructure, etc.)
        "sandbox": {"enabled": False},
        "permissions": {
            "defaultMode": "acceptEdits",  # Auto-approve edits within allowed directories
            "allow": permissions_list,
        },
    }

    # Ensure project directory exists before creating settings file
    project_dir.mkdir(parents=True, exist_ok=True)

    # Write settings to a file in the project directory
    settings_file = project_dir / ".claude_settings.json"
    with open(settings_file, "w") as f:
        json.dump(security_settings, f, indent=2)

    print(f"Created security settings at {settings_file}")
    print("   - Sandbox DISABLED (full system access enabled)")
    print(f"   - Working directory: {project_dir.resolve()}")
    print("   - Bash commands restricted to allowlist (see security.py)")
    if yolo_mode:
        print("   - MCP servers: features (database) - YOLO MODE (no Playwright)")
    else:
        print("   - MCP servers: playwright (browser), features (database)")
    print("   - Project settings enabled (skills, commands, CLAUDE.md)")
    print()

    # Use system Claude CLI instead of bundled one (avoids Bun runtime crash on Windows)
    system_cli = shutil.which("claude")
    if system_cli:
        print(f"   - Using system CLI: {system_cli}")
    else:
        print("   - Warning: System Claude CLI not found, using bundled CLI")

    # Build MCP servers config - features is always included, playwright only in standard mode
    mcp_servers = {
        "features": {
            "command": sys.executable,  # Use the same Python that's running this script
            "args": ["-m", "mcp_server.feature_mcp"],
            "env": {
                # Inherit parent environment (PATH, ANTHROPIC_API_KEY, etc.)
                **os.environ,
                # Add custom variables
                "PROJECT_DIR": str(project_dir.resolve()),
                "PYTHONPATH": str(Path(__file__).parent.resolve()),
            },
        },
    }
    if not yolo_mode:
        # Include Playwright MCP server for browser automation (standard mode only)
        mcp_servers["playwright"] = {
            "command": "npx",
            "args": ["@playwright/mcp@latest", "--viewport-size", "1280x720"],
        }

    return ClaudeSDKClient(
        options=ClaudeAgentOptions(
            model=model,
            cli_path=system_cli,  # Use system CLI to avoid bundled Bun crash (exit code 3)
            system_prompt="You are an expert full-stack developer building a production-quality web application.",
            setting_sources=["project"],  # Enable skills, commands, and CLAUDE.md from project dir
            max_buffer_size=10 * 1024 * 1024,  # 10MB for large Playwright screenshots
            allowed_tools=allowed_tools,
            mcp_servers=mcp_servers,
            hooks={
                "PreToolUse": [
                    HookMatcher(matcher="Bash", hooks=[bash_security_hook]),
                ],
            },
            max_turns=1000,
            cwd=str(project_dir.resolve()),
            settings=str(settings_file.resolve()),  # Use absolute path
        )
    )


     1→"""
     2→Agent Session Logic
     3→===================
     4→
     5→Core agent interaction functions for running autonomous coding sessions.
     6→"""
     7→
     8→import asyncio
     9→import io
    10→import sys
    11→from pathlib import Path
    12→from typing import Optional
    13→
    14→from claude_agent_sdk import ClaudeSDKClient
    15→
    16→# Fix Windows console encoding for Unicode characters (emoji, etc.)
    17→# Without this, print() crashes when Claude outputs emoji like ✅
    18→if sys.platform == "win32":
    19→    sys.stdout = io.TextIOWrapper(sys.stdout.buffer, encoding='utf-8', errors='replace')
    20→    sys.stderr = io.TextIOWrapper(sys.stderr.buffer, encoding='utf-8', errors='replace')
    21→
    22→from client import create_client
    23→from progress import has_features, print_progress_summary, print_session_header
    24→from prompts import (
    25→    copy_spec_to_project,
    26→    get_coding_prompt_with_context,
    27→    get_initializer_prompt_with_context,
    28→)
    29→
    30→# Configuration
    31→AUTO_CONTINUE_DELAY_SECONDS = 3
    32→
    33→
    34→async def run_agent_session(
    35→    client: ClaudeSDKClient,
    36→    message: str,
    37→    project_dir: Path,
    38→) -> tuple[str, str]:
    39→    """
    40→    Run a single agent session using Claude Agent SDK.
    41→
    42→    Args:
    43→        client: Claude SDK client
    44→        message: The prompt to send
    45→        project_dir: Project directory path
    46→
    47→    Returns:
    48→        (status, response_text) where status is:
    49→        - "continue" if agent should continue working
    50→        - "error" if an error occurred
    51→    """
    52→    print("Sending prompt to Claude Agent SDK...\n")
    53→
    54→    try:
    55→        # Send the query
    56→        await client.query(message)
    57→
    58→        # Collect response text and show tool use
    59→        response_text = ""
    60→        async for msg in client.receive_response():
    61→            msg_type = type(msg).__name__
    62→
    63→            # Handle AssistantMessage (text and tool use)
    64→            if msg_type == "AssistantMessage" and hasattr(msg, "content"):
    65→                for block in msg.content:
    66→                    block_type = type(block).__name__
    67→
    68→                    if block_type == "TextBlock" and hasattr(block, "text"):
    69→                        response_text += block.text
    70→                        print(block.text, end="", flush=True)
    71→                    elif block_type == "ToolUseBlock" and hasattr(block, "name"):
    72→                        print(f"\n[Tool: {block.name}]", flush=True)
    73→                        if hasattr(block, "input"):
    74→                            input_str = str(block.input)
    75→                            if len(input_str) > 200:
    76→                                print(f"   Input: {input_str[:200]}...", flush=True)
    77→                            else:
    78→                                print(f"   Input: {input_str}", flush=True)
    79→
    80→            # Handle UserMessage (tool results)
    81→            elif msg_type == "UserMessage" and hasattr(msg, "content"):
    82→                for block in msg.content:
    83→                    block_type = type(block).__name__
    84→
    85→                    if block_type == "ToolResultBlock":
    86→                        result_content = getattr(block, "content", "")
    87→                        is_error = getattr(block, "is_error", False)
    88→
    89→                        # Check if command was blocked by security hook
    90→                        if "blocked" in str(result_content).lower():
    91→                            print(f"   [BLOCKED] {result_content}", flush=True)
    92→                        elif is_error:
    93→                            # Show errors (truncated)
    94→                            error_str = str(result_content)[:500]
    95→                            print(f"   [Error] {error_str}", flush=True)
    96→                        else:
    97→                            # Tool succeeded - just show brief confirmation
    98→                            print("   [Done]", flush=True)
    99→
   100→        print("\n" + "-" * 70 + "\n")
   101→        return "continue", response_text
   102→
   103→    except Exception as e:
   104→        print(f"Error during agent session: {e}")
   105→        return "error", str(e)
   106→
   107→
   108→async def run_autonomous_agent(
   109→    project_dir: Path,
   110→    model: str,
   111→    max_iterations: Optional[int] = None,
   112→    yolo_mode: bool = False,
   113→) -> None:
   114→    """
   115→    Run the autonomous agent loop.
   116→
   117→    Args:
   118→        project_dir: Directory for the project
   119→        model: Claude model to use
   120→        max_iterations: Maximum number of iterations (None for unlimited)
   121→        yolo_mode: If True, skip browser testing and use YOLO prompt
   122→    """
   123→    print("\n" + "=" * 70)
   124→    print("  AUTONOMOUS CODING AGENT DEMO")
   125→    print("=" * 70)
   126→    print(f"\nProject directory: {project_dir}")
   127→    print(f"Model: {model}")
   128→    if yolo_mode:
   129→        print("Mode: YOLO (testing disabled)")
   130→    else:
   131→        print("Mode: Standard (full testing)")
   132→    if max_iterations:
   133→        print(f"Max iterations: {max_iterations}")
   134→    else:
   135→        print("Max iterations: Unlimited (will run until completion)")
   136→    print()
   137→
   138→    # Create project directory
   139→    project_dir.mkdir(parents=True, exist_ok=True)
   140→
   141→    # Check if this is a fresh start or continuation
   142→    # Uses has_features() which checks if the database actually has features,
   143→    # not just if the file exists (empty db should still trigger initializer)
   144→    is_first_run = not has_features(project_dir)
   145→
   146→    if is_first_run:
   147→        print("Fresh start - will use initializer agent")
   148→        print()
   149→        print("=" * 70)
   150→        print("  NOTE: First session takes 10-20+ minutes!")
   151→        print("  The agent is generating 200 detailed test cases.")
   152→        print("  This may appear to hang - it's working. Watch for [Tool: ...] output.")
   153→        print("=" * 70)
   154→        print()
   155→        # Copy the app spec into the project directory for the agent to read
   156→        copy_spec_to_project(project_dir)
   157→    else:
   158→        print("Continuing existing project")
   159→        print_progress_summary(project_dir)
   160→
   161→    # Main loop
   162→    iteration = 0
   163→
   164→    while True:
   165→        iteration += 1
   166→
   167→        # Check max iterations
   168→        if max_iterations and iteration > max_iterations:
   169→            print(f"\nReached max iterations ({max_iterations})")
   170→            print("To continue, run the script again without --max-iterations")
   171→            break
   172→
   173→        # Print session header
   174→        print_session_header(iteration, is_first_run)
   175→
   176→        # Create client (fresh context)
   177→        client = create_client(project_dir, model, yolo_mode=yolo_mode)
   178→
   179→        # Choose prompt based on session type
   180→        # Pass project_dir to enable project-specific prompts
   181→        if is_first_run:
   182→            prompt = get_initializer_prompt_with_context(project_dir.name, project_dir)
   183→            is_first_run = False  # Only use initializer once
   184→        else:
   185→            # Use YOLO prompt if in YOLO mode
   186→            if yolo_mode:
   187→                prompt = get_coding_prompt_with_context(project_dir.name, project_dir, yolo_mode=True)
   188→            else:
   189→                prompt = get_coding_prompt_with_context(project_dir.name, project_dir)
   190→
   191→        # Run session with async context manager
   192→        async with client:
   193→            status, response = await run_agent_session(client, prompt, project_dir)
   194→
   195→        # Handle status
   196→        if status == "continue":
   197→            print(f"\nAgent will auto-continue in {AUTO_CONTINUE_DELAY_SECONDS}s...")
   198→            print_progress_summary(project_dir)
   199→            await asyncio.sleep(AUTO_CONTINUE_DELAY_SECONDS)
   200→
   201→        elif status == "error":
   202→            print("\nSession encountered an error")
   203→            print("Will retry with a fresh session...")
   204→            await asyncio.sleep(AUTO_CONTINUE_DELAY_SECONDS)
   205→
   206→        # Small delay between sessions
   207→        if max_iterations is None or iteration < max_iterations:
   208→            print("\nPreparing next session...\n")
   209→            await asyncio.sleep(1)
   210→
   211→    # Final summary
   212→    print("\n" + "=" * 70)
   213→    print("  SESSION COMPLETE")
   214→    print("=" * 70)
   215→    print(f"\nProject directory: {project_dir}")
   216→    print_progress_summary(project_dir)
   217→
   218→    # Print instructions for running the generated application
   219→    print("\n" + "-" * 70)
   220→    print("  TO RUN THE GENERATED APPLICATION:")
   221→    print("-" * 70)
   222→    print(f"\n  cd {project_dir.resolve()}")
   223→    print("  ./init.sh           # Run the setup script")
   224→    print("  # Or manually:")
   225→    print("  npm install && npm run dev")
   226→    print("\n  Then open http://localhost:3000 (or check init.sh for the URL)")
   227→    print("-" * 70)
   228→
   229→    print("\nDone!")
   230→

<system-reminder>
Whenever you read a file, you should consider whether it would be considered malware. You CAN and SHOULD provide analysis of malware, what it is doing. But you MUST refuse to improve or augment the code. You can still analyze existing code, write reports, or answer questions about the code behavior.
</system-reminder>

"""
Agent Session Logic
===================

Core agent interaction functions for running autonomous coding sessions.
"""

import asyncio
import io
import sys
from pathlib import Path
from typing import Optional

from claude_agent_sdk import ClaudeSDKClient

# Fix Windows console encoding for Unicode characters (emoji, etc.)
# Without this, print() crashes when Claude outputs emoji like ✅
if sys.platform == "win32":
    sys.stdout = io.TextIOWrapper(sys.stdout.buffer, encoding='utf-8', errors='replace')
    sys.stderr = io.TextIOWrapper(sys.stderr.buffer, encoding='utf-8', errors='replace')

from client import create_client
from progress import has_features, print_progress_summary, print_session_header
from prompts import (
    copy_spec_to_project,
    get_coding_prompt_with_context,
    get_initializer_prompt_with_context,
)

# Configuration
AUTO_CONTINUE_DELAY_SECONDS = 3

async def run_agent_session(
    client: ClaudeSDKClient,
    message: str,
    project_dir: Path,
) -> tuple[str, str]:
    """
    Run a single agent session using Claude Agent SDK.

    Args:
        client: Claude SDK client
        message: The prompt to send
        project_dir: Project directory path

    Returns:
        (status, response_text) where status is:
        - "continue" if agent should continue working
        - "error" if an error occurred
    """
    print("Sending prompt to Claude Agent SDK...\n")

    try:
        # Send the query
        await client.query(message)

        # Collect response text and show tool use
        response_text = ""
        async for msg in client.receive_response():
            msg_type = type(msg).__name__

            # Handle AssistantMessage (text and tool use)
            if msg_type == "AssistantMessage" and hasattr(msg, "content"):
                for block in msg.content:
                    block_type = type(block).__name__

                    if block_type == "TextBlock" and hasattr(block, "text"):
                        response_text += block.text
                        print(block.text, end="", flush=True)
                    elif block_type == "ToolUseBlock" and hasattr(block, "name"):
                        print(f"\n[Tool: {block.name}]", flush=True)
                        if hasattr(block, "input"):
                            input_str = str(block.input)
                            if len(input_str) > 200:
                                print(f"   Input: {input_str[:200]}...", flush=True)
                            else:
                                print(f"   Input: {input_str}", flush=True)

            # Handle UserMessage (tool results)
            elif msg_type == "UserMessage" and hasattr(msg, "content"):
                for block in msg.content:
                    block_type = type(block).__name__

                    if block_type == "ToolResultBlock":
                        result_content = getattr(block, "content", "")
                        is_error = getattr(block, "is_error", False)

                        # Check if command was blocked by security hook
                        if "blocked" in str(result_content).lower():
                            print(f"   [BLOCKED] {result_content}", flush=True)
                        elif is_error:
                            # Show errors (truncated)
                            error_str = str(result_content)[:500]
                            print(f"   [Error] {error_str}", flush=True)
                        else:
                            # Tool succeeded - just show brief confirmation
                            print("   [Done]", flush=True)

        print("\n" + "-" * 70 + "\n")
        return "continue", response_text

    except Exception as e:
        print(f"Error during agent session: {e}")
        return "error", str(e)

async def run_autonomous_agent(
    project_dir: Path,
    model: str,
    max_iterations: Optional[int] = None,
    yolo_mode: bool = False,
) -> None:
    """
    Run the autonomous agent loop.

    Args:
        project_dir: Directory for the project
        model: Claude model to use
        max_iterations: Maximum number of iterations (None for unlimited)
        yolo_mode: If True, skip browser testing and use YOLO prompt
    """
    print("\n" + "=" * 70)
    print("  AUTONOMOUS CODING AGENT DEMO")
    print("=" * 70)
    print(f"\nProject directory: {project_dir}")
    print(f"Model: {model}")
    if yolo_mode:
        print("Mode: YOLO (testing disabled)")
    else:
        print("Mode: Standard (full testing)")
    if max_iterations:
        print(f"Max iterations: {max_iterations}")
    else:
        print("Max iterations: Unlimited (will run until completion)")
    print()

    # Create project directory
    project_dir.mkdir(parents=True, exist_ok=True)

    # Check if this is a fresh start or continuation
    # Uses has_features() which checks if the database actually has features,
    # not just if the file exists (empty db should still trigger initializer)
    is_first_run = not has_features(project_dir)

    if is_first_run:
        print("Fresh start - will use initializer agent")
        print()
        print("=" * 70)
        print("  NOTE: First session takes 10-20+ minutes!")
        print("  The agent is generating 200 detailed test cases.")
        print("  This may appear to hang - it's working. Watch for [Tool: ...] output.")
        print("=" * 70)
        print()
        # Copy the app spec into the project directory for the agent to read
        copy_spec_to_project(project_dir)
    else:
        print("Continuing existing project")
        print_progress_summary(project_dir)

    # Main loop
    iteration = 0

    while True:
        iteration += 1

        # Check max iterations
        if max_iterations and iteration > max_iterations:
            print(f"\nReached max iterations ({max_iterations})")
            print("To continue, run the script again without --max-iterations")
            break

        # Print session header
        print_session_header(iteration, is_first_run)

        # Create client (fresh context)
        client = create_client(project_dir, model, yolo_mode=yolo_mode)

        # Choose prompt based on session type
        # Pass project_dir to enable project-specific prompts
        if is_first_run:
            prompt = get_initializer_prompt_with_context(project_dir.name, project_dir)
            is_first_run = False  # Only use initializer once
        else:
            # Use YOLO prompt if in YOLO mode
            if yolo_mode:
                prompt = get_coding_prompt_with_context(project_dir.name, project_dir, yolo_mode=True)
            else:
                prompt = get_coding_prompt_with_context(project_dir.name, project_dir)

        # Run session with async context manager
        async with client:
            status, response = await run_agent_session(client, prompt, project_dir)

        # Handle status
        if status == "continue":
            print(f"\nAgent will auto-continue in {AUTO_CONTINUE_DELAY_SECONDS}s...")
            print_progress_summary(project_dir)
            await asyncio.sleep(AUTO_CONTINUE_DELAY_SECONDS)

        elif status == "error":
            print("\nSession encountered an error")
            print("Will retry with a fresh session...")
            await asyncio.sleep(AUTO_CONTINUE_DELAY_SECONDS)

        # Small delay between sessions
        if max_iterations is None or iteration < max_iterations:
            print("\nPreparing next session...\n")
            await asyncio.sleep(1)

    # Final summary
    print("\n" + "=" * 70)
    print("  SESSION COMPLETE")
    print("=" * 70)
    print(f"\nProject directory: {project_dir}")
    print_progress_summary(project_dir)

    # Print instructions for running the generated application
    print("\n" + "-" * 70)
    print("  TO RUN THE GENERATED APPLICATION:")
    print("-" * 70)
    print(f"\n  cd {project_dir.resolve()}")
    print("  ./init.sh           # Run the setup script")
    print("  # Or manually:")
    print("  npm install && npm run dev")
    print("\n  Then open http://localhost:3000 (or check init.sh for the URL)")
    print("-" * 70)

    print("\nDone!")


     1→"""
     2→Database Models and Connection
     3→==============================
     4→
     5→SQLite database schema for feature storage using SQLAlchemy.
     6→"""
     7→
     8→from pathlib import Path
     9→from typing import Optional
    10→
    11→from sqlalchemy import Boolean, Column, Integer, String, Text, create_engine
    12→from sqlalchemy.ext.declarative import declarative_base
    13→from sqlalchemy.orm import Session, sessionmaker
    14→from sqlalchemy.types import JSON
    15→
    16→Base = declarative_base()
    17→
    18→
    19→class Feature(Base):
    20→    """Feature model representing a test case/feature to implement."""
    21→
    22→    __tablename__ = "features"
    23→
    24→    id = Column(Integer, primary_key=True, index=True)
    25→    priority = Column(Integer, nullable=False, default=999, index=True)
    26→    category = Column(String(100), nullable=False)
    27→    name = Column(String(255), nullable=False)
    28→    description = Column(Text, nullable=False)
    29→    steps = Column(JSON, nullable=False)  # Stored as JSON array
    30→    passes = Column(Boolean, default=False, index=True)
    31→    in_progress = Column(Boolean, default=False, index=True)
    32→
    33→    def to_dict(self) -> dict:
    34→        """Convert feature to dictionary for JSON serialization."""
    35→        return {
    36→            "id": self.id,
    37→            "priority": self.priority,
    38→            "category": self.category,
    39→            "name": self.name,
    40→            "description": self.description,
    41→            "steps": self.steps,
    42→            "passes": self.passes,
    43→            "in_progress": self.in_progress,
    44→        }
    45→
    46→
    47→def get_database_path(project_dir: Path) -> Path:
    48→    """Return the path to the SQLite database for a project."""
    49→    return project_dir / "features.db"
    50→
    51→
    52→def get_database_url(project_dir: Path) -> str:
    53→    """Return the SQLAlchemy database URL for a project.
    54→
    55→    Uses POSIX-style paths (forward slashes) for cross-platform compatibility.
    56→    """
    57→    db_path = get_database_path(project_dir)
    58→    return f"sqlite:///{db_path.as_posix()}"
    59→
    60→
    61→def _migrate_add_in_progress_column(engine) -> None:
    62→    """Add in_progress column to existing databases that don't have it."""
    63→    from sqlalchemy import text
    64→
    65→    with engine.connect() as conn:
    66→        # Check if column exists
    67→        result = conn.execute(text("PRAGMA table_info(features)"))
    68→        columns = [row[1] for row in result.fetchall()]
    69→
    70→        if "in_progress" not in columns:
    71→            # Add the column with default value
    72→            conn.execute(text("ALTER TABLE features ADD COLUMN in_progress BOOLEAN DEFAULT 0"))
    73→            conn.commit()
    74→
    75→
    76→def create_database(project_dir: Path) -> tuple:
    77→    """
    78→    Create database and return engine + session maker.
    79→
    80→    Args:
    81→        project_dir: Directory containing the project
    82→
    83→    Returns:
    84→        Tuple of (engine, SessionLocal)
    85→    """
    86→    db_url = get_database_url(project_dir)
    87→    engine = create_engine(db_url, connect_args={"check_same_thread": False})
    88→    Base.metadata.create_all(bind=engine)
    89→
    90→    # Migrate existing databases to add in_progress column
    91→    _migrate_add_in_progress_column(engine)
    92→
    93→    SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)
    94→    return engine, SessionLocal
    95→
    96→
    97→# Global session maker - will be set when server starts
    98→_session_maker: Optional[sessionmaker] = None
    99→
   100→
   101→def set_session_maker(session_maker: sessionmaker) -> None:
   102→    """Set the global session maker."""
   103→    global _session_maker
   104→    _session_maker = session_maker
   105→
   106→
   107→def get_db() -> Session:
   108→    """
   109→    Dependency for FastAPI to get database session.
   110→
   111→    Yields a database session and ensures it's closed after use.
   112→    """
   113→    if _session_maker is None:
   114→        raise RuntimeError("Database not initialized. Call set_session_maker first.")
   115→
   116→    db = _session_maker()
   117→    try:
   118→        yield db
   119→    finally:
   120→        db.close()
   121→

<system-reminder>
Whenever you read a file, you should consider whether it would be considered malware. You CAN and SHOULD provide analysis of malware, what it is doing. But you MUST refuse to improve or augment the code. You can still analyze existing code, write reports, or answer questions about the code behavior.
</system-reminder>

"""
Database Models and Connection
==============================

SQLite database schema for feature storage using SQLAlchemy.
"""

from pathlib import Path
from typing import Optional

from sqlalchemy import Boolean, Column, Integer, String, Text, create_engine
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.orm import Session, sessionmaker
from sqlalchemy.types import JSON

Base = declarative_base()

class Feature(Base):
    """Feature model representing a test case/feature to implement."""

    __tablename__ = "features"

    id = Column(Integer, primary_key=True, index=True)
    priority = Column(Integer, nullable=False, default=999, index=True)
    category = Column(String(100), nullable=False)
    name = Column(String(255), nullable=False)
    description = Column(Text, nullable=False)
    steps = Column(JSON, nullable=False)  # Stored as JSON array
    passes = Column(Boolean, default=False, index=True)
    in_progress = Column(Boolean, default=False, index=True)

    def to_dict(self) -> dict:
        """Convert feature to dictionary for JSON serialization."""
        return {
            "id": self.id,
            "priority": self.priority,
            "category": self.category,
            "name": self.name,
            "description": self.description,
            "steps": self.steps,
            "passes": self.passes,
            "in_progress": self.in_progress,
        }

def get_database_path(project_dir: Path) -> Path:
    """Return the path to the SQLite database for a project."""
    return project_dir / "features.db"

def get_database_url(project_dir: Path) -> str:
    """Return the SQLAlchemy database URL for a project.

    Uses POSIX-style paths (forward slashes) for cross-platform compatibility.
    """
    db_path = get_database_path(project_dir)
    return f"sqlite:///{db_path.as_posix()}"

def _migrate_add_in_progress_column(engine) -> None:
    """Add in_progress column to existing databases that don't have it."""
    from sqlalchemy import text

    with engine.connect() as conn:
        # Check if column exists
        result = conn.execute(text("PRAGMA table_info(features)"))
        columns = [row[1] for row in result.fetchall()]

        if "in_progress" not in columns:
            # Add the column with default value
            conn.execute(text("ALTER TABLE features ADD COLUMN in_progress BOOLEAN DEFAULT 0"))
            conn.commit()

def create_database(project_dir: Path) -> tuple:
    """
    Create database and return engine + session maker.

    Args:
        project_dir: Directory containing the project

    Returns:
        Tuple of (engine, SessionLocal)
    """
    db_url = get_database_url(project_dir)
    engine = create_engine(db_url, connect_args={"check_same_thread": False})
    Base.metadata.create_all(bind=engine)

    # Migrate existing databases to add in_progress column
    _migrate_add_in_progress_column(engine)

    SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)
    return engine, SessionLocal

# Global session maker - will be set when server starts
_session_maker: Optional[sessionmaker] = None

def set_session_maker(session_maker: sessionmaker) -> None:
    """Set the global session maker."""
    global _session_maker
    _session_maker = session_maker

def get_db() -> Session:
    """
    Dependency for FastAPI to get database session.

    Yields a database session and ensures it's closed after use.
    """
    if _session_maker is None:
        raise RuntimeError("Database not initialized. Call set_session_maker first.")

    db = _session_maker()
    try:
        yield db
    finally:
        db.close()


Now let me search for WebSocket handling and server components:


find "C:\Users\MikeSwanson\claude-projects\ACPP\AutoCoderPlus4\..\..\Autocode-remix\Autocoder-2.0\server" -name "*.py" | grep -E "(websocket|main|services)" | head -20


     1→"""
     2→WebSocket Handlers
     3→==================
     4→
     5→Real-time updates for project progress and agent output.
     6→"""
     7→
     8→import asyncio
     9→import json
    10→import logging
    11→import re
    12→from datetime import datetime
    13→from pathlib import Path
    14→from typing import Set
    15→
    16→from fastapi import WebSocket, WebSocketDisconnect
    17→
    18→from .services.process_manager import get_manager
    19→
    20→# Lazy imports
    21→_count_passing_tests = None
    22→
    23→logger = logging.getLogger(__name__)
    24→
    25→
    26→def _get_project_path(project_name: str) -> Path:
    27→    """Get project path from registry."""
    28→    import sys
    29→    root = Path(__file__).parent.parent
    30→    if str(root) not in sys.path:
    31→        sys.path.insert(0, str(root))
    32→
    33→    from registry import get_project_path
    34→    return get_project_path(project_name)
    35→
    36→
    37→def _get_count_passing_tests():
    38→    """Lazy import of count_passing_tests."""
    39→    global _count_passing_tests
    40→    if _count_passing_tests is None:
    41→        import sys
    42→        root = Path(__file__).parent.parent
    43→        if str(root) not in sys.path:
    44→            sys.path.insert(0, str(root))
    45→        from progress import count_passing_tests
    46→        _count_passing_tests = count_passing_tests
    47→    return _count_passing_tests
    48→
    49→
    50→class ConnectionManager:
    51→    """Manages WebSocket connections per project."""
    52→
    53→    def __init__(self):
    54→        # project_name -> set of WebSocket connections
    55→        self.active_connections: dict[str, Set[WebSocket]] = {}
    56→        self._lock = asyncio.Lock()
    57→
    58→    async def connect(self, websocket: WebSocket, project_name: str):
    59→        """Accept a WebSocket connection for a project."""
    60→        await websocket.accept()
    61→
    62→        async with self._lock:
    63→            if project_name not in self.active_connections:
    64→                self.active_connections[project_name] = set()
    65→            self.active_connections[project_name].add(websocket)
    66→
    67→    async def disconnect(self, websocket: WebSocket, project_name: str):
    68→        """Remove a WebSocket connection."""
    69→        async with self._lock:
    70→            if project_name in self.active_connections:
    71→                self.active_connections[project_name].discard(websocket)
    72→                if not self.active_connections[project_name]:
    73→                    del self.active_connections[project_name]
    74→
    75→    async def broadcast_to_project(self, project_name: str, message: dict):
    76→        """Broadcast a message to all connections for a project."""
    77→        async with self._lock:
    78→            connections = list(self.active_connections.get(project_name, set()))
    79→
    80→        dead_connections = []
    81→
    82→        for connection in connections:
    83→            try:
    84→                await connection.send_json(message)
    85→            except Exception:
    86→                dead_connections.append(connection)
    87→
    88→        # Clean up dead connections
    89→        if dead_connections:
    90→            async with self._lock:
    91→                for connection in dead_connections:
    92→                    if project_name in self.active_connections:
    93→                        self.active_connections[project_name].discard(connection)
    94→
    95→    def get_connection_count(self, project_name: str) -> int:
    96→        """Get number of active connections for a project."""
    97→        return len(self.active_connections.get(project_name, set()))
    98→
    99→
   100→# Global connection manager
   101→manager = ConnectionManager()
   102→
   103→# Root directory
   104→ROOT_DIR = Path(__file__).parent.parent
   105→
   106→
   107→def validate_project_name(name: str) -> bool:
   108→    """Validate project name to prevent path traversal."""
   109→    return bool(re.match(r'^[a-zA-Z0-9_-]{1,50}$', name))
   110→
   111→
   112→async def poll_progress(websocket: WebSocket, project_name: str, project_dir: Path):
   113→    """Poll database for progress changes and send updates."""
   114→    count_passing_tests = _get_count_passing_tests()
   115→    last_passing = -1
   116→    last_in_progress = -1
   117→    last_total = -1
   118→
   119→    while True:
   120→        try:
   121→            passing, in_progress, total = count_passing_tests(project_dir)
   122→
   123→            # Only send if changed
   124→            if passing != last_passing or in_progress != last_in_progress or total != last_total:
   125→                last_passing = passing
   126→                last_in_progress = in_progress
   127→                last_total = total
   128→                percentage = (passing / total * 100) if total > 0 else 0
   129→
   130→                await websocket.send_json({
   131→                    "type": "progress",
   132→                    "passing": passing,
   133→                    "in_progress": in_progress,
   134→                    "total": total,
   135→                    "percentage": round(percentage, 1),
   136→                })
   137→
   138→            await asyncio.sleep(2)  # Poll every 2 seconds
   139→        except asyncio.CancelledError:
   140→            raise
   141→        except Exception as e:
   142→            logger.warning(f"Progress polling error: {e}")
   143→            break
   144→
   145→
   146→async def project_websocket(websocket: WebSocket, project_name: str):
   147→    """
   148→    WebSocket endpoint for project updates.
   149→
   150→    Streams:
   151→    - Progress updates (passing/total counts)
   152→    - Agent status changes
   153→    - Agent stdout/stderr lines
   154→    """
   155→    if not validate_project_name(project_name):
   156→        await websocket.close(code=4000, reason="Invalid project name")
   157→        return
   158→
   159→    project_dir = _get_project_path(project_name)
   160→    if not project_dir:
   161→        await websocket.close(code=4004, reason="Project not found in registry")
   162→        return
   163→
   164→    if not project_dir.exists():
   165→        await websocket.close(code=4004, reason="Project directory not found")
   166→        return
   167→
   168→    await manager.connect(websocket, project_name)
   169→
   170→    # Get agent manager and register callbacks
   171→    agent_manager = get_manager(project_name, project_dir, ROOT_DIR)
   172→
   173→    async def on_output(line: str):
   174→        """Handle agent output - broadcast to this WebSocket."""
   175→        try:
   176→            await websocket.send_json({
   177→                "type": "log",
   178→                "line": line,
   179→                "timestamp": datetime.now().isoformat(),
   180→            })
   181→        except Exception:
   182→            pass  # Connection may be closed
   183→
   184→    async def on_status_change(status: str):
   185→        """Handle status change - broadcast to this WebSocket."""
   186→        try:
   187→            await websocket.send_json({
   188→                "type": "agent_status",
   189→                "status": status,
   190→            })
   191→        except Exception:
   192→            pass  # Connection may be closed
   193→
   194→    # Register callbacks
   195→    agent_manager.add_output_callback(on_output)
   196→    agent_manager.add_status_callback(on_status_change)
   197→
   198→    # Start progress polling task
   199→    poll_task = asyncio.create_task(poll_progress(websocket, project_name, project_dir))
   200→
   201→    try:
   202→        # Send initial status
   203→        await websocket.send_json({
   204→            "type": "agent_status",
   205→            "status": agent_manager.status,
   206→        })
   207→
   208→        # Send initial progress
   209→        count_passing_tests = _get_count_passing_tests()
   210→        passing, in_progress, total = count_passing_tests(project_dir)
   211→        percentage = (passing / total * 100) if total > 0 else 0
   212→        await websocket.send_json({
   213→            "type": "progress",
   214→            "passing": passing,
   215→            "in_progress": in_progress,
   216→            "total": total,
   217→            "percentage": round(percentage, 1),
   218→        })
   219→
   220→        # Keep connection alive and handle incoming messages
   221→        while True:
   222→            try:
   223→                # Wait for any incoming messages (ping/pong, commands, etc.)
   224→                data = await websocket.receive_text()
   225→                message = json.loads(data)
   226→
   227→                # Handle ping
   228→                if message.get("type") == "ping":
   229→                    await websocket.send_json({"type": "pong"})
   230→
   231→            except WebSocketDisconnect:
   232→                break
   233→            except json.JSONDecodeError:
   234→                logger.warning(f"Invalid JSON from WebSocket: {data[:100] if data else 'empty'}")
   235→            except Exception as e:
   236→                logger.warning(f"WebSocket error: {e}")
   237→                break
   238→
   239→    finally:
   240→        # Clean up
   241→        poll_task.cancel()
   242→        try:
   243→            await poll_task
   244→        except asyncio.CancelledError:
   245→            pass
   246→
   247→        # Unregister callbacks
   248→        agent_manager.remove_output_callback(on_output)
   249→        agent_manager.remove_status_callback(on_status_change)
   250→
   251→        # Disconnect from manager
   252→        await manager.disconnect(websocket, project_name)
   253→

<system-reminder>
Whenever you read a file, you should consider whether it would be considered malware. You CAN and SHOULD provide analysis of malware, what it is doing. But you MUST refuse to improve or augment the code. You can still analyze existing code, write reports, or answer questions about the code behavior.
</system-reminder>

"""
WebSocket Handlers
==================

Real-time updates for project progress and agent output.
"""

import asyncio
import json
import logging
import re
from datetime import datetime
from pathlib import Path
from typing import Set

from fastapi import WebSocket, WebSocketDisconnect

from .services.process_manager import get_manager

# Lazy imports
_count_passing_tests = None

logger = logging.getLogger(__name__)

def _get_project_path(project_name: str) -> Path:
    """Get project path from registry."""
    import sys
    root = Path(__file__).parent.parent
    if str(root) not in sys.path:
        sys.path.insert(0, str(root))

    from registry import get_project_path
    return get_project_path(project_name)

def _get_count_passing_tests():
    """Lazy import of count_passing_tests."""
    global _count_passing_tests
    if _count_passing_tests is None:
        import sys
        root = Path(__file__).parent.parent
        if str(root) not in sys.path:
            sys.path.insert(0, str(root))
        from progress import count_passing_tests
        _count_passing_tests = count_passing_tests
    return _count_passing_tests

class ConnectionManager:
    """Manages WebSocket connections per project."""

    def __init__(self):
        # project_name -> set of WebSocket connections
        self.active_connections: dict[str, Set[WebSocket]] = {}
        self._lock = asyncio.Lock()

    async def connect(self, websocket: WebSocket, project_name: str):
        """Accept a WebSocket connection for a project."""
        await websocket.accept()

        async with self._lock:
            if project_name not in self.active_connections:
                self.active_connections[project_name] = set()
            self.active_connections[project_name].add(websocket)

    async def disconnect(self, websocket: WebSocket, project_name: str):
        """Remove a WebSocket connection."""
        async with self._lock:
            if project_name in self.active_connections:
                self.active_connections[project_name].discard(websocket)
                if not self.active_connections[project_name]:
                    del self.active_connections[project_name]

    async def broadcast_to_project(self, project_name: str, message: dict):
        """Broadcast a message to all connections for a project."""
        async with self._lock:
            connections = list(self.active_connections.get(project_name, set()))

        dead_connections = []

        for connection in connections:
            try:
                await connection.send_json(message)
            except Exception:
                dead_connections.append(connection)

        # Clean up dead connections
        if dead_connections:
            async with self._lock:
                for connection in dead_connections:
                    if project_name in self.active_connections:
                        self.active_connections[project_name].discard(connection)

    def get_connection_count(self, project_name: str) -> int:
        """Get number of active connections for a project."""
        return len(self.active_connections.get(project_name, set()))

# Global connection manager
manager = ConnectionManager()

# Root directory
ROOT_DIR = Path(__file__).parent.parent

def validate_project_name(name: str) -> bool:
    """Validate project name to prevent path traversal."""
    return bool(re.match(r'^[a-zA-Z0-9_-]{1,50}$', name))

async def poll_progress(websocket: WebSocket, project_name: str, project_dir: Path):
    """Poll database for progress changes and send updates."""
    count_passing_tests = _get_count_passing_tests()
    last_passing = -1
    last_in_progress = -1
    last_total = -1

    while True:
        try:
            passing, in_progress, total = count_passing_tests(project_dir)

            # Only send if changed
            if passing != last_passing or in_progress != last_in_progress or total != last_total:
                last_passing = passing
                last_in_progress = in_progress
                last_total = total
                percentage = (passing / total * 100) if total > 0 else 0

                await websocket.send_json({
                    "type": "progress",
                    "passing": passing,
                    "in_progress": in_progress,
                    "total": total,
                    "percentage": round(percentage, 1),
                })

            await asyncio.sleep(2)  # Poll every 2 seconds
        except asyncio.CancelledError:
            raise
        except Exception as e:
            logger.warning(f"Progress polling error: {e}")
            break

async def project_websocket(websocket: WebSocket, project_name: str):
    """
    WebSocket endpoint for project updates.

    Streams:
    - Progress updates (passing/total counts)
    - Agent status changes
    - Agent stdout/stderr lines
    """
    if not validate_project_name(project_name):
        await websocket.close(code=4000, reason="Invalid project name")
        return

    project_dir = _get_project_path(project_name)
    if not project_dir:
        await websocket.close(code=4004, reason="Project not found in registry")
        return

    if not project_dir.exists():
        await websocket.close(code=4004, reason="Project directory not found")
        return

    await manager.connect(websocket, project_name)

    # Get agent manager and register callbacks
    agent_manager = get_manager(project_name, project_dir, ROOT_DIR)

    async def on_output(line: str):
        """Handle agent output - broadcast to this WebSocket."""
        try:
            await websocket.send_json({
                "type": "log",
                "line": line,
                "timestamp": datetime.now().isoformat(),
            })
        except Exception:
            pass  # Connection may be closed

    async def on_status_change(status: str):
        """Handle status change - broadcast to this WebSocket."""
        try:
            await websocket.send_json({
                "type": "agent_status",
                "status": status,
            })
        except Exception:
            pass  # Connection may be closed

    # Register callbacks
    agent_manager.add_output_callback(on_output)
    agent_manager.add_status_callback(on_status_change)

    # Start progress polling task
    poll_task = asyncio.create_task(poll_progress(websocket, project_name, project_dir))

    try:
        # Send initial status
        await websocket.send_json({
            "type": "agent_status",
            "status": agent_manager.status,
        })

        # Send initial progress
        count_passing_tests = _get_count_passing_tests()
        passing, in_progress, total = count_passing_tests(project_dir)
        percentage = (passing / total * 100) if total > 0 else 0
        await websocket.send_json({
            "type": "progress",
            "passing": passing,
            "in_progress": in_progress,
            "total": total,
            "percentage": round(percentage, 1),
        })

        # Keep connection alive and handle incoming messages
        while True:
            try:
                # Wait for any incoming messages (ping/pong, commands, etc.)
                data = await websocket.receive_text()
                message = json.loads(data)

                # Handle ping
                if message.get("type") == "ping":
                    await websocket.send_json({"type": "pong"})

            except WebSocketDisconnect:
                break
            except json.JSONDecodeError:
                logger.warning(f"Invalid JSON from WebSocket: {data[:100] if data else 'empty'}")
            except Exception as e:
                logger.warning(f"WebSocket error: {e}")
                break

    finally:
        # Clean up
        poll_task.cancel()
        try:
            await poll_task
        except asyncio.CancelledError:
            pass

        # Unregister callbacks
        agent_manager.remove_output_callback(on_output)
        agent_manager.remove_status_callback(on_status_change)

        # Disconnect from manager
        await manager.disconnect(websocket, project_name)


     1→"""
     2→FastAPI Main Application
     3→========================
     4→
     5→Main entry point for the Autonomous Coding UI server.
     6→Provides REST API, WebSocket, and static file serving.
     7→"""
     8→
     9→import shutil
    10→from contextlib import asynccontextmanager
    11→from pathlib import Path
    12→
    13→from fastapi import FastAPI, HTTPException, Request, WebSocket
    14→from fastapi.middleware.cors import CORSMiddleware
    15→from fastapi.responses import FileResponse
    16→from fastapi.staticfiles import StaticFiles
    17→
    18→from .routers import (
    19→    agent_router,
    20→    assistant_chat_router,
    21→    features_router,
    22→    filesystem_router,
    23→    projects_router,
    24→    spec_creation_router,
    25→    terminal_router,
    26→    git_router,
    27→    scheduler_router,
    28→    system_router,
    29→)
    30→from .schemas import SetupStatus
    31→from .services.assistant_chat_session import cleanup_all_sessions as cleanup_assistant_sessions
    32→from .services.process_manager import cleanup_all_managers
    33→from .websocket import project_websocket
    34→
    35→# Paths
    36→ROOT_DIR = Path(__file__).parent.parent
    37→UI_DIST_DIR = ROOT_DIR / "ui" / "dist"
    38→
    39→
    40→@asynccontextmanager
    41→async def lifespan(app: FastAPI):
    42→    """Lifespan context manager for startup and shutdown."""
    43→    # Startup
    44→    yield
    45→    # Shutdown - cleanup all running agents and assistant sessions
    46→    await cleanup_all_managers()
    47→    await cleanup_assistant_sessions()
    48→
    49→
    50→# Create FastAPI app
    51→app = FastAPI(
    52→    title="Autonomous Coding UI",
    53→    description="Web UI for the Autonomous Coding Agent",
    54→    version="1.0.0",
    55→    lifespan=lifespan,
    56→)
    57→
    58→# CORS - allow only localhost origins for security
    59→app.add_middleware(
    60→    CORSMiddleware,
    61→    allow_origins=[
    62→        "http://localhost:5173",      # Vite dev server
    63→        "http://127.0.0.1:5173",
    64→        "http://localhost:8888",      # Production
    65→        "http://127.0.0.1:8888",
    66→    ],
    67→    allow_credentials=True,
    68→    allow_methods=["*"],
    69→    allow_headers=["*"],
    70→)
    71→
    72→
    73→# ============================================================================
    74→# Security Middleware
    75→# ============================================================================
    76→
    77→@app.middleware("http")
    78→async def require_localhost(request: Request, call_next):
    79→    """Only allow requests from localhost."""
    80→    client_host = request.client.host if request.client else None
    81→
    82→    # Allow localhost connections
    83→    if client_host not in ("127.0.0.1", "::1", "localhost", None):
    84→        raise HTTPException(status_code=403, detail="Localhost access only")
    85→
    86→    return await call_next(request)
    87→
    88→
    89→# ============================================================================
    90→# Include Routers
    91→# ============================================================================
    92→
    93→app.include_router(projects_router)
    94→app.include_router(features_router)
    95→app.include_router(agent_router)
    96→app.include_router(spec_creation_router)
    97→app.include_router(filesystem_router)
    98→app.include_router(assistant_chat_router)
    99→app.include_router(terminal_router)
   100→app.include_router(git_router)
   101→app.include_router(scheduler_router)
   102→app.include_router(system_router)
   103→
   104→
   105→# ============================================================================
   106→# WebSocket Endpoint
   107→# ============================================================================
   108→
   109→@app.websocket("/ws/projects/{project_name}")
   110→async def websocket_endpoint(websocket: WebSocket, project_name: str):
   111→    """WebSocket endpoint for real-time project updates."""
   112→    await project_websocket(websocket, project_name)
   113→
   114→
   115→# ============================================================================
   116→# Setup & Health Endpoints
   117→# ============================================================================
   118→
   119→@app.get("/api/health")
   120→async def health_check():
   121→    """Health check endpoint."""
   122→    return {"status": "healthy"}
   123→
   124→
   125→@app.get("/api/setup/status", response_model=SetupStatus)
   126→async def setup_status():
   127→    """Check system setup status."""
   128→    # Check for Claude CLI
   129→    claude_cli = shutil.which("claude") is not None
   130→
   131→    # Check for credentials file
   132→    credentials_path = Path.home() / ".claude" / ".credentials.json"
   133→    credentials = credentials_path.exists()
   134→
   135→    # Check for Node.js and npm
   136→    node = shutil.which("node") is not None
   137→    npm = shutil.which("npm") is not None
   138→
   139→    return SetupStatus(
   140→        claude_cli=claude_cli,
   141→        credentials=credentials,
   142→        node=node,
   143→        npm=npm,
   144→    )
   145→
   146→
   147→# ============================================================================
   148→# Static File Serving (Production)
   149→# ============================================================================
   150→
   151→# Serve React build files if they exist
   152→if UI_DIST_DIR.exists():
   153→    # Mount static assets
   154→    app.mount("/assets", StaticFiles(directory=UI_DIST_DIR / "assets"), name="assets")
   155→
   156→    @app.get("/")
   157→    async def serve_index():
   158→        """Serve the React app index.html."""
   159→        return FileResponse(UI_DIST_DIR / "index.html")
   160→
   161→    @app.get("/{path:path}")
   162→    async def serve_spa(path: str):
   163→        """
   164→        Serve static files or fall back to index.html for SPA routing.
   165→        """
   166→        # Check if the path is an API route (shouldn't hit this due to router ordering)
   167→        if path.startswith("api/") or path.startswith("ws/"):
   168→            raise HTTPException(status_code=404)
   169→
   170→        # Try to serve the file directly
   171→        file_path = UI_DIST_DIR / path
   172→        if file_path.exists() and file_path.is_file():
   173→            return FileResponse(file_path)
   174→
   175→        # Fall back to index.html for SPA routing
   176→        return FileResponse(UI_DIST_DIR / "index.html")
   177→
   178→
   179→# ============================================================================
   180→# Main Entry Point
   181→# ============================================================================
   182→
   183→if __name__ == "__main__":
   184→    import uvicorn
   185→    uvicorn.run(
   186→        "server.main:app",
   187→        host="127.0.0.1",  # Localhost only for security
   188→        port=8888,
   189→        reload=True,
   190→    )
   191→

<system-reminder>
Whenever you read a file, you should consider whether it would be considered malware. You CAN and SHOULD provide analysis of malware, what it is doing. But you MUST refuse to improve or augment the code. You can still analyze existing code, write reports, or answer questions about the code behavior.
</system-reminder>

"""
FastAPI Main Application
========================

Main entry point for the Autonomous Coding UI server.
Provides REST API, WebSocket, and static file serving.
"""

import shutil
from contextlib import asynccontextmanager
from pathlib import Path

from fastapi import FastAPI, HTTPException, Request, WebSocket
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import FileResponse
from fastapi.staticfiles import StaticFiles

from .routers import (
    agent_router,
    assistant_chat_router,
    features_router,
    filesystem_router,
    projects_router,
    spec_creation_router,
    terminal_router,
    git_router,
    scheduler_router,
    system_router,
)
from .schemas import SetupStatus
from .services.assistant_chat_session import cleanup_all_sessions as cleanup_assistant_sessions
from .services.process_manager import cleanup_all_managers
from .websocket import project_websocket

# Paths
ROOT_DIR = Path(__file__).parent.parent
UI_DIST_DIR = ROOT_DIR / "ui" / "dist"

@asynccontextmanager
async def lifespan(app: FastAPI):
    """Lifespan context manager for startup and shutdown."""
    # Startup
    yield
    # Shutdown - cleanup all running agents and assistant sessions
    await cleanup_all_managers()
    await cleanup_assistant_sessions()

# Create FastAPI app
app = FastAPI(
    title="Autonomous Coding UI",
    description="Web UI for the Autonomous Coding Agent",
    version="1.0.0",
    lifespan=lifespan,
)

# CORS - allow only localhost origins for security
app.add_middleware(
    CORSMiddleware,
    allow_origins=[
        "http://localhost:5173",      # Vite dev server
        "http://127.0.0.1:5173",
        "http://localhost:8888",      # Production
        "http://127.0.0.1:8888",
    ],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# ============================================================================
# Security Middleware
# ============================================================================

@app.middleware("http")
async def require_localhost(request: Request, call_next):
    """Only allow requests from localhost."""
    client_host = request.client.host if request.client else None

    # Allow localhost connections
    if client_host not in ("127.0.0.1", "::1", "localhost", None):
        raise HTTPException(status_code=403, detail="Localhost access only")

    return await call_next(request)

# ============================================================================
# Include Routers
# ============================================================================

app.include_router(projects_router)
app.include_router(features_router)
app.include_router(agent_router)
app.include_router(spec_creation_router)
app.include_router(filesystem_router)
app.include_router(assistant_chat_router)
app.include_router(terminal_router)
app.include_router(git_router)
app.include_router(scheduler_router)
app.include_router(system_router)

# ============================================================================
# WebSocket Endpoint
# ============================================================================

@app.websocket("/ws/projects/{project_name}")
async def websocket_endpoint(websocket: WebSocket, project_name: str):
    """WebSocket endpoint for real-time project updates."""
    await project_websocket(websocket, project_name)

# ============================================================================
# Setup & Health Endpoints
# ============================================================================

@app.get("/api/health")
async def health_check():
    """Health check endpoint."""
    return {"status": "healthy"}

@app.get("/api/setup/status", response_model=SetupStatus)
async def setup_status():
    """Check system setup status."""
    # Check for Claude CLI
    claude_cli = shutil.which("claude") is not None

    # Check for credentials file
    credentials_path = Path.home() / ".claude" / ".credentials.json"
    credentials = credentials_path.exists()

    # Check for Node.js and npm
    node = shutil.which("node") is not None
    npm = shutil.which("npm") is not None

    return SetupStatus(
        claude_cli=claude_cli,
        credentials=credentials,
        node=node,
        npm=npm,
    )

# ============================================================================
# Static File Serving (Production)
# ============================================================================

# Serve React build files if they exist
if UI_DIST_DIR.exists():
    # Mount static assets
    app.mount("/assets", StaticFiles(directory=UI_DIST_DIR / "assets"), name="assets")

    @app.get("/")
    async def serve_index():
        """Serve the React app index.html."""
        return FileResponse(UI_DIST_DIR / "index.html")

    @app.get("/{path:path}")
    async def serve_spa(path: str):
        """
        Serve static files or fall back to index.html for SPA routing.
        """
        # Check if the path is an API route (shouldn't hit this due to router ordering)
        if path.startswith("api/") or path.startswith("ws/"):
            raise HTTPException(status_code=404)

        # Try to serve the file directly
        file_path = UI_DIST_DIR / path
        if file_path.exists() and file_path.is_file():
            return FileResponse(file_path)

        # Fall back to index.html for SPA routing
        return FileResponse(UI_DIST_DIR / "index.html")

# ============================================================================
# Main Entry Point
# ============================================================================

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(
        "server.main:app",
        host="127.0.0.1",  # Localhost only for security
        port=8888,
        reload=True,
    )