import asyncio import json import logging import uuid from concurrent.futures import ThreadPoolExecutor from typing import AsyncGenerator, Dict from fastapi import APIRouter, Depends, HTTPException from fastapi.responses import StreamingResponse from pydantic import BaseModel, field_validator from agent.events import AgentEventEmitter, BUILD_TOPOLOGY from agent.graph import run_build from agent.scaffold import coerce_framework from auth.dependencies import get_current_user from db.database import SessionLocal from db.models import Project, User logger = logging.getLogger(__name__) router = APIRouter(prefix="/generate", tags=["generate"]) _executor = ThreadPoolExecutor(max_workers=4) _job_queues: Dict[str, asyncio.Queue] = {} # NOTE: single-consumer per job; redesign planned in a later task. _DONE = "__DONE__" _ERR = "__ERR__:" def shutdown_executor() -> None: """Shut down the thread pool gracefully (call from app lifespan on shutdown).""" _executor.shutdown(wait=False) class GenerateRequest(BaseModel): brief: str session_id: str | None = None # ponytail: kept for backward compat, ignored for ownership premium: bool = False # honored ONLY for SUPER_ADMIN (server-side gate below) framework: str = "react" # "react" (default) | "static" | "nextjs"; coerced + validated @field_validator("framework", mode="before") @classmethod def _coerce_framework(cls, v): return coerce_framework(v) def _resolve_premium(requested: bool, user) -> bool: """Server-side gate: premium production-grade builds are honored ONLY for SUPER_ADMIN. ponytail: never trust the client flag — a non-admin asking for premium gets standard. """ return bool(requested) and getattr(user, "role", "") == "SUPER_ADMIN" def _persist_project(user_id: str | None, title: str, brief: str, files: dict, framework: str = "react") -> str | None: """Save a Project row; returns the new UUID string or None on failure.""" if not user_id: return None try: with SessionLocal() as db: row = Project(session_id="", user_id=uuid.UUID(user_id), title=title, brief=brief, files=files, framework=framework) db.add(row) db.commit() db.refresh(row) return str(row.id) except Exception as exc: # noqa: BLE001 ponytail: DB hiccup must not break the build logger.error("DB persist failed (non-fatal): %s", exc) return None def _run_job(job_id, queue, loop, brief, user_id=None, premium=False, framework="react"): emitter = AgentEventEmitter(loop, queue, BUILD_TOPOLOGY) try: out = run_build(brief, emitter=emitter, premium=premium, framework=framework) fw = out.get("framework", framework) project_id = _persist_project(user_id, out["plan"].name, brief, out["files"], fw) payload: dict = {"type": "result", "name": out["plan"].name, "files": out["files"], "framework": fw} if project_id: payload["project_id"] = project_id asyncio.run_coroutine_threadsafe(queue.put(json.dumps(payload)), loop) asyncio.run_coroutine_threadsafe(queue.put(_DONE), loop) except Exception as exc: # noqa: BLE001 logger.error("Generate job %s failed: %s", job_id, exc) asyncio.run_coroutine_threadsafe(queue.put(f"{_ERR}{exc}"), loop) @router.post("/start") async def generate_start( req: GenerateRequest, current_user: User = Depends(get_current_user), ) -> dict: if not req.brief.strip(): raise HTTPException(status_code=400, detail="brief is required") job_id = str(uuid.uuid4()) loop = asyncio.get_running_loop() queue: asyncio.Queue = asyncio.Queue() _job_queues[job_id] = queue user_id = str(current_user.id) premium = _resolve_premium(req.premium, current_user) _executor.submit(_run_job, job_id, queue, loop, req.brief, user_id, premium, req.framework) return {"job_id": job_id} async def _sse(job_id, queue) -> AsyncGenerator[str, None]: try: while True: try: msg = await asyncio.wait_for(queue.get(), timeout=15.0) except asyncio.TimeoutError: yield ": keepalive\n\n" continue if msg == _DONE: yield f"data: {json.dumps({'type': 'done'})}\n\n" break if isinstance(msg, str) and msg.startswith(_ERR): yield f"data: {json.dumps({'type': 'error', 'message': msg[len(_ERR):]})}\n\n" break yield f"data: {msg}\n\n" finally: _job_queues.pop(job_id, None) # ponytail: events stream is open (no auth) — job_id UUID is already a nonce @router.get("/events/{job_id}") async def generate_events(job_id: str) -> StreamingResponse: queue = _job_queues.get(job_id) if not queue: raise HTTPException(status_code=404, detail="Job not found or already completed") return StreamingResponse( _sse(job_id, queue), media_type="text/event-stream", headers={ "Cache-Control": "no-cache", "X-Accel-Buffering": "no", "Connection": "keep-alive", }, )