feat: Add "Expand Project" for bulk AI-powered feature creation

Adds the ability to add multiple features to an existing project through
a natural language conversation with Claude, similar to how initial spec
creation works.

Features:
- New "Expand" button in header (keyboard shortcut: E)
- Full-screen chat interface for describing new features
- Claude reads existing app_spec.txt for context
- Features created directly in database after user approval
- Bulk feature creation endpoint for batch operations

New files:
- .claude/commands/expand-project.md - Claude skill for expansion
- server/services/expand_chat_session.py - Chat session service
- server/routers/expand_project.py - WebSocket endpoint
- ui/src/components/ExpandProjectChat.tsx - Chat UI
- ui/src/components/ExpandProjectModal.tsx - Modal wrapper
- ui/src/hooks/useExpandChat.ts - WebSocket hook

Modified:
- Added POST /bulk endpoint to features router
- Added FeatureBulkCreate schemas
- Integrated Expand button and modal in App.tsx

Co-Authored-By: Claude <noreply@anthropic.com>

server/main.py

@@ -18,6 +18,7 @@ from fastapi.staticfiles import StaticFiles
 from .routers import (
     agent_router,
     assistant_chat_router,
+    expand_project_router,
     features_router,
     filesystem_router,
     projects_router,
@@ -25,6 +26,7 @@ from .routers import (
 )
 from .schemas import SetupStatus
 from .services.assistant_chat_session import cleanup_all_sessions as cleanup_assistant_sessions
+from .services.expand_chat_session import cleanup_all_expand_sessions
 from .services.process_manager import cleanup_all_managers
 from .websocket import project_websocket
 
@@ -38,9 +40,10 @@ async def lifespan(app: FastAPI):
     """Lifespan context manager for startup and shutdown."""
     # Startup
     yield
-    # Shutdown - cleanup all running agents and assistant sessions
+    # Shutdown - cleanup all running agents and sessions
     await cleanup_all_managers()
     await cleanup_assistant_sessions()
+    await cleanup_all_expand_sessions()
 
 
 # Create FastAPI app
@@ -90,6 +93,7 @@ app.include_router(projects_router)
 app.include_router(features_router)
 app.include_router(agent_router)
 app.include_router(spec_creation_router)
+app.include_router(expand_project_router)
 app.include_router(filesystem_router)
 app.include_router(assistant_chat_router)
 

server/routers/__init__.py

@@ -7,6 +7,7 @@ FastAPI routers for different API endpoints.
 
 from .agent import router as agent_router
 from .assistant_chat import router as assistant_chat_router
+from .expand_project import router as expand_project_router
 from .features import router as features_router
 from .filesystem import router as filesystem_router
 from .projects import router as projects_router
@@ -17,6 +18,7 @@ __all__ = [
     "features_router",
     "agent_router",
     "spec_creation_router",
+    "expand_project_router",
     "filesystem_router",
     "assistant_chat_router",
 ]

server/routers/expand_project.py

@@ -0,0 +1,246 @@
+"""
+Expand Project Router
+=====================
+
+WebSocket and REST endpoints for interactive project expansion with Claude.
+Allows adding multiple features to existing projects via natural language.
+"""
+
+import json
+import logging
+import re
+from pathlib import Path
+from typing import Optional
+
+from fastapi import APIRouter, HTTPException, WebSocket, WebSocketDisconnect
+from pydantic import BaseModel, ValidationError
+
+from ..schemas import ImageAttachment
+from ..services.expand_chat_session import (
+    ExpandChatSession,
+    create_expand_session,
+    get_expand_session,
+    list_expand_sessions,
+    remove_expand_session,
+)
+
+logger = logging.getLogger(__name__)
+
+router = APIRouter(prefix="/api/expand", tags=["expand-project"])
+
+# Root directory
+ROOT_DIR = Path(__file__).parent.parent.parent
+
+
+def _get_project_path(project_name: str) -> Path:
+    """Get project path from registry."""
+    import sys
+    root = Path(__file__).parent.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 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))
+
+
+# ============================================================================
+# REST Endpoints
+# ============================================================================
+
+class ExpandSessionStatus(BaseModel):
+    """Status of an expansion session."""
+    project_name: str
+    is_active: bool
+    is_complete: bool
+    features_created: int
+    message_count: int
+
+
+@router.get("/sessions", response_model=list[str])
+async def list_expand_sessions_endpoint():
+    """List all active expansion sessions."""
+    return list_expand_sessions()
+
+
+@router.get("/sessions/{project_name}", response_model=ExpandSessionStatus)
+async def get_expand_session_status(project_name: str):
+    """Get status of an expansion session."""
+    if not validate_project_name(project_name):
+        raise HTTPException(status_code=400, detail="Invalid project name")
+
+    session = get_expand_session(project_name)
+    if not session:
+        raise HTTPException(status_code=404, detail="No active expansion session for this project")
+
+    return ExpandSessionStatus(
+        project_name=project_name,
+        is_active=True,
+        is_complete=session.is_complete(),
+        features_created=session.get_features_created(),
+        message_count=len(session.get_messages()),
+    )
+
+
+@router.delete("/sessions/{project_name}")
+async def cancel_expand_session(project_name: str):
+    """Cancel and remove an expansion session."""
+    if not validate_project_name(project_name):
+        raise HTTPException(status_code=400, detail="Invalid project name")
+
+    session = get_expand_session(project_name)
+    if not session:
+        raise HTTPException(status_code=404, detail="No active expansion session for this project")
+
+    await remove_expand_session(project_name)
+    return {"success": True, "message": "Expansion session cancelled"}
+
+
+# ============================================================================
+# WebSocket Endpoint
+# ============================================================================
+
+@router.websocket("/ws/{project_name}")
+async def expand_project_websocket(websocket: WebSocket, project_name: str):
+    """
+    WebSocket endpoint for interactive project expansion chat.
+
+    Message protocol:
+
+    Client -> Server:
+    - {"type": "start"} - Start the expansion session
+    - {"type": "message", "content": "..."} - Send user message
+    - {"type": "ping"} - Keep-alive ping
+
+    Server -> Client:
+    - {"type": "text", "content": "..."} - Text chunk from Claude
+    - {"type": "features_created", "count": N, "features": [...]} - Features added
+    - {"type": "expansion_complete", "total_added": N} - Session complete
+    - {"type": "response_done"} - Response complete
+    - {"type": "error", "content": "..."} - Error message
+    - {"type": "pong"} - Keep-alive pong
+    """
+    if not validate_project_name(project_name):
+        await websocket.close(code=4000, reason="Invalid project name")
+        return
+
+    # Look up project directory from registry
+    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
+
+    # Verify project has app_spec.txt
+    spec_path = project_dir / "prompts" / "app_spec.txt"
+    if not spec_path.exists():
+        await websocket.close(code=4004, reason="Project has no spec. Create spec first.")
+        return
+
+    await websocket.accept()
+
+    session: Optional[ExpandChatSession] = None
+
+    try:
+        while True:
+            try:
+                # Receive message from client
+                data = await websocket.receive_text()
+                message = json.loads(data)
+                msg_type = message.get("type")
+
+                if msg_type == "ping":
+                    await websocket.send_json({"type": "pong"})
+                    continue
+
+                elif msg_type == "start":
+                    # Create and start a new expansion session
+                    session = await create_expand_session(project_name, project_dir)
+
+                    # Stream the initial greeting
+                    async for chunk in session.start():
+                        await websocket.send_json(chunk)
+
+                elif msg_type == "message":
+                    # User sent a message
+                    if not session:
+                        session = get_expand_session(project_name)
+                        if not session:
+                            await websocket.send_json({
+                                "type": "error",
+                                "content": "No active session. Send 'start' first."
+                            })
+                            continue
+
+                    user_content = message.get("content", "").strip()
+
+                    # Parse attachments if present
+                    attachments: list[ImageAttachment] = []
+                    raw_attachments = message.get("attachments", [])
+                    if raw_attachments:
+                        try:
+                            for raw_att in raw_attachments:
+                                attachments.append(ImageAttachment(**raw_att))
+                        except (ValidationError, Exception) as e:
+                            logger.warning(f"Invalid attachment data: {e}")
+                            await websocket.send_json({
+                                "type": "error",
+                                "content": f"Invalid attachment: {str(e)}"
+                            })
+                            continue
+
+                    # Allow empty content if attachments are present
+                    if not user_content and not attachments:
+                        await websocket.send_json({
+                            "type": "error",
+                            "content": "Empty message"
+                        })
+                        continue
+
+                    # Stream Claude's response
+                    async for chunk in session.send_message(user_content, attachments if attachments else None):
+                        await websocket.send_json(chunk)
+
+                elif msg_type == "done":
+                    # User is done adding features
+                    if session:
+                        await websocket.send_json({
+                            "type": "expansion_complete",
+                            "total_added": session.get_features_created()
+                        })
+
+                else:
+                    await websocket.send_json({
+                        "type": "error",
+                        "content": f"Unknown message type: {msg_type}"
+                    })
+
+            except json.JSONDecodeError:
+                await websocket.send_json({
+                    "type": "error",
+                    "content": "Invalid JSON"
+                })
+
+    except WebSocketDisconnect:
+        logger.info(f"Expand chat WebSocket disconnected for {project_name}")
+
+    except Exception as e:
+        logger.exception(f"Expand chat WebSocket error for {project_name}")
+        try:
+            await websocket.send_json({
+                "type": "error",
+                "content": f"Server error: {str(e)}"
+            })
+        except Exception:
+            pass
+
+    finally:
+        # Don't remove the session on disconnect - allow resume
+        pass

