     1→#!/usr/bin/env python3
     2→"""
     3→MCP Server for Feature Management
     4→==================================
     5→
     6→Provides tools to manage features in the autonomous coding system,
     7→replacing the previous FastAPI-based REST API.
     8→
     9→Tools:
    10→- feature_get_stats: Get progress statistics
    11→- feature_get_next: Get next feature to implement
    12→- feature_get_for_regression: Get random passing features for testing
    13→- feature_mark_passing: Mark a feature as passing
    14→- feature_skip: Skip a feature (move to end of queue)
    15→- feature_mark_in_progress: Mark a feature as in-progress
    16→- feature_clear_in_progress: Clear in-progress status
    17→- feature_create_bulk: Create multiple features at once
    18→"""
    19→
    20→import json
    21→import os
    22→import sys
    23→from contextlib import asynccontextmanager
    24→from pathlib import Path
    25→from typing import Annotated
    26→
    27→from mcp.server.fastmcp import FastMCP
    28→from pydantic import BaseModel, Field
    29→from sqlalchemy.sql.expression import func
    30→
    31→# Add parent directory to path so we can import from api module
    32→sys.path.insert(0, str(Path(__file__).parent.parent))
    33→
    34→from api.database import Feature, create_database
    35→from api.migration import migrate_json_to_sqlite
    36→
    37→# Configuration from environment
    38→PROJECT_DIR = Path(os.environ.get("PROJECT_DIR", ".")).resolve()
    39→
    40→
    41→# Pydantic models for input validation
    42→class MarkPassingInput(BaseModel):
    43→    """Input for marking a feature as passing."""
    44→    feature_id: int = Field(..., description="The ID of the feature to mark as passing", ge=1)
    45→
    46→
    47→class SkipFeatureInput(BaseModel):
    48→    """Input for skipping a feature."""
    49→    feature_id: int = Field(..., description="The ID of the feature to skip", ge=1)
    50→
    51→
    52→class MarkInProgressInput(BaseModel):
    53→    """Input for marking a feature as in-progress."""
    54→    feature_id: int = Field(..., description="The ID of the feature to mark as in-progress", ge=1)
    55→
    56→
    57→class ClearInProgressInput(BaseModel):
    58→    """Input for clearing in-progress status."""
    59→    feature_id: int = Field(..., description="The ID of the feature to clear in-progress status", ge=1)
    60→
    61→
    62→class RegressionInput(BaseModel):
    63→    """Input for getting regression features."""
    64→    limit: int = Field(default=3, ge=1, le=10, description="Maximum number of passing features to return")
    65→
    66→
    67→class FeatureCreateItem(BaseModel):
    68→    """Schema for creating a single feature."""
    69→    category: str = Field(..., min_length=1, max_length=100, description="Feature category")
    70→    name: str = Field(..., min_length=1, max_length=255, description="Feature name")
    71→    description: str = Field(..., min_length=1, description="Detailed description")
    72→    steps: list[str] = Field(..., min_length=1, description="Implementation/test steps")
    73→
    74→
    75→class BulkCreateInput(BaseModel):
    76→    """Input for bulk creating features."""
    77→    features: list[FeatureCreateItem] = Field(..., min_length=1, description="List of features to create")
    78→
    79→
    80→# Global database session maker (initialized on startup)
    81→_session_maker = None
    82→_engine = None
    83→
    84→
    85→@asynccontextmanager
    86→async def server_lifespan(server: FastMCP):
    87→    """Initialize database on startup, cleanup on shutdown."""
    88→    global _session_maker, _engine
    89→
    90→    # Create project directory if it doesn't exist
    91→    PROJECT_DIR.mkdir(parents=True, exist_ok=True)
    92→
    93→    # Initialize database
    94→    _engine, _session_maker = create_database(PROJECT_DIR)
    95→
    96→    # Run migration if needed (converts legacy JSON to SQLite)
    97→    migrate_json_to_sqlite(PROJECT_DIR, _session_maker)
    98→
    99→    yield
   100→
   101→    # Cleanup
   102→    if _engine:
   103→        _engine.dispose()
   104→
   105→
   106→# Initialize the MCP server
   107→mcp = FastMCP("features", lifespan=server_lifespan)
   108→
   109→
   110→def get_session():
   111→    """Get a new database session."""
   112→    if _session_maker is None:
   113→        raise RuntimeError("Database not initialized")
   114→    return _session_maker()
   115→
   116→
   117→@mcp.tool()
   118→def feature_get_stats() -> str:
   119→    """Get statistics about feature completion progress.
   120→
   121→    Returns the number of passing features, in-progress features, total features,
   122→    and completion percentage. Use this to track overall progress of the implementation.
   123→
   124→    Returns:
   125→        JSON with: passing (int), in_progress (int), total (int), percentage (float)
   126→    """
   127→    session = get_session()
   128→    try:
   129→        total = session.query(Feature).count()
   130→        passing = session.query(Feature).filter(Feature.passes == True).count()
   131→        in_progress = session.query(Feature).filter(Feature.in_progress == True).count()
   132→        percentage = round((passing / total) * 100, 1) if total > 0 else 0.0
   133→
   134→        return json.dumps({
   135→            "passing": passing,
   136→            "in_progress": in_progress,
   137→            "total": total,
   138→            "percentage": percentage
   139→        }, indent=2)
   140→    finally:
   141→        session.close()
   142→
   143→
   144→@mcp.tool()
   145→def feature_get_next() -> str:
   146→    """Get the highest-priority pending feature to work on.
   147→
   148→    Returns the feature with the lowest priority number that has passes=false.
   149→    Use this at the start of each coding session to determine what to implement next.
   150→
   151→    Returns:
   152→        JSON with feature details (id, priority, category, name, description, steps, passes, in_progress)
   153→        or error message if all features are passing.
   154→    """
   155→    session = get_session()
   156→    try:
   157→        feature = (
   158→            session.query(Feature)
   159→            .filter(Feature.passes == False)
   160→            .order_by(Feature.priority.asc(), Feature.id.asc())
   161→            .first()
   162→        )
   163→
   164→        if feature is None:
   165→            return json.dumps({"error": "All features are passing! No more work to do."})
   166→
   167→        return json.dumps(feature.to_dict(), indent=2)
   168→    finally:
   169→        session.close()
   170→
   171→
   172→@mcp.tool()
   173→def feature_get_for_regression(
   174→    limit: Annotated[int, Field(default=3, ge=1, le=10, description="Maximum number of passing features to return")] = 3
   175→) -> str:
   176→    """Get random passing features for regression testing.
   177→
   178→    Returns a random selection of features that are currently passing.
   179→    Use this to verify that previously implemented features still work
   180→    after making changes.
   181→
   182→    Args:
   183→        limit: Maximum number of features to return (1-10, default 3)
   184→
   185→    Returns:
   186→        JSON with: features (list of feature objects), count (int)
   187→    """
   188→    session = get_session()
   189→    try:
   190→        features = (
   191→            session.query(Feature)
   192→            .filter(Feature.passes == True)
   193→            .order_by(func.random())
   194→            .limit(limit)
   195→            .all()
   196→        )
   197→
   198→        return json.dumps({
   199→            "features": [f.to_dict() for f in features],
   200→            "count": len(features)
   201→        }, indent=2)
   202→    finally:
   203→        session.close()
   204→
   205→
   206→@mcp.tool()
   207→def feature_mark_passing(
   208→    feature_id: Annotated[int, Field(description="The ID of the feature to mark as passing", ge=1)]
   209→) -> str:
   210→    """Mark a feature as passing after successful implementation.
   211→
   212→    Updates the feature's passes field to true and clears the in_progress flag.
   213→    Use this after you have implemented the feature and verified it works correctly.
   214→
   215→    Args:
   216→        feature_id: The ID of the feature to mark as passing
   217→
   218→    Returns:
   219→        JSON with the updated feature details, or error if not found.
   220→    """
   221→    session = get_session()
   222→    try:
   223→        feature = session.query(Feature).filter(Feature.id == feature_id).first()
   224→
   225→        if feature is None:
   226→            return json.dumps({"error": f"Feature with ID {feature_id} not found"})
   227→
   228→        feature.passes = True
   229→        feature.in_progress = False
   230→        session.commit()
   231→        session.refresh(feature)
   232→
   233→        return json.dumps(feature.to_dict(), indent=2)
   234→    finally:
   235→        session.close()
   236→
   237→
   238→@mcp.tool()
   239→def feature_skip(
   240→    feature_id: Annotated[int, Field(description="The ID of the feature to skip", ge=1)]
   241→) -> str:
   242→    """Skip a feature by moving it to the end of the priority queue.
   243→
   244→    Use this when a feature cannot be implemented yet due to:
   245→    - Dependencies on other features that aren't implemented yet
   246→    - External blockers (missing assets, unclear requirements)
   247→    - Technical prerequisites that need to be addressed first
   248→
   249→    The feature's priority is set to max_priority + 1, so it will be
   250→    worked on after all other pending features. Also clears the in_progress
   251→    flag so the feature returns to "pending" status.
   252→
   253→    Args:
   254→        feature_id: The ID of the feature to skip
   255→
   256→    Returns:
   257→        JSON with skip details: id, name, old_priority, new_priority, message
   258→    """
   259→    session = get_session()
   260→    try:
   261→        feature = session.query(Feature).filter(Feature.id == feature_id).first()
   262→
   263→        if feature is None:
   264→            return json.dumps({"error": f"Feature with ID {feature_id} not found"})
   265→
   266→        if feature.passes:
   267→            return json.dumps({"error": "Cannot skip a feature that is already passing"})
   268→
   269→        old_priority = feature.priority
   270→
   271→        # Get max priority and set this feature to max + 1
   272→        max_priority_result = session.query(Feature.priority).order_by(Feature.priority.desc()).first()
   273→        new_priority = (max_priority_result[0] + 1) if max_priority_result else 1
   274→
   275→        feature.priority = new_priority
   276→        feature.in_progress = False
   277→        session.commit()
   278→        session.refresh(feature)
   279→
   280→        return json.dumps({
   281→            "id": feature.id,
   282→            "name": feature.name,
   283→            "old_priority": old_priority,
   284→            "new_priority": new_priority,
   285→            "message": f"Feature '{feature.name}' moved to end of queue"
   286→        }, indent=2)
   287→    finally:
   288→        session.close()
   289→
   290→
   291→@mcp.tool()
   292→def feature_mark_in_progress(
   293→    feature_id: Annotated[int, Field(description="The ID of the feature to mark as in-progress", ge=1)]
   294→) -> str:
   295→    """Mark a feature as in-progress. Call immediately after feature_get_next().
   296→
   297→    This prevents other agent sessions from working on the same feature.
   298→    Use this as soon as you retrieve a feature to work on.
   299→
   300→    Args:
   301→        feature_id: The ID of the feature to mark as in-progress
   302→
   303→    Returns:
   304→        JSON with the updated feature details, or error if not found or already in-progress.
   305→    """
   306→    session = get_session()
   307→    try:
   308→        feature = session.query(Feature).filter(Feature.id == feature_id).first()
   309→
   310→        if feature is None:
   311→            return json.dumps({"error": f"Feature with ID {feature_id} not found"})
   312→
   313→        if feature.passes:
   314→            return json.dumps({"error": f"Feature with ID {feature_id} is already passing"})
   315→
   316→        if feature.in_progress:
   317→            return json.dumps({"error": f"Feature with ID {feature_id} is already in-progress"})
   318→
   319→        feature.in_progress = True
   320→        session.commit()
   321→        session.refresh(feature)
   322→
   323→        return json.dumps(feature.to_dict(), indent=2)
   324→    finally:
   325→        session.close()
   326→
   327→
   328→@mcp.tool()
   329→def feature_clear_in_progress(
   330→    feature_id: Annotated[int, Field(description="The ID of the feature to clear in-progress status", ge=1)]
   331→) -> str:
   332→    """Clear in-progress status from a feature.
   333→
   334→    Use this when abandoning a feature or manually unsticking a stuck feature.
   335→    The feature will return to the pending queue.
   336→
   337→    Args:
   338→        feature_id: The ID of the feature to clear in-progress status
   339→
   340→    Returns:
   341→        JSON with the updated feature details, or error if not found.
   342→    """
   343→    session = get_session()
   344→    try:
   345→        feature = session.query(Feature).filter(Feature.id == feature_id).first()
   346→
   347→        if feature is None:
   348→            return json.dumps({"error": f"Feature with ID {feature_id} not found"})
   349→
   350→        feature.in_progress = False
   351→        session.commit()
   352→        session.refresh(feature)
   353→
   354→        return json.dumps(feature.to_dict(), indent=2)
   355→    finally:
   356→        session.close()
   357→
   358→
   359→@mcp.tool()
   360→def feature_create_bulk(
   361→    features: Annotated[list[dict], Field(description="List of features to create, each with category, name, description, and steps")]
   362→) -> str:
   363→    """Create multiple features in a single operation.
   364→
   365→    Features are assigned sequential priorities based on their order.
   366→    All features start with passes=false.
   367→
   368→    This is typically used by the initializer agent to set up the initial
   369→    feature list from the app specification.
   370→
   371→    Args:
   372→        features: List of features to create, each with:
   373→            - category (str): Feature category
   374→            - name (str): Feature name
   375→            - description (str): Detailed description
   376→            - steps (list[str]): Implementation/test steps
   377→
   378→    Returns:
   379→        JSON with: created (int) - number of features created
   380→    """
   381→    session = get_session()
   382→    try:
   383→        # Get the starting priority
   384→        max_priority_result = session.query(Feature.priority).order_by(Feature.priority.desc()).first()
   385→        start_priority = (max_priority_result[0] + 1) if max_priority_result else 1
   386→
   387→        created_count = 0
   388→        for i, feature_data in enumerate(features):
   389→            # Validate required fields
   390→            if not all(key in feature_data for key in ["category", "name", "description", "steps"]):
   391→                return json.dumps({
   392→                    "error": f"Feature at index {i} missing required fields (category, name, description, steps)"
   393→                })
   394→
   395→            db_feature = Feature(
   396→                priority=start_priority + i,
   397→                category=feature_data["category"],
   398→                name=feature_data["name"],
   399→                description=feature_data["description"],
   400→                steps=feature_data["steps"],
   401→                passes=False,
   402→            )
   403→            session.add(db_feature)
   404→            created_count += 1
   405→
   406→        session.commit()
   407→
   408→        return json.dumps({"created": created_count}, indent=2)
   409→    except Exception as e:
   410→        session.rollback()
   411→        return json.dumps({"error": str(e)})
   412→    finally:
   413→        session.close()
   414→
   415→
   416→if __name__ == "__main__":
   417→    mcp.run()
   418→

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

