src/agents/state_store.py

"""AnalysisStateStore — read/write the per-analysis session state.

The orchestrator gate + Help skill read `AnalysisState` (the locked contract in
`gate.py`) every turn; the Problem Statement skill writes `problem_validated`. The
row shares its id with the chat `rooms` row — one session = one analysis = one
conversation (`analysis_id == room_id`).

Mirrors `PostgresReportInputStore`: each call opens its own `AsyncSession`.
"""

from __future__ import annotations

from sqlalchemy.dialects.postgresql import insert

from src.agents.gate import AnalysisState
from src.db.postgres.connection import AsyncSessionLocal
from src.db.postgres.models import AnalysisStateRow
from src.middlewares.logging import get_logger

logger = get_logger("analysis_state_store")


def _row_to_state(row: AnalysisStateRow) -> AnalysisState:
    """Map a DB row to the frozen `AnalysisState` contract."""
    return AnalysisState(
        id=row.id,
        analysis_title=row.analysis_title,
        problem_statement=row.problem_statement,
        problem_validated=row.problem_validated,
        user_id=row.user_id,
        report_id=row.report_id,
        created_at=row.created_at,
        updated_at=row.updated_at,
    )


class AnalysisStateStore:
    """Read/write the dedorch `analysis` table, keyed by the shared session id."""

    async def get(self, analysis_id: str) -> AnalysisState | None:
        async with AsyncSessionLocal() as session:
            row = await session.get(AnalysisStateRow, analysis_id)
            return _row_to_state(row) if row is not None else None

    async def ensure(
        self,
        analysis_id: str,
        user_id: str,
        analysis_title: str = "New analysis",
    ) -> AnalysisState:
        """Get-or-create the state row for a session (idempotent, race-safe).

        Sessions born from `/room/create` have no `analysis_states` row; without
        one the gate redirect-loops and `problem_statement` / `report_id` writes
        silently no-op. Called per turn (analysis_id == room_id) so any session is
        gate-ready. `INSERT ... ON CONFLICT DO NOTHING` makes concurrent first
        turns safe; the row is then read back. Legacy rows created this way carry
        no source bindings — binding scoping fail-opens to the whole catalog.
        """
        async with AsyncSessionLocal() as session:
            stmt = (
                insert(AnalysisStateRow)
                .values(
                    id=analysis_id,
                    user_id=user_id,
                    analysis_title=analysis_title,
                    problem_statement="",
                    problem_validated=False,
                )
                .on_conflict_do_nothing(index_elements=[AnalysisStateRow.id])
            )
            await session.execute(stmt)
            await session.commit()
            row = await session.get(AnalysisStateRow, analysis_id)
            return _row_to_state(row)

    async def create(
        self,
        *,
        analysis_id: str,
        user_id: str,
        analysis_title: str = "New analysis",
        problem_statement: str = "",
    ) -> AnalysisState:
        """Create the state row for a new analysis (id shared with its chat room)."""
        async with AsyncSessionLocal() as session:
            row = AnalysisStateRow(
                id=analysis_id,
                user_id=user_id,
                analysis_title=analysis_title,
                problem_statement=problem_statement,
                problem_validated=False,
            )
            session.add(row)
            await session.commit()
            await session.refresh(row)
            return _row_to_state(row)

    async def update(
        self,
        analysis_id: str,
        *,
        problem_statement: str | None = None,
        problem_validated: bool | None = None,
        report_id: str | None = None,
    ) -> AnalysisState | None:
        """Patch the given fields (only non-None args are written). Returns the row.

        Used by the Problem Statement skill (`problem_validated`) and the report
        flow (`report_id`). Returns None if the analysis doesn't exist.
        """
        async with AsyncSessionLocal() as session:
            row = await session.get(AnalysisStateRow, analysis_id)
            if row is None:
                logger.warning(
                    "analysis row missing — update skipped",
                    analysis_id=analysis_id,
                )
                return None
            if problem_statement is not None:
                row.problem_statement = problem_statement
            if problem_validated is not None:
                row.problem_validated = problem_validated
            if report_id is not None:
                row.report_id = report_id
            await session.commit()
            await session.refresh(row)
            return _row_to_state(row)