"""Canonical POST /api/projects/{id}/chat — single-run lifecycle + active-run rules.""" from __future__ import annotations import json from typing import Any, Optional from uuid import UUID import structlog from pydantic import BaseModel, model_validator from sqlalchemy.ext.asyncio import AsyncSession from src.core.exceptions import ( NotFoundError, RunBusyError, RunContinuationError, ValidationError, ) from src.domain.auth.models import User from src.config import get_settings from src.domain.codegen.prompt_wrapper import merge_plan_user_prompt from src.domain.codegen.repository import CodegenRunRepository from src.domain.codegen.run_lifecycle import ( CONTINUABLE_STATUSES, RUN_BUSY_STATUSES, USER_INPUT_BLOCKING_STATUSES, ) from src.domain.codegen.run_chat_repository import RunChatMessageRepository from src.domain.codegen.service import CodegenService from src.domain.instructions.repository import InstructionRepository, PromptRepository, PromptSessionRepository from src.domain.projects.paths import get_project_repo_path from src.domain.projects.repository import ProjectRepository, ProjectSettingsRepository from src.domain.prompt_media.repository import PromptMediaRepository from src.domain.visual_edit.deterministic import validate_visual_edit_mapping logger = structlog.get_logger() class ElementContext(BaseModel): tag: Optional[str] = None text_snippet: Optional[str] = None xpath: Optional[str] = None selector: Optional[str] = None element_hint: Optional[str] = None class ChatRequest(BaseModel): """v2 chat contract (attachments = prompt_media asset ids).""" message: str = "" attachments: Optional[list[UUID]] = None context: Optional[dict[str, Any]] = None session_id: Optional[UUID] = None run_id: Optional[UUID] = None clarification_answers: Optional[dict[str, Any]] = None voice_transcript: Optional[str] = None approve_plan: bool = False @model_validator(mode="after") def _require_payload(self) -> "ChatRequest": text = (self.message or "").strip() voice = (self.voice_transcript or "").strip() ids = self.attachments or [] structured = self.clarification_answers is not None if self.run_id is not None: if text or voice or ids or structured or self.approve_plan: return self raise ValueError( "Continuation requires message, voice_transcript, attachments, clarification_answers, or approve_plan" ) mode = (self.context or {}).get("mode") if mode == "visual_edit": sel = (self.context or {}).get("selected_elements") if isinstance(sel, list) and len(sel) > 0: return self if not text and not voice and not ids: raise ValueError("Provide message, voice_transcript, and/or attachments") return self class ChatResponse(BaseModel): message_id: UUID run_id: UUID session_id: UUID status: str = "queued" def _flatten_clarification_answers(raw: Optional[dict[str, Any]]) -> dict[str, str]: if not raw: return {} out: dict[str, str] = {} for k, v in raw.items(): key = str(k) if isinstance(v, dict): sel = v.get("selected") ot = v.get("other_text") if sel is not None and str(sel).strip(): out[key] = str(sel).strip() elif ot is not None and str(ot).strip(): out[key] = str(ot).strip() else: out[key] = json.dumps(v, ensure_ascii=False) else: out[key] = str(v) return out def _ctx_from_chat(body: ChatRequest) -> tuple[dict[str, Any], Optional[str], Optional[str], Optional[ElementContext]]: ctx = dict(body.context) if body.context else {} route = ctx.pop("route", None) or ctx.pop("selected_route", None) viewport = ctx.pop("viewport", None) raw_el = ctx.pop("selected_elements", None) element_context: Optional[ElementContext] = None mode = ctx.get("mode") if raw_el and isinstance(raw_el, list) and len(raw_el) > 0 and mode == "visual_edit": ctx["selected_elements"] = raw_el elif raw_el and isinstance(raw_el, list) and len(raw_el) > 0: first = raw_el[0] if isinstance(first, dict): try: element_context = ElementContext.model_validate(first) except Exception: element_context = None if body.attachments: ctx["attachment_ids"] = [str(x) for x in body.attachments] vt = (body.voice_transcript or "").strip() if vt: ctx["voice_transcript"] = vt return ctx, route, viewport, element_context async def handle_project_chat( db: AsyncSession, project_id: UUID, user: User, body: ChatRequest, *, codegen_service: CodegenService, ) -> ChatResponse: project_repo = ProjectRepository(db) session_repo = PromptSessionRepository(db) prompt_repo = PromptRepository(db) codegen_repo = CodegenRunRepository(db) instruction_repo = InstructionRepository(db) run_chat_repo = RunChatMessageRepository(db) await codegen_repo.lock_project_for_run_creation(project_id) from src.domain.projects.discovery_repository import DiscoveryRepository proj = await DiscoveryRepository(db).get(project_id) if not proj: raise NotFoundError("Project not found") from src.domain.projects.discovery_service import require_discovery_complete require_discovery_complete(proj) understanding = None active = await codegen_repo.get_active_reserving_run(project_id) if body.run_id is None: if active: await codegen_service.supersede_active_run_for_new_chat(project_id) else: run = await codegen_repo.get_by_id(body.run_id) if not run or run.project_id != project_id: raise RunContinuationError("Run not found for this project.") if run.status in RUN_BUSY_STATUSES: raise RunBusyError() if run.status not in CONTINUABLE_STATUSES: raise RunContinuationError( f"This run cannot accept chat continuation (status: {run.status})." ) ctx_merged, selected_route, viewport, element_context = _ctx_from_chat(body) ctx_merged.pop("project_understanding", None) if understanding: ctx_merged["project_understanding"] = understanding.model_dump() sid = body.session_id or proj.active_session_id if not sid: s = await session_repo.create(project_id, user.id, None) sid = s.id await project_repo.update(project_id, active_session_id=sid) else: sess = await session_repo.get_by_id(sid) if not sess or sess.project_id != project_id: raise NotFoundError("Session not found") # ── Continuation paths (explicit run_id) ───────────────────────────── if body.run_id is not None: run = await codegen_repo.get_by_id(body.run_id) assert run is not None if run.status in USER_INPUT_BLOCKING_STATUSES: vals = _flatten_clarification_answers(body.clarification_answers) if not vals: raise ValidationError( "clarification_answers (or structured values) are required to continue this run." ) display = (body.message or "").strip() if not display: display = f"[Configuration: {len(vals)} value(s)]" p = await prompt_repo.create(sid, display, user.id, ctx_merged) await run_chat_repo.create( project_id=project_id, session_id=sid, run_id=run.id, role="user", content=display, message_type="message", prompt_id=p.id, ) run.prompt_id = p.id await db.flush() updated = await codegen_service.supply_user_inputs( run.id, project_id, vals, user.id ) await db.commit() from src.workers.enqueue import enqueue_codegen_run await enqueue_codegen_run(str(updated.id)) logger.info( "chat_continued_supply", run_id=str(updated.id), message_id=str(p.id), ) return ChatResponse( message_id=p.id, run_id=updated.id, session_id=sid, status="queued", ) if run.status == "plan_ready": display = (body.message or "").strip() or "[Approve plan]" p = await prompt_repo.create(sid, display, user.id, ctx_merged) await run_chat_repo.create( project_id=project_id, session_id=sid, run_id=run.id, role="user", content=display, message_type="message", prompt_id=p.id, ) pl = dict(run.plan or {}) ctx0 = dict(pl.get("context") or {}) if body.message.strip(): ctx0["approval_notes"] = body.message.strip() pl["context"] = ctx0 await codegen_repo.update_status(run.id, "plan_ready", plan=pl) await db.flush() run.prompt_id = p.id await db.flush() approved = await codegen_service.execute_plan(run.id, project_id, user.id) await db.commit() from src.workers.enqueue import enqueue_codegen_run await enqueue_codegen_run(str(approved.id)) logger.info("chat_continued_plan", run_id=str(approved.id), message_id=str(p.id)) return ChatResponse( message_id=p.id, run_id=approved.id, session_id=sid, status="queued", ) if run.status == "clarification_required": if not body.clarification_answers and not (body.message or "").strip(): raise ValidationError( "Provide clarification_answers and/or a follow-up message.", ) display = (body.message or "").strip() or "[Clarification reply]" p = await prompt_repo.create(sid, display, user.id, ctx_merged) await run_chat_repo.create( project_id=project_id, session_id=sid, run_id=run.id, role="user", content=display, message_type="message", prompt_id=p.id, ) pl = dict(run.plan or {}) ctx0 = dict(pl.get("context") or {}) history = list(ctx0.get("clarification_history") or []) round_entry: dict[str, Any] = {"answers": body.clarification_answers or {}} if body.message.strip(): round_entry["freeform_followup"] = body.message.strip() history.append(round_entry) ctx0["clarification_history"] = history pl["context"] = ctx0 pl["user_prompt"] = merge_plan_user_prompt( str(pl.get("original_user_prompt") or ""), ctx0, ) run.prompt_id = p.id await db.flush() updated = await codegen_repo.requeue_with_plan(run.id, pl) if not updated: raise RunContinuationError() await db.commit() from src.workers.enqueue import enqueue_codegen_run await enqueue_codegen_run(str(updated.id)) return ChatResponse( message_id=p.id, run_id=updated.id, session_id=sid, status="queued", ) if run.status == "queued": display = (body.message or "").strip() if not display and body.attachments: display = f"[{len(body.attachments)} attachment(s)]" if not display and body.clarification_answers: display = "[Follow-up]" p = await prompt_repo.create(sid, display, user.id, ctx_merged) await run_chat_repo.create( project_id=project_id, session_id=sid, run_id=run.id, role="user", content=display, message_type="message", prompt_id=p.id, ) pl = dict(run.plan or {}) ctx0 = dict(pl.get("context") or {}) new_part = (body.message or "").strip() if new_part: apx = (ctx0.get("user_message_appendix") or "").strip() ctx0["user_message_appendix"] = ( f"{apx}\n\n{new_part}".strip() if apx else new_part ) if body.attachments: ctx0["attachment_ids"] = [str(x) for x in body.attachments] if body.voice_transcript: ctx0["voice_transcript"] = body.voice_transcript.strip() pl["context"] = ctx0 pl["user_prompt"] = merge_plan_user_prompt( str(pl.get("original_user_prompt") or ""), ctx0, ) await codegen_repo.update_status(run.id, "queued", plan=pl) await codegen_repo.get_by_id(run.id) await db.flush() run.prompt_id = p.id await db.commit() logger.info( "chat_merged_queued", run_id=str(run.id), message_id=str(p.id), ) return ChatResponse( message_id=p.id, run_id=run.id, session_id=sid, status="queued", ) raise RunContinuationError(f"Unsupported continuation status: {run.status}") # ── New run ───────────────────────────────────────────────────────── media_repo = PromptMediaRepository(db) att = body.attachments or [] if att: assets = await media_repo.list_by_ids_for_project(att, project_id) if len(assets) != len(set(att)): raise ValidationError("Unknown attachment id") if any(a.prompt_id is not None for a in assets): raise ValidationError("Attachment already linked to a prompt") vt = (body.voice_transcript or "").strip() display = (body.message or "").strip() if not display: hint: list[str] = [] if ctx_merged.get("mode") == "visual_edit": hint.append("[Visual edit]") if vt: hint.append("[Voice input]") if att: hint.append(f"[{len(att)} attachment(s)]") display = " ".join(hint) p = await prompt_repo.create(sid, display, user.id, ctx_merged) if att: linked = await media_repo.attach_to_prompt(p.id, att, project_id) if linked != len(att): raise ValidationError("Could not link one or more attachments") ve_mode = ctx_merged.get("mode") == "visual_edit" sel_ve = ctx_merged.get("selected_elements") if ve_mode else None if ve_mode and isinstance(sel_ve, list) and len(sel_ve) > 0: st = get_settings() rpath = get_project_repo_path(st.repo_workspace, proj.workspace_id, project_id) verr = validate_visual_edit_mapping(rpath, sel_ve) if verr: raise ValidationError( "Could not deterministically map the selected element to an editable source location.", error="mapping_failed", ) run_mode = "visual_edit" if ve_mode else "agent" run = await codegen_repo.create( project_id=project_id, prompt_id=p.id, created_by=user.id, run_mode=run_mode, selected_elements=list(sel_ve) if ve_mode and isinstance(sel_ve, list) else None, ) await run_chat_repo.create( project_id=project_id, session_id=sid, run_id=run.id, role="user", content=display, message_type="message", prompt_id=p.id, ) instruction_content = await instruction_repo.get_current_content(project_id) settings_repo = ProjectSettingsRepository(db) persona = await settings_repo.get_persona(project_id) existing_prompts = await prompt_repo.list_for_session(sid) prompt_history = [ep.content for ep in existing_prompts if ep.id != p.id][-10:] from src.domain.codegen.platform_backend import build_platform_backend_capabilities from src.domain.codegen.prompt_wrapper import build_system_prompt from src.domain.manifest.repository import ManifestRunRepository from src.domain.manifest.service import ManifestService manifest = await ManifestService( db=db, run_repo=ManifestRunRepository(db), ).get_manifest(project_id, compact=False) platform_backend = build_platform_backend_capabilities(manifest) form_policy = await settings_repo.get_form_policy(project_id) system_prompt = build_system_prompt( project_type=proj.type, instruction_content=instruction_content, selected_route=selected_route, viewport=viewport, element_context=element_context.model_dump() if element_context else None, persona=persona if persona else None, prompt_history=prompt_history if prompt_history else None, selected_elements=list(sel_ve) if ve_mode and isinstance(sel_ve, list) else None, platform_backend=platform_backend, form_policy=form_policy, ) plan_ctx: dict[str, Any] = { "selected_route": selected_route, "viewport": viewport, "project_type": proj.type, "element_context": element_context.model_dump() if element_context else None, "run_mode": run_mode, } if ve_mode and isinstance(sel_ve, list): plan_ctx["mode"] = "visual_edit" plan_ctx["visual_edit_elements"] = list(sel_ve) if att: plan_ctx["attachment_ids"] = [str(x) for x in att] if vt: plan_ctx["voice_transcript"] = vt orig = (body.message or "").strip() plan = { "system_prompt": system_prompt, "user_prompt": merge_plan_user_prompt(orig, plan_ctx), "original_user_prompt": orig, "context": plan_ctx, } await codegen_repo.update_status(run.id, "queued", plan=plan) discovery = getattr(proj, "discovery_json", None) if discovery and discovery.get("status") == "completed" and not discovery.get("initial_run_id"): proj.discovery_json = {**discovery, "initial_run_id": str(run.id)} await db.commit() from src.workers.enqueue import enqueue_codegen_run await enqueue_codegen_run(str(run.id)) logger.info("chat_new_run", run_id=str(run.id), message_id=str(p.id)) return ChatResponse( message_id=p.id, run_id=run.id, session_id=sid, status="queued", )