#!/usr/bin/env python3
"""
MCP Server for Feature Management
==================================

Provides tools to manage features in the autonomous coding system,
replacing the previous FastAPI-based REST API.

Tools:
- feature_get_stats: Get progress statistics
- feature_get_next: Get next feature to implement
- feature_get_for_regression: Get random passing features for testing
- feature_mark_passing: Mark a feature as passing
- feature_skip: Skip a feature (move to end of queue)
- feature_mark_in_progress: Mark a feature as in-progress
- feature_clear_in_progress: Clear in-progress status
- feature_create_bulk: Create multiple features at once
"""

import json
import os
import sys
from contextlib import asynccontextmanager
from pathlib import Path
from typing import Annotated

from mcp.server.fastmcp import FastMCP
from pydantic import BaseModel, Field
from sqlalchemy.sql.expression import func

# Add parent directory to path so we can import from api module
sys.path.insert(0, str(Path(__file__).parent.parent))

from api.database import Feature, create_database
from api.migration import migrate_json_to_sqlite

# Configuration from environment
PROJECT_DIR = Path(os.environ.get("PROJECT_DIR", ".")).resolve()

# Pydantic models for input validation
class MarkPassingInput(BaseModel):
    """Input for marking a feature as passing."""
    feature_id: int = Field(..., description="The ID of the feature to mark as passing", ge=1)

class SkipFeatureInput(BaseModel):
    """Input for skipping a feature."""
    feature_id: int = Field(..., description="The ID of the feature to skip", ge=1)

class MarkInProgressInput(BaseModel):
    """Input for marking a feature as in-progress."""
    feature_id: int = Field(..., description="The ID of the feature to mark as in-progress", ge=1)

class ClearInProgressInput(BaseModel):
    """Input for clearing in-progress status."""
    feature_id: int = Field(..., description="The ID of the feature to clear in-progress status", ge=1)

class RegressionInput(BaseModel):
    """Input for getting regression features."""
    limit: int = Field(default=3, ge=1, le=10, description="Maximum number of passing features to return")

class FeatureCreateItem(BaseModel):
    """Schema for creating a single feature."""
    category: str = Field(..., min_length=1, max_length=100, description="Feature category")
    name: str = Field(..., min_length=1, max_length=255, description="Feature name")
    description: str = Field(..., min_length=1, description="Detailed description")
    steps: list[str] = Field(..., min_length=1, description="Implementation/test steps")

class BulkCreateInput(BaseModel):
    """Input for bulk creating features."""
    features: list[FeatureCreateItem] = Field(..., min_length=1, description="List of features to create")

# Global database session maker (initialized on startup)
_session_maker = None
_engine = None

@asynccontextmanager
async def server_lifespan(server: FastMCP):
    """Initialize database on startup, cleanup on shutdown."""
    global _session_maker, _engine

    # Create project directory if it doesn't exist
    PROJECT_DIR.mkdir(parents=True, exist_ok=True)

    # Initialize database
    _engine, _session_maker = create_database(PROJECT_DIR)

    # Run migration if needed (converts legacy JSON to SQLite)
    migrate_json_to_sqlite(PROJECT_DIR, _session_maker)

    yield

    # Cleanup
    if _engine:
        _engine.dispose()

# Initialize the MCP server
mcp = FastMCP("features", lifespan=server_lifespan)

def get_session():
    """Get a new database session."""
    if _session_maker is None:
        raise RuntimeError("Database not initialized")
    return _session_maker()

@mcp.tool()
def feature_get_stats() -> str:
    """Get statistics about feature completion progress.

    Returns the number of passing features, in-progress features, total features,
    and completion percentage. Use this to track overall progress of the implementation.

    Returns:
        JSON with: passing (int), in_progress (int), total (int), percentage (float)
    """
    session = get_session()
    try:
        total = session.query(Feature).count()
        passing = session.query(Feature).filter(Feature.passes == True).count()
        in_progress = session.query(Feature).filter(Feature.in_progress == True).count()
        percentage = round((passing / total) * 100, 1) if total > 0 else 0.0

        return json.dumps({
            "passing": passing,
            "in_progress": in_progress,
            "total": total,
            "percentage": percentage
        }, indent=2)
    finally:
        session.close()

@mcp.tool()
def feature_get_next() -> str:
    """Get the highest-priority pending feature to work on.

    Returns the feature with the lowest priority number that has passes=false.
    Use this at the start of each coding session to determine what to implement next.

    Returns:
        JSON with feature details (id, priority, category, name, description, steps, passes, in_progress)
        or error message if all features are passing.
    """
    session = get_session()
    try:
        feature = (
            session.query(Feature)
            .filter(Feature.passes == False)
            .order_by(Feature.priority.asc(), Feature.id.asc())
            .first()
        )

        if feature is None:
            return json.dumps({"error": "All features are passing! No more work to do."})

        return json.dumps(feature.to_dict(), indent=2)
    finally:
        session.close()

@mcp.tool()
def feature_get_for_regression(
    limit: Annotated[int, Field(default=3, ge=1, le=10, description="Maximum number of passing features to return")] = 3
) -> str:
    """Get random passing features for regression testing.

    Returns a random selection of features that are currently passing.
    Use this to verify that previously implemented features still work
    after making changes.

    Args:
        limit: Maximum number of features to return (1-10, default 3)

    Returns:
        JSON with: features (list of feature objects), count (int)
    """
    session = get_session()
    try:
        features = (
            session.query(Feature)
            .filter(Feature.passes == True)
            .order_by(func.random())
            .limit(limit)
            .all()
        )

        return json.dumps({
            "features": [f.to_dict() for f in features],
            "count": len(features)
        }, indent=2)
    finally:
        session.close()

@mcp.tool()
def feature_mark_passing(
    feature_id: Annotated[int, Field(description="The ID of the feature to mark as passing", ge=1)]
) -> str:
    """Mark a feature as passing after successful implementation.

    Updates the feature's passes field to true and clears the in_progress flag.
    Use this after you have implemented the feature and verified it works correctly.

    Args:
        feature_id: The ID of the feature to mark as passing

    Returns:
        JSON with the updated feature details, or error if not found.
    """
    session = get_session()
    try:
        feature = session.query(Feature).filter(Feature.id == feature_id).first()

        if feature is None:
            return json.dumps({"error": f"Feature with ID {feature_id} not found"})

        feature.passes = True
        feature.in_progress = False
        session.commit()
        session.refresh(feature)

        return json.dumps(feature.to_dict(), indent=2)
    finally:
        session.close()

@mcp.tool()
def feature_skip(
    feature_id: Annotated[int, Field(description="The ID of the feature to skip", ge=1)]
) -> str:
    """Skip a feature by moving it to the end of the priority queue.

    Use this when a feature cannot be implemented yet due to:
    - Dependencies on other features that aren't implemented yet
    - External blockers (missing assets, unclear requirements)
    - Technical prerequisites that need to be addressed first

    The feature's priority is set to max_priority + 1, so it will be
    worked on after all other pending features. Also clears the in_progress
    flag so the feature returns to "pending" status.

    Args:
        feature_id: The ID of the feature to skip

    Returns:
        JSON with skip details: id, name, old_priority, new_priority, message
    """
    session = get_session()
    try:
        feature = session.query(Feature).filter(Feature.id == feature_id).first()

        if feature is None:
            return json.dumps({"error": f"Feature with ID {feature_id} not found"})

        if feature.passes:
            return json.dumps({"error": "Cannot skip a feature that is already passing"})

        old_priority = feature.priority

        # Get max priority and set this feature to max + 1
        max_priority_result = session.query(Feature.priority).order_by(Feature.priority.desc()).first()
        new_priority = (max_priority_result[0] + 1) if max_priority_result else 1

        feature.priority = new_priority
        feature.in_progress = False
        session.commit()
        session.refresh(feature)

        return json.dumps({
            "id": feature.id,
            "name": feature.name,
            "old_priority": old_priority,
            "new_priority": new_priority,
            "message": f"Feature '{feature.name}' moved to end of queue"
        }, indent=2)
    finally:
        session.close()

@mcp.tool()
def feature_mark_in_progress(
    feature_id: Annotated[int, Field(description="The ID of the feature to mark as in-progress", ge=1)]
) -> str:
    """Mark a feature as in-progress. Call immediately after feature_get_next().

    This prevents other agent sessions from working on the same feature.
    Use this as soon as you retrieve a feature to work on.

    Args:
        feature_id: The ID of the feature to mark as in-progress

    Returns:
        JSON with the updated feature details, or error if not found or already in-progress.
    """
    session = get_session()
    try:
        feature = session.query(Feature).filter(Feature.id == feature_id).first()

        if feature is None:
            return json.dumps({"error": f"Feature with ID {feature_id} not found"})

        if feature.passes:
            return json.dumps({"error": f"Feature with ID {feature_id} is already passing"})

        if feature.in_progress:
            return json.dumps({"error": f"Feature with ID {feature_id} is already in-progress"})

        feature.in_progress = True
        session.commit()
        session.refresh(feature)

        return json.dumps(feature.to_dict(), indent=2)
    finally:
        session.close()

@mcp.tool()
def feature_clear_in_progress(
    feature_id: Annotated[int, Field(description="The ID of the feature to clear in-progress status", ge=1)]
) -> str:
    """Clear in-progress status from a feature.

    Use this when abandoning a feature or manually unsticking a stuck feature.
    The feature will return to the pending queue.

    Args:
        feature_id: The ID of the feature to clear in-progress status

    Returns:
        JSON with the updated feature details, or error if not found.
    """
    session = get_session()
    try:
        feature = session.query(Feature).filter(Feature.id == feature_id).first()

        if feature is None:
            return json.dumps({"error": f"Feature with ID {feature_id} not found"})

        feature.in_progress = False
        session.commit()
        session.refresh(feature)

        return json.dumps(feature.to_dict(), indent=2)
    finally:
        session.close()