server/routers/features.py

@@ -13,6 +13,8 @@ from pathlib import Path
 from fastapi import APIRouter, HTTPException
 
 from ..schemas import (
+    FeatureBulkCreate,
+    FeatureBulkCreateResponse,
     FeatureCreate,
     FeatureListResponse,
     FeatureResponse,
@@ -295,3 +297,83 @@ async def skip_feature(project_name: str, feature_id: int):
     except Exception:
         logger.exception("Failed to skip feature")
         raise HTTPException(status_code=500, detail="Failed to skip feature")
+
+
+@router.post("/bulk", response_model=FeatureBulkCreateResponse)
+async def create_features_bulk(project_name: str, bulk: FeatureBulkCreate):
+    """
+    Create multiple features at once.
+
+    Features are assigned sequential priorities starting from:
+    - starting_priority if specified
+    - max(existing priorities) + 1 if not specified
+
+    This is useful for:
+    - Expanding a project with new features via AI
+    - Importing features from external sources
+    - Batch operations
+
+    Returns:
+        {"created": N, "features": [...]}
+    """
+    project_name = validate_project_name(project_name)
+    project_dir = _get_project_path(project_name)
+
+    if not project_dir:
+        raise HTTPException(status_code=404, detail=f"Project '{project_name}' not found in registry")
+
+    if not project_dir.exists():
+        raise HTTPException(status_code=404, detail="Project directory not found")
+
+    if not bulk.features:
+        return FeatureBulkCreateResponse(created=0, features=[])
+
+    _, Feature = _get_db_classes()
+
+    try:
+        with get_db_session(project_dir) as session:
+            # Determine starting priority
+            if bulk.starting_priority is not None:
+                current_priority = bulk.starting_priority
+            else:
+                max_priority_feature = session.query(Feature).order_by(Feature.priority.desc()).first()
+                current_priority = (max_priority_feature.priority + 1) if max_priority_feature else 1
+
+            created_features = []
+
+            for feature_data in bulk.features:
+                db_feature = Feature(
+                    priority=current_priority,
+                    category=feature_data.category,
+                    name=feature_data.name,
+                    description=feature_data.description,
+                    steps=feature_data.steps,
+                    passes=False,
+                )
+                session.add(db_feature)
+                current_priority += 1
+
+            session.commit()
+
+            # Refresh to get IDs and return responses
+            for db_feature in session.query(Feature).order_by(Feature.priority.desc()).limit(len(bulk.features)).all():
+                created_features.insert(0, feature_to_response(db_feature))
+
+            # Re-query to get the actual created features in order
+            created_features = []
+            start_priority = current_priority - len(bulk.features)
+            for db_feature in session.query(Feature).filter(
+                Feature.priority >= start_priority,
+                Feature.priority < current_priority
+            ).order_by(Feature.priority).all():
+                created_features.append(feature_to_response(db_feature))
+
+            return FeatureBulkCreateResponse(
+                created=len(created_features),
+                features=created_features
+            )
+    except HTTPException:
+        raise
+    except Exception:
+        logger.exception("Failed to bulk create features")
+        raise HTTPException(status_code=500, detail="Failed to bulk create features")

server/schemas.py

@@ -96,6 +96,18 @@ class FeatureListResponse(BaseModel):
     done: list[FeatureResponse]
 
 
+class FeatureBulkCreate(BaseModel):
+    """Request schema for bulk creating features."""
+    features: list[FeatureCreate]
+    starting_priority: int | None = None  # If None, appends after max priority
+
+
+class FeatureBulkCreateResponse(BaseModel):
+    """Response for bulk feature creation."""
+    created: int
+    features: list[FeatureResponse]
+
+
 # ============================================================================
 # Agent Schemas
 # ============================================================================

server/services/expand_chat_session.py

@@ -0,0 +1,444 @@
+"""
+Expand Chat Session
+===================
+
+Manages interactive project expansion conversation with Claude.
+Uses the expand-project.md skill to help users add features to existing projects.
+"""
+
+import json
+import logging
+import re
+import shutil
+import threading
+from datetime import datetime
+from pathlib import Path
+from typing import AsyncGenerator, Optional
+
+from claude_agent_sdk import ClaudeAgentOptions, ClaudeSDKClient
+
+from ..schemas import ImageAttachment
+
+logger = logging.getLogger(__name__)
+
+
+async def _make_multimodal_message(content_blocks: list[dict]) -> AsyncGenerator[dict, None]:
+    """
+    Create an async generator that yields a properly formatted multimodal message.
+    """
+    yield {
+        "type": "user",
+        "message": {"role": "user", "content": content_blocks},
+        "parent_tool_use_id": None,
+        "session_id": "default",
+    }
+
+
+# Root directory of the project
+ROOT_DIR = Path(__file__).parent.parent.parent
+
+
+class ExpandChatSession:
+    """
+    Manages a project expansion conversation.
+
+    Unlike SpecChatSession which writes spec files, this session:
+    1. Reads existing app_spec.txt for context
+    2. Parses feature definitions from Claude's output
+    3. Creates features via REST API
+    4. Tracks which features were created during the session
+    """
+
+    def __init__(self, project_name: str, project_dir: Path):
+        """
+        Initialize the session.
+
+        Args:
+            project_name: Name of the project being expanded
+            project_dir: Absolute path to the project directory
+        """
+        self.project_name = project_name
+        self.project_dir = project_dir
+        self.client: Optional[ClaudeSDKClient] = None
+        self.messages: list[dict] = []
+        self.complete: bool = False
+        self.created_at = datetime.now()
+        self._conversation_id: Optional[str] = None
+        self._client_entered: bool = False
+        self.features_created: int = 0
+        self.created_feature_ids: list[int] = []
+
+    async def close(self) -> None:
+        """Clean up resources and close the Claude client."""
+        if self.client and self._client_entered:
+            try:
+                await self.client.__aexit__(None, None, None)
+            except Exception as e:
+                logger.warning(f"Error closing Claude client: {e}")
+            finally:
+                self._client_entered = False
+                self.client = None
+
+    async def start(self) -> AsyncGenerator[dict, None]:
+        """
+        Initialize session and get initial greeting from Claude.
+
+        Yields message chunks as they stream in.
+        """
+        # Load the expand-project skill
+        skill_path = ROOT_DIR / ".claude" / "commands" / "expand-project.md"
+
+        if not skill_path.exists():
+            yield {
+                "type": "error",
+                "content": f"Expand project skill not found at {skill_path}"
+            }
+            return
+
+        # Verify project has existing spec
+        spec_path = self.project_dir / "prompts" / "app_spec.txt"
+        if not spec_path.exists():
+            yield {
+                "type": "error",
+                "content": "Project has no app_spec.txt. Please create it first using spec creation."
+            }
+            return
+
+        try:
+            skill_content = skill_path.read_text(encoding="utf-8")
+        except UnicodeDecodeError:
+            skill_content = skill_path.read_text(encoding="utf-8", errors="replace")
+
+        # Create security settings file
+        security_settings = {
+            "sandbox": {"enabled": False},
+            "permissions": {
+                "defaultMode": "acceptEdits",
+                "allow": [
+                    "Read(./**)",
+                    "Glob(./**)",
+                ],
+            },
+        }
+        settings_file = self.project_dir / ".claude_settings.json"
+        with open(settings_file, "w") as f:
+            json.dump(security_settings, f, indent=2)
+
+        # Replace $ARGUMENTS with absolute project path
+        project_path = str(self.project_dir.resolve())
+        system_prompt = skill_content.replace("$ARGUMENTS", project_path)
+
+        # Create Claude SDK client
+        system_cli = shutil.which("claude")
+        try:
+            self.client = ClaudeSDKClient(
+                options=ClaudeAgentOptions(
+                    model="claude-opus-4-5-20251101",
+                    cli_path=system_cli,
+                    system_prompt=system_prompt,
+                    allowed_tools=[
+                        "Read",
+                        "Glob",
+                    ],
+                    permission_mode="acceptEdits",
+                    max_turns=100,
+                    cwd=str(self.project_dir.resolve()),
+                    settings=str(settings_file.resolve()),
+                )
+            )
+            await self.client.__aenter__()
+            self._client_entered = True
+        except Exception as e:
+            logger.exception("Failed to create Claude client")
+            yield {
+                "type": "error",
+                "content": f"Failed to initialize Claude: {str(e)}"
+            }
+            return
+
+        # Start the conversation
+        try:
+            async for chunk in self._query_claude("Begin the project expansion process."):
+                yield chunk
+            yield {"type": "response_done"}
+        except Exception as e:
+            logger.exception("Failed to start expand chat")
+            yield {
+                "type": "error",
+                "content": f"Failed to start conversation: {str(e)}"
+            }
+
+    async def send_message(
+        self,
+        user_message: str,
+        attachments: list[ImageAttachment] | None = None
+    ) -> AsyncGenerator[dict, None]:
+        """
+        Send user message and stream Claude's response.
+
+        Args:
+            user_message: The user's response
+            attachments: Optional list of image attachments
+
+        Yields:
+            Message chunks of various types:
+            - {"type": "text", "content": str}
+            - {"type": "features_created", "count": N, "features": [...]}
+            - {"type": "expansion_complete", "total_added": N}
+            - {"type": "error", "content": str}
+        """
+        if not self.client:
+            yield {
+                "type": "error",
+                "content": "Session not initialized. Call start() first."
+            }
+            return
+
+        # Store the user message
+        self.messages.append({
+            "role": "user",
+            "content": user_message,
+            "has_attachments": bool(attachments),
+            "timestamp": datetime.now().isoformat()
+        })
+
+        try:
+            async for chunk in self._query_claude(user_message, attachments):
+                yield chunk
+            yield {"type": "response_done"}
+        except Exception as e:
+            logger.exception("Error during Claude query")
+            yield {
+                "type": "error",
+                "content": f"Error: {str(e)}"
+            }
+
+    async def _query_claude(
+        self,
+        message: str,
+        attachments: list[ImageAttachment] | None = None
+    ) -> AsyncGenerator[dict, None]:
+        """
+        Internal method to query Claude and stream responses.
+
+        Handles text responses and detects feature creation blocks.
+        """
+        if not self.client:
+            return
+
+        # Build the message content
+        if attachments and len(attachments) > 0:
+            content_blocks = []
+            if message:
+                content_blocks.append({"type": "text", "text": message})
+            for att in attachments:
+                content_blocks.append({
+                    "type": "image",
+                    "source": {
+                        "type": "base64",
+                        "media_type": att.mimeType,
+                        "data": att.base64Data,
+                    }
+                })
+            await self.client.query(_make_multimodal_message(content_blocks))
+            logger.info(f"Sent multimodal message with {len(attachments)} image(s)")
+        else:
+            await self.client.query(message)
+
+        # Accumulate full response to detect feature blocks
+        full_response = ""
+
+        # Stream the response
+        async for msg in self.client.receive_response():
+            msg_type = type(msg).__name__
+
+            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"):
+                        text = block.text
+                        if text:
+                            full_response += text
+                            yield {"type": "text", "content": text}
+
+                            self.messages.append({
+                                "role": "assistant",
+                                "content": text,
+                                "timestamp": datetime.now().isoformat()
+                            })
+
+        # Check for feature creation block in full response
+        features_match = re.search(
+            r'<features_to_create>\s*(\[[\s\S]*?\])\s*</features_to_create>',
+            full_response
+        )
+
+        if features_match:
+            try:
+                features_json = features_match.group(1)
+                features_data = json.loads(features_json)
+
+                if features_data and isinstance(features_data, list):
+                    # Create features via REST API
+                    created = await self._create_features_bulk(features_data)
+
+                    if created:
+                        self.features_created += len(created)
+                        self.created_feature_ids.extend([f["id"] for f in created])
+
+                        yield {
+                            "type": "features_created",
+                            "count": len(created),
+                            "features": created
+                        }
+
+                        logger.info(f"Created {len(created)} features for {self.project_name}")
+            except json.JSONDecodeError as e:
+                logger.error(f"Failed to parse features JSON: {e}")
+                yield {
+                    "type": "error",
+                    "content": f"Failed to parse feature definitions: {str(e)}"
+                }
+            except Exception as e:
+                logger.exception("Failed to create features")
+                yield {
+                    "type": "error",
+                    "content": f"Failed to create features: {str(e)}"
+                }
+
+    async def _create_features_bulk(self, features: list[dict]) -> list[dict]:
+        """
+        Create features directly in the database.
+
+        Args:
+            features: List of feature dictionaries with category, name, description, steps
+
+        Returns:
+            List of created feature dictionaries with IDs
+        """
+        # Import database classes
+        import sys
+        root = Path(__file__).parent.parent.parent
+        if str(root) not in sys.path:
+            sys.path.insert(0, str(root))
+
+        from api.database import Feature, create_database
+
+        # Get database session
+        _, SessionLocal = create_database(self.project_dir)
+        session = SessionLocal()
+
+        try:
+            # Determine starting priority
+            max_priority_feature = session.query(Feature).order_by(Feature.priority.desc()).first()
+            current_priority = (max_priority_feature.priority + 1) if max_priority_feature else 1
+
+            created_features = []
+
+            for f in features:
+                db_feature = Feature(
+                    priority=current_priority,
+                    category=f.get("category", "functional"),
+                    name=f.get("name", "Unnamed feature"),
+                    description=f.get("description", ""),
+                    steps=f.get("steps", []),
+                    passes=False,
+                )
+                session.add(db_feature)
+                current_priority += 1
+
+            session.commit()
+
+            # Re-query to get the created features with IDs
+            start_priority = current_priority - len(features)
+            for db_feature in session.query(Feature).filter(
+                Feature.priority >= start_priority,
+                Feature.priority < current_priority
+            ).order_by(Feature.priority).all():
+                created_features.append({
+                    "id": db_feature.id,
+                    "name": db_feature.name,
+                    "category": db_feature.category,
+                })
+
+            return created_features
+
+        finally:
+            session.close()
+
+    def get_features_created(self) -> int:
+        """Get the total number of features created in this session."""
+        return self.features_created
+
+    def is_complete(self) -> bool:
+        """Check if expansion session is complete."""
+        return self.complete
+
+    def get_messages(self) -> list[dict]:
+        """Get all messages in the conversation."""
+        return self.messages.copy()
+
+
+# Session registry with thread safety
+_expand_sessions: dict[str, ExpandChatSession] = {}
+_expand_sessions_lock = threading.Lock()
+
+
+def get_expand_session(project_name: str) -> Optional[ExpandChatSession]:
+    """Get an existing expansion session for a project."""
+    with _expand_sessions_lock:
+        return _expand_sessions.get(project_name)
+
+
+async def create_expand_session(project_name: str, project_dir: Path) -> ExpandChatSession:
+    """Create a new expansion session for a project, closing any existing one."""
+    old_session: Optional[ExpandChatSession] = None
+
+    with _expand_sessions_lock:
+        old_session = _expand_sessions.pop(project_name, None)
+        session = ExpandChatSession(project_name, project_dir)
+        _expand_sessions[project_name] = session
+
+    if old_session:
+        try:
+            await old_session.close()
+        except Exception as e:
+            logger.warning(f"Error closing old expand session for {project_name}: {e}")
+
+    return session
+
+
+async def remove_expand_session(project_name: str) -> None:
+    """Remove and close an expansion session."""
+    session: Optional[ExpandChatSession] = None
+
+    with _expand_sessions_lock:
+        session = _expand_sessions.pop(project_name, None)
+
+    if session:
+        try:
+            await session.close()
+        except Exception as e:
+            logger.warning(f"Error closing expand session for {project_name}: {e}")
+
+
+def list_expand_sessions() -> list[str]:
+    """List all active expansion session project names."""
+    with _expand_sessions_lock:
+        return list(_expand_sessions.keys())
+
+
+async def cleanup_all_expand_sessions() -> None:
+    """Close all active expansion sessions. Called on server shutdown."""
+    sessions_to_close: list[ExpandChatSession] = []
+
+    with _expand_sessions_lock:
+        sessions_to_close = list(_expand_sessions.values())
+        _expand_sessions.clear()
+
+    for session in sessions_to_close:
+        try:
+            await session.close()
+        except Exception as e:
+            logger.warning(f"Error closing expand session {session.project_name}: {e}")