""" Spec Creation Router ==================== WebSocket and REST endpoints for interactive spec creation with Claude. """ import asyncio import json import logging import re from pathlib import Path from typing import Any, Optional from fastapi import APIRouter, WebSocket, WebSocketDisconnect, HTTPException from pydantic import BaseModel from ..services.spec_chat_session import ( SpecChatSession, get_session, create_session, remove_session, list_sessions, ) logger = logging.getLogger(__name__) router = APIRouter(prefix="/api/spec", tags=["spec-creation"]) # Root directory ROOT_DIR = Path(__file__).parent.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)) # ============================================================================ # REST Endpoints # ============================================================================ class SpecSessionStatus(BaseModel): """Status of a spec creation session.""" project_name: str is_active: bool is_complete: bool message_count: int @router.get("/sessions", response_model=list[str]) async def list_spec_sessions(): """List all active spec creation sessions.""" return list_sessions() @router.get("/sessions/{project_name}", response_model=SpecSessionStatus) async def get_session_status(project_name: str): """Get status of a spec creation session.""" if not validate_project_name(project_name): raise HTTPException(status_code=400, detail="Invalid project name") session = get_session(project_name) if not session: raise HTTPException(status_code=404, detail="No active session for this project") return SpecSessionStatus( project_name=project_name, is_active=True, is_complete=session.is_complete(), message_count=len(session.get_messages()), ) @router.delete("/sessions/{project_name}") async def cancel_session(project_name: str): """Cancel and remove a spec creation session.""" if not validate_project_name(project_name): raise HTTPException(status_code=400, detail="Invalid project name") session = get_session(project_name) if not session: raise HTTPException(status_code=404, detail="No active session for this project") await remove_session(project_name) return {"success": True, "message": "Session cancelled"} # ============================================================================ # WebSocket Endpoint # ============================================================================ @router.websocket("/ws/{project_name}") async def spec_chat_websocket(websocket: WebSocket, project_name: str): """ WebSocket endpoint for interactive spec creation chat. Message protocol: Client -> Server: - {"type": "start"} - Start the spec creation session - {"type": "message", "content": "..."} - Send user message - {"type": "answer", "answers": {...}, "tool_id": "..."} - Answer structured question - {"type": "ping"} - Keep-alive ping Server -> Client: - {"type": "text", "content": "..."} - Text chunk from Claude - {"type": "question", "questions": [...], "tool_id": "..."} - Structured question - {"type": "spec_complete", "path": "..."} - Spec file created - {"type": "file_written", "path": "..."} - Other file written - {"type": "complete"} - Session 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 await websocket.accept() session: Optional[SpecChatSession] = 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 session session = await create_session(project_name) # Stream the initial greeting async for chunk in session.start(): await websocket.send_json(chunk) # Check for completion if chunk.get("type") == "spec_complete": await websocket.send_json({"type": "complete"}) elif msg_type == "message": # User sent a message if not session: session = get_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() if not user_content: await websocket.send_json({ "type": "error", "content": "Empty message" }) continue # Stream Claude's response async for chunk in session.send_message(user_content): await websocket.send_json(chunk) # Check for completion if chunk.get("type") == "spec_complete": await websocket.send_json({"type": "complete"}) elif msg_type == "answer": # User answered a structured question if not session: session = get_session(project_name) if not session: await websocket.send_json({ "type": "error", "content": "No active session" }) continue # Format the answers as a natural response answers = message.get("answers", {}) if isinstance(answers, dict): # Convert structured answers to a message response_parts = [] for question_idx, answer_value in answers.items(): if isinstance(answer_value, list): response_parts.append(", ".join(answer_value)) else: response_parts.append(str(answer_value)) user_response = "; ".join(response_parts) if response_parts else "OK" else: user_response = str(answers) # Stream Claude's response async for chunk in session.send_message(user_response): await websocket.send_json(chunk) if chunk.get("type") == "spec_complete": await websocket.send_json({"type": "complete"}) 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"Spec chat WebSocket disconnected for {project_name}") except Exception as e: logger.exception(f"Spec 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