@mcp.tool()
def feature_create_bulk(
    features: Annotated[list[dict], Field(description="List of features to create, each with category, name, description, and steps")]
) -> str:
    """Create multiple features in a single operation.

    Features are assigned sequential priorities based on their order.
    All features start with passes=false.

    This is typically used by the initializer agent to set up the initial
    feature list from the app specification.

    Args:
        features: List of features to create, each with:
            - category (str): Feature category
            - name (str): Feature name
            - description (str): Detailed description
            - steps (list[str]): Implementation/test steps

    Returns:
        JSON with: created (int) - number of features created
    """
    session = get_session()
    try:
        # Get the starting priority
        max_priority_result = session.query(Feature.priority).order_by(Feature.priority.desc()).first()
        start_priority = (max_priority_result[0] + 1) if max_priority_result else 1

        created_count = 0
        for i, feature_data in enumerate(features):
            # Validate required fields
            if not all(key in feature_data for key in ["category", "name", "description", "steps"]):
                return json.dumps({
                    "error": f"Feature at index {i} missing required fields (category, name, description, steps)"
                })

            db_feature = Feature(
                priority=start_priority + i,
                category=feature_data["category"],
                name=feature_data["name"],
                description=feature_data["description"],
                steps=feature_data["steps"],
                passes=False,
            )
            session.add(db_feature)
            created_count += 1

        session.commit()

        return json.dumps({"created": created_count}, indent=2)
    except Exception as e:
        session.rollback()
        return json.dumps({"error": str(e)})
    finally:
        session.close()

if __name__ == "__main__":
    mcp.run()


     1→"""
     2→Security Hooks for Autonomous Coding Agent
     3→==========================================
     4→
     5→Pre-tool-use hooks that validate bash commands for security.
     6→Uses an allowlist approach - only explicitly permitted commands can run.
     7→"""
     8→
     9→import os
    10→import shlex
    11→
    12→# Allowed commands for development tasks
    13→# Minimal set needed for the autonomous coding demo
    14→ALLOWED_COMMANDS = {
    15→    # File inspection
    16→    "ls",
    17→    "cat",
    18→    "head",
    19→    "tail",
    20→    "wc",
    21→    "grep",
    22→    # File operations (agent uses SDK tools for most file ops, but cp/mkdir needed occasionally)
    23→    "cp",
    24→    "mkdir",
    25→    "chmod",  # For making scripts executable; validated separately
    26→    # Directory
    27→    "pwd",
    28→    # Output
    29→    "echo",
    30→    # Node.js development
    31→    "npm",
    32→    "npx",
    33→    "pnpm",  # Project uses pnpm
    34→    "node",
    35→    # Version control
    36→    "git",
    37→    # Docker (for PostgreSQL)
    38→    "docker",
    39→    # Process management
    40→    "ps",
    41→    "lsof",
    42→    "sleep",
    43→    "kill",  # Kill by PID
    44→    "pkill",  # For killing dev servers; validated separately
    45→    # Network/API testing
    46→    "curl",
    47→    # File operations
    48→    "mv",
    49→    "rm",  # Use with caution
    50→    "touch",
    51→    # Shell scripts
    52→    "sh",
    53→    "bash",
    54→    # Script execution
    55→    "init.sh",  # Init scripts; validated separately
    56→}
    57→
    58→# Commands that need additional validation even when in the allowlist
    59→COMMANDS_NEEDING_EXTRA_VALIDATION = {"pkill", "chmod", "init.sh"}
    60→
    61→
    62→def split_command_segments(command_string: str) -> list[str]:
    63→    """
    64→    Split a compound command into individual command segments.
    65→
    66→    Handles command chaining (&&, ||, ;) but not pipes (those are single commands).
    67→
    68→    Args:
    69→        command_string: The full shell command
    70→
    71→    Returns:
    72→        List of individual command segments
    73→    """
    74→    import re
    75→
    76→    # Split on && and || while preserving the ability to handle each segment
    77→    # This regex splits on && or || that aren't inside quotes
    78→    segments = re.split(r"\s*(?:&&|\|\|)\s*", command_string)
    79→
    80→    # Further split on semicolons
    81→    result = []
    82→    for segment in segments:
    83→        sub_segments = re.split(r'(?<!["\'])\s*;\s*(?!["\'])', segment)
    84→        for sub in sub_segments:
    85→            sub = sub.strip()
    86→            if sub:
    87→                result.append(sub)
    88→
    89→    return result
    90→
    91→
    92→def extract_commands(command_string: str) -> list[str]:
    93→    """
    94→    Extract command names from a shell command string.
    95→
    96→    Handles pipes, command chaining (&&, ||, ;), and subshells.
    97→    Returns the base command names (without paths).
    98→
    99→    Args:
   100→        command_string: The full shell command
   101→
   102→    Returns:
   103→        List of command names found in the string
   104→    """
   105→    commands = []
   106→
   107→    # shlex doesn't treat ; as a separator, so we need to pre-process
   108→    import re
   109→
   110→    # Split on semicolons that aren't inside quotes (simple heuristic)
   111→    # This handles common cases like "echo hello; ls"
   112→    segments = re.split(r'(?<!["\'])\s*;\s*(?!["\'])', command_string)
   113→
   114→    for segment in segments:
   115→        segment = segment.strip()
   116→        if not segment:
   117→            continue
   118→
   119→        try:
   120→            tokens = shlex.split(segment)
   121→        except ValueError:
   122→            # Malformed command (unclosed quotes, etc.)
   123→            # Return empty to trigger block (fail-safe)
   124→            return []
   125→
   126→        if not tokens:
   127→            continue
   128→
   129→        # Track when we expect a command vs arguments
   130→        expect_command = True
   131→
   132→        for token in tokens:
   133→            # Shell operators indicate a new command follows
   134→            if token in ("|", "||", "&&", "&"):
   135→                expect_command = True
   136→                continue
   137→
   138→            # Skip shell keywords that precede commands
   139→            if token in (
   140→                "if",
   141→                "then",
   142→                "else",
   143→                "elif",
   144→                "fi",
   145→                "for",
   146→                "while",
   147→                "until",
   148→                "do",
   149→                "done",
   150→                "case",
   151→                "esac",
   152→                "in",
   153→                "!",
   154→                "{",
   155→                "}",
   156→            ):
   157→                continue
   158→
   159→            # Skip flags/options
   160→            if token.startswith("-"):
   161→                continue
   162→
   163→            # Skip variable assignments (VAR=value)
   164→            if "=" in token and not token.startswith("="):
   165→                continue
   166→
   167→            if expect_command:
   168→                # Extract the base command name (handle paths like /usr/bin/python)
   169→                cmd = os.path.basename(token)
   170→                commands.append(cmd)
   171→                expect_command = False
   172→
   173→    return commands
   174→
   175→
   176→def validate_pkill_command(command_string: str) -> tuple[bool, str]:
   177→    """
   178→    Validate pkill commands - only allow killing dev-related processes.
   179→
   180→    Uses shlex to parse the command, avoiding regex bypass vulnerabilities.
   181→
   182→    Returns:
   183→        Tuple of (is_allowed, reason_if_blocked)
   184→    """
   185→    # Allowed process names for pkill
   186→    allowed_process_names = {
   187→        "node",
   188→        "npm",
   189→        "npx",
   190→        "vite",
   191→        "next",
   192→    }
   193→
   194→    try:
   195→        tokens = shlex.split(command_string)
   196→    except ValueError:
   197→        return False, "Could not parse pkill command"
   198→
   199→    if not tokens:
   200→        return False, "Empty pkill command"
   201→
   202→    # Separate flags from arguments
   203→    args = []
   204→    for token in tokens[1:]:
   205→        if not token.startswith("-"):
   206→            args.append(token)
   207→
   208→    if not args:
   209→        return False, "pkill requires a process name"
   210→
   211→    # The target is typically the last non-flag argument
   212→    target = args[-1]
   213→
   214→    # For -f flag (full command line match), extract the first word as process name
   215→    # e.g., "pkill -f 'node server.js'" -> target is "node server.js", process is "node"
   216→    if " " in target:
   217→        target = target.split()[0]
   218→
   219→    if target in allowed_process_names:
   220→        return True, ""
   221→    return False, f"pkill only allowed for dev processes: {allowed_process_names}"
   222→
   223→
   224→def validate_chmod_command(command_string: str) -> tuple[bool, str]:
   225→    """
   226→    Validate chmod commands - only allow making files executable with +x.
   227→
   228→    Returns:
   229→        Tuple of (is_allowed, reason_if_blocked)
   230→    """
   231→    try:
   232→        tokens = shlex.split(command_string)
   233→    except ValueError:
   234→        return False, "Could not parse chmod command"
   235→
   236→    if not tokens or tokens[0] != "chmod":
   237→        return False, "Not a chmod command"
   238→
   239→    # Look for the mode argument
   240→    # Valid modes: +x, u+x, a+x, etc. (anything ending with +x for execute permission)
   241→    mode = None
   242→    files = []
   243→
   244→    for token in tokens[1:]:
   245→        if token.startswith("-"):
   246→            # Skip flags like -R (we don't allow recursive chmod anyway)
   247→            return False, "chmod flags are not allowed"
   248→        elif mode is None:
   249→            mode = token
   250→        else:
   251→            files.append(token)
   252→
   253→    if mode is None:
   254→        return False, "chmod requires a mode"
   255→
   256→    if not files:
   257→        return False, "chmod requires at least one file"
   258→
   259→    # Only allow +x variants (making files executable)
   260→    # This matches: +x, u+x, g+x, o+x, a+x, ug+x, etc.
   261→    import re
   262→
   263→    if not re.match(r"^[ugoa]*\+x$", mode):
   264→        return False, f"chmod only allowed with +x mode, got: {mode}"
   265→
   266→    return True, ""
   267→
   268→
   269→def validate_init_script(command_string: str) -> tuple[bool, str]:
   270→    """
   271→    Validate init.sh script execution - only allow ./init.sh.
   272→
   273→    Returns:
   274→        Tuple of (is_allowed, reason_if_blocked)
   275→    """
   276→    try:
   277→        tokens = shlex.split(command_string)
   278→    except ValueError:
   279→        return False, "Could not parse init script command"
   280→
   281→    if not tokens:
   282→        return False, "Empty command"
   283→
   284→    # The command should be exactly ./init.sh (possibly with arguments)
   285→    script = tokens[0]
   286→
   287→    # Allow ./init.sh or paths ending in /init.sh
   288→    if script == "./init.sh" or script.endswith("/init.sh"):
   289→        return True, ""
   290→
   291→    return False, f"Only ./init.sh is allowed, got: {script}"
   292→
   293→
   294→def get_command_for_validation(cmd: str, segments: list[str]) -> str:
   295→    """
   296→    Find the specific command segment that contains the given command.
   297→
   298→    Args:
   299→        cmd: The command name to find
   300→        segments: List of command segments
   301→
   302→    Returns:
   303→        The segment containing the command, or empty string if not found
   304→    """
   305→    for segment in segments:
   306→        segment_commands = extract_commands(segment)
   307→        if cmd in segment_commands:
   308→            return segment
   309→    return ""
   310→
   311→
   312→async def bash_security_hook(input_data, tool_use_id=None, context=None):
   313→    """
   314→    Pre-tool-use hook that validates bash commands using an allowlist.
   315→
   316→    Only commands in ALLOWED_COMMANDS are permitted.
   317→
   318→    Args:
   319→        input_data: Dict containing tool_name and tool_input
   320→        tool_use_id: Optional tool use ID
   321→        context: Optional context
   322→
   323→    Returns:
   324→        Empty dict to allow, or {"decision": "block", "reason": "..."} to block
   325→    """
   326→    if input_data.get("tool_name") != "Bash":
   327→        return {}
   328→
   329→    command = input_data.get("tool_input", {}).get("command", "")
   330→    if not command:
   331→        return {}
   332→
   333→    # Extract all commands from the command string
   334→    commands = extract_commands(command)
   335→
   336→    if not commands:
   337→        # Could not parse - fail safe by blocking
   338→        return {
   339→            "decision": "block",
   340→            "reason": f"Could not parse command for security validation: {command}",
   341→        }
   342→
   343→    # Split into segments for per-command validation
   344→    segments = split_command_segments(command)
   345→
   346→    # Check each command against the allowlist
   347→    for cmd in commands:
   348→        if cmd not in ALLOWED_COMMANDS:
   349→            return {
   350→                "decision": "block",
   351→                "reason": f"Command '{cmd}' is not in the allowed commands list",
   352→            }
   353→
   354→        # Additional validation for sensitive commands
   355→        if cmd in COMMANDS_NEEDING_EXTRA_VALIDATION:
   356→            # Find the specific segment containing this command
   357→            cmd_segment = get_command_for_validation(cmd, segments)
   358→            if not cmd_segment:
   359→                cmd_segment = command  # Fallback to full command
   360→
   361→            if cmd == "pkill":
   362→                allowed, reason = validate_pkill_command(cmd_segment)
   363→                if not allowed:
   364→                    return {"decision": "block", "reason": reason}
   365→            elif cmd == "chmod":
   366→                allowed, reason = validate_chmod_command(cmd_segment)
   367→                if not allowed:
   368→                    return {"decision": "block", "reason": reason}
   369→            elif cmd == "init.sh":
   370→                allowed, reason = validate_init_script(cmd_segment)
   371→                if not allowed:
   372→                    return {"decision": "block", "reason": reason}
   373→
   374→    return {}
   375→

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

"""
Security Hooks for Autonomous Coding Agent
==========================================

Pre-tool-use hooks that validate bash commands for security.
Uses an allowlist approach - only explicitly permitted commands can run.
"""

import os
import shlex

# Allowed commands for development tasks
# Minimal set needed for the autonomous coding demo
ALLOWED_COMMANDS = {
    # File inspection
    "ls",
    "cat",
    "head",
    "tail",
    "wc",
    "grep",
    # File operations (agent uses SDK tools for most file ops, but cp/mkdir needed occasionally)
    "cp",
    "mkdir",
    "chmod",  # For making scripts executable; validated separately
    # Directory
    "pwd",
    # Output
    "echo",
    # Node.js development
    "npm",
    "npx",
    "pnpm",  # Project uses pnpm
    "node",
    # Version control
    "git",
    # Docker (for PostgreSQL)
    "docker",
    # Process management
    "ps",
    "lsof",
    "sleep",
    "kill",  # Kill by PID
    "pkill",  # For killing dev servers; validated separately
    # Network/API testing
    "curl",
    # File operations
    "mv",
    "rm",  # Use with caution
    "touch",
    # Shell scripts
    "sh",
    "bash",
    # Script execution
    "init.sh",  # Init scripts; validated separately
}

# Commands that need additional validation even when in the allowlist
COMMANDS_NEEDING_EXTRA_VALIDATION = {"pkill", "chmod", "init.sh"}

def split_command_segments(command_string: str) -> list[str]:
    """
    Split a compound command into individual command segments.

    Handles command chaining (&&, ||, ;) but not pipes (those are single commands).

    Args:
        command_string: The full shell command

    Returns:
        List of individual command segments
    """
    import re

    # Split on && and || while preserving the ability to handle each segment
    # This regex splits on && or || that aren't inside quotes
    segments = re.split(r"\s*(?:&&|\|\|)\s*", command_string)

    # Further split on semicolons
    result = []
    for segment in segments:
        sub_segments = re.split(r'(?<!["\'])\s*;\s*(?!["\'])', segment)
        for sub in sub_segments:
            sub = sub.strip()
            if sub:
                result.append(sub)

    return result

def extract_commands(command_string: str) -> list[str]:
    """
    Extract command names from a shell command string.

    Handles pipes, command chaining (&&, ||, ;), and subshells.
    Returns the base command names (without paths).

    Args:
        command_string: The full shell command

    Returns:
        List of command names found in the string
    """
    commands = []

    # shlex doesn't treat ; as a separator, so we need to pre-process
    import re

    # Split on semicolons that aren't inside quotes (simple heuristic)
    # This handles common cases like "echo hello; ls"
    segments = re.split(r'(?<!["\'])\s*;\s*(?!["\'])', command_string)

    for segment in segments:
        segment = segment.strip()
        if not segment:
            continue

        try:
            tokens = shlex.split(segment)
        except ValueError:
            # Malformed command (unclosed quotes, etc.)
            # Return empty to trigger block (fail-safe)
            return []

        if not tokens:
            continue

        # Track when we expect a command vs arguments
        expect_command = True

        for token in tokens:
            # Shell operators indicate a new command follows
            if token in ("|", "||", "&&", "&"):
                expect_command = True
                continue

            # Skip shell keywords that precede commands
            if token in (
                "if",
                "then",
                "else",
                "elif",
                "fi",
                "for",
                "while",
                "until",
                "do",
                "done",
                "case",
                "esac",
                "in",
                "!",
                "{",
                "}",
            ):
                continue

            # Skip flags/options
            if token.startswith("-"):
                continue

            # Skip variable assignments (VAR=value)
            if "=" in token and not token.startswith("="):
                continue

            if expect_command:
                # Extract the base command name (handle paths like /usr/bin/python)
                cmd = os.path.basename(token)
                commands.append(cmd)
                expect_command = False

    return commands

def validate_pkill_command(command_string: str) -> tuple[bool, str]:
    """
    Validate pkill commands - only allow killing dev-related processes.

    Uses shlex to parse the command, avoiding regex bypass vulnerabilities.

    Returns:
        Tuple of (is_allowed, reason_if_blocked)
    """
    # Allowed process names for pkill
    allowed_process_names = {
        "node",
        "npm",
        "npx",
        "vite",
        "next",
    }

    try:
        tokens = shlex.split(command_string)
    except ValueError:
        return False, "Could not parse pkill command"

    if not tokens:
        return False, "Empty pkill command"

    # Separate flags from arguments
    args = []
    for token in tokens[1:]:
        if not token.startswith("-"):
            args.append(token)

    if not args:
        return False, "pkill requires a process name"

    # The target is typically the last non-flag argument
    target = args[-1]

    # For -f flag (full command line match), extract the first word as process name
    # e.g., "pkill -f 'node server.js'" -> target is "node server.js", process is "node"
    if " " in target:
        target = target.split()[0]

    if target in allowed_process_names:
        return True, ""
    return False, f"pkill only allowed for dev processes: {allowed_process_names}"

def validate_chmod_command(command_string: str) -> tuple[bool, str]:
    """
    Validate chmod commands - only allow making files executable with +x.

    Returns:
        Tuple of (is_allowed, reason_if_blocked)
    """
    try:
        tokens = shlex.split(command_string)
    except ValueError:
        return False, "Could not parse chmod command"

    if not tokens or tokens[0] != "chmod":
        return False, "Not a chmod command"

    # Look for the mode argument
    # Valid modes: +x, u+x, a+x, etc. (anything ending with +x for execute permission)
    mode = None
    files = []

    for token in tokens[1:]:
        if token.startswith("-"):
            # Skip flags like -R (we don't allow recursive chmod anyway)
            return False, "chmod flags are not allowed"
        elif mode is None:
            mode = token
        else:
            files.append(token)

    if mode is None:
        return False, "chmod requires a mode"

    if not files:
        return False, "chmod requires at least one file"

    # Only allow +x variants (making files executable)
    # This matches: +x, u+x, g+x, o+x, a+x, ug+x, etc.
    import re

    if not re.match(r"^[ugoa]*\+x$", mode):
        return False, f"chmod only allowed with +x mode, got: {mode}"

    return True, ""

def validate_init_script(command_string: str) -> tuple[bool, str]:
    """
    Validate init.sh script execution - only allow ./init.sh.

    Returns:
        Tuple of (is_allowed, reason_if_blocked)
    """
    try:
        tokens = shlex.split(command_string)
    except ValueError:
        return False, "Could not parse init script command"

    if not tokens:
        return False, "Empty command"

    # The command should be exactly ./init.sh (possibly with arguments)
    script = tokens[0]

    # Allow ./init.sh or paths ending in /init.sh
    if script == "./init.sh" or script.endswith("/init.sh"):
        return True, ""

    return False, f"Only ./init.sh is allowed, got: {script}"

def get_command_for_validation(cmd: str, segments: list[str]) -> str:
    """
    Find the specific command segment that contains the given command.

    Args:
        cmd: The command name to find
        segments: List of command segments

    Returns:
        The segment containing the command, or empty string if not found
    """
    for segment in segments:
        segment_commands = extract_commands(segment)
        if cmd in segment_commands:
            return segment
    return ""

async def bash_security_hook(input_data, tool_use_id=None, context=None):
    """
    Pre-tool-use hook that validates bash commands using an allowlist.

    Only commands in ALLOWED_COMMANDS are permitted.

    Args:
        input_data: Dict containing tool_name and tool_input
        tool_use_id: Optional tool use ID
        context: Optional context

    Returns:
        Empty dict to allow, or {"decision": "block", "reason": "..."} to block
    """
    if input_data.get("tool_name") != "Bash":
        return {}

    command = input_data.get("tool_input", {}).get("command", "")
    if not command:
        return {}

    # Extract all commands from the command string
    commands = extract_commands(command)

    if not commands:
        # Could not parse - fail safe by blocking
        return {
            "decision": "block",
            "reason": f"Could not parse command for security validation: {command}",
        }

    # Split into segments for per-command validation
    segments = split_command_segments(command)

    # Check each command against the allowlist
    for cmd in commands:
        if cmd not in ALLOWED_COMMANDS:
            return {
                "decision": "block",
                "reason": f"Command '{cmd}' is not in the allowed commands list",
            }

        # Additional validation for sensitive commands
        if cmd in COMMANDS_NEEDING_EXTRA_VALIDATION:
            # Find the specific segment containing this command
            cmd_segment = get_command_for_validation(cmd, segments)
            if not cmd_segment:
                cmd_segment = command  # Fallback to full command

            if cmd == "pkill":
                allowed, reason = validate_pkill_command(cmd_segment)
                if not allowed:
                    return {"decision": "block", "reason": reason}
            elif cmd == "chmod":
                allowed, reason = validate_chmod_command(cmd_segment)
                if not allowed:
                    return {"decision": "block", "reason": reason}
            elif cmd == "init.sh":
                allowed, reason = validate_init_script(cmd_segment)
                if not allowed:
                    return {"decision": "block", "reason": reason}

    return {}


     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)


Let me also explore the UI code and the server routers:


Bash

List UI src directory


List server routers


List server services


total 34
drwxr-xr-x 1 AzureAD+MikeSwanson 4096    0 Jan  8 13:56 .
drwxr-xr-x 1 AzureAD+MikeSwanson 4096    0 Jan  8 13:56 ..
-rw-r--r-- 1 AzureAD+MikeSwanson 4096 9696 Jan  8 08:02 App.tsx
drwxr-xr-x 1 AzureAD+MikeSwanson 4096    0 Jan  8 13:56 components
drwxr-xr-x 1 AzureAD+MikeSwanson 4096    0 Jan  8 13:56 hooks
drwxr-xr-x 1 AzureAD+MikeSwanson 4096    0 Jan  8 13:56 lib
-rw-r--r-- 1 AzureAD+MikeSwanson 4096  531 Jan  8 08:02 main.tsx
drwxr-xr-x 1 AzureAD+MikeSwanson 4096    0 Jan  8 13:56 styles
-rw-r--r-- 1 AzureAD+MikeSwanson 4096   38 Jan  8 08:02 vite-env.d.ts


total 85
drwxr-xr-x 1 AzureAD+MikeSwanson 4096     0 Jan  8 13:53 .
drwxr-xr-x 1 AzureAD+MikeSwanson 4096     0 Jan  8 13:53 ..
-rw-r--r-- 1 AzureAD+MikeSwanson 4096   550 Jan  8 08:02 __init__.py
drwxr-xr-x 1 AzureAD+MikeSwanson 4096     0 Jan  8 13:53 __pycache__
-rw-r--r-- 1 AzureAD+MikeSwanson 4096  3608 Jan  8 08:02 agent.py
-rw-r--r-- 1 AzureAD+MikeSwanson 4096 11974 Jan  8 08:02 assistant_chat.py
-rw-r--r-- 1 AzureAD+MikeSwanson 4096  9684 Jan  8 08:02 features.py
-rw-r--r-- 1 AzureAD+MikeSwanson 4096 14883 Jan  8 08:02 filesystem.py
-rw-r--r-- 1 AzureAD+MikeSwanson 4096 10334 Jan  8 08:02 projects.py
-rw-r--r-- 1 AzureAD+MikeSwanson 4096 14324 Jan  8 08:02 spec_creation.py


total 73
drwxr-xr-x 1 AzureAD+MikeSwanson 4096     0 Jan  8 13:53 .
drwxr-xr-x 1 AzureAD+MikeSwanson 4096     0 Jan  8 13:53 ..
-rw-r--r-- 1 AzureAD+MikeSwanson 4096   176 Jan  8 08:02 __init__.py
drwxr-xr-x 1 AzureAD+MikeSwanson 4096     0 Jan  8 13:53 __pycache__
-rw-r--r-- 1 AzureAD+MikeSwanson 4096 12981 Jan  8 08:02 assistant_chat_session.py
-rw-r--r-- 1 AzureAD+MikeSwanson 4096  7786 Jan  8 08:02 assistant_database.py
-rw-r--r-- 1 AzureAD+MikeSwanson 4096 14512 Jan  8 08:02 process_manager.py
-rw-r--r-- 1 AzureAD+MikeSwanson 4096 19490 Jan  8 08:02 spec_chat_session.py