src/traceability/store.py

"""TraceabilityStore — the seam the chat pipeline persists provenance through (KM-691).

`ChatHandler` (and the v2 chat endpoint's greeting/cache branches) flush one
`TraceabilityPayload` per assistant turn through this seam, right before the `done`
SSE event; `GET /api/v1/traceability` reads it back by (analysis_id, message_id).

- `NullTraceabilityStore` logs and stores nothing (tests / disabled persistence).
- `PostgresTraceabilityStore` writes one `message_traceability` row per turn
  (dedorch, `AsyncSessionLocal`), mirroring `PostgresReportInputStore`.

`save` must NEVER raise on the caller's path — a persistence failure must not break
the user's answer. `get` is the endpoint read and returns `None` on a miss (→ 404).
"""

from __future__ import annotations

from typing import Protocol, runtime_checkable

from sqlalchemy import select
from sqlalchemy.dialects.postgresql import insert

from src.db.postgres.connection import AsyncSessionLocal
from src.db.postgres.models import MessageTraceabilityRow
from src.middlewares.logging import get_logger

from .schemas import TraceabilityPayload

logger = get_logger("traceability_store")


@runtime_checkable
class TraceabilityStore(Protocol):
    """Persist + read one provenance record per assistant `message_id`.

    `save` must never raise on the caller's path. `get` returns the payload for one
    turn, or `None` if none exists yet (the turn is still running or the id is unknown).
    """

    async def save(self, payload: TraceabilityPayload) -> None: ...

    async def get(
        self, analysis_id: str, message_id: str, user_id: str | None = None
    ) -> TraceabilityPayload | None: ...


class NullTraceabilityStore:
    """No-op store: logs the payload, persists nothing. Reads return `None`."""

    async def save(self, payload: TraceabilityPayload) -> None:
        logger.info(
            "traceability produced (not persisted — NullTraceabilityStore)",
            message_id=payload.message_id,
            intent=payload.intent,
            n_tool_calls=len(payload.tool_calls),
        )

    async def get(
        self, analysis_id: str, message_id: str, user_id: str | None = None
    ) -> TraceabilityPayload | None:
        return None


class PostgresTraceabilityStore:
    """Writes/reads `message_traceability` jsonb rows. Upsert on `message_id`."""

    async def save(self, payload: TraceabilityPayload) -> None:
        try:
            data = payload.model_dump(mode="json", by_alias=True)
            async with AsyncSessionLocal() as session:
                stmt = insert(MessageTraceabilityRow).values(
                    message_id=payload.message_id,
                    analysis_id=payload.analysis_id,
                    user_id=payload.user_id,
                    intent=payload.intent,
                    data=data,
                )
                # Idempotent: a re-flushed turn overwrites its own row.
                stmt = stmt.on_conflict_do_update(
                    index_elements=[MessageTraceabilityRow.message_id],
                    set_={"data": stmt.excluded.data, "intent": stmt.excluded.intent},
                )
                await session.execute(stmt)
                await session.commit()
            logger.info(
                "traceability persisted",
                message_id=payload.message_id,
                analysis_id=payload.analysis_id,
                intent=payload.intent,
            )
        except Exception as exc:  # never break the user's answer
            logger.error(
                "traceability persist failed",
                degraded_seam="traceability_persist",
                message_id=payload.message_id,
                error=repr(exc),
            )

    async def get(
        self, analysis_id: str, message_id: str, user_id: str | None = None
    ) -> TraceabilityPayload | None:
        """One turn's payload, or None on a miss.

        `user_id` scopes the read to the turn's owner when a caller supplies one.

        **It stays unused by the endpoint — lead decision, reaffirmed 2026-07-23.**
        `GET /api/v1/traceability` looks up by `(analysis_id, message_id)` alone, and
        `GET /api/v1/charts` by `message_id` alone (the 2026-07-13 decision). Review
        finding F-3 proposed adding the parameter to both; that was **declined**, so
        do not "finish" this by wiring it into the endpoint — reopening it needs the
        same sign-off path as any other locked decision.

        The consequence is deliberate and accepted: this row is reachable by
        `(analysis_id, message_id)` alone, and the payload itself carries `user_id`.
        **The service-secret gate is the control that protects these endpoints**, not
        this predicate — which is why `dataeyond__service__secret` being set matters
        more than it would otherwise (DEV_PLAN #37).

        The parameter is kept because it costs nothing and a future Go-forwarded
        identity (#43) would use it.
        """
        async with AsyncSessionLocal() as session:
            where = [
                MessageTraceabilityRow.message_id == message_id,
                MessageTraceabilityRow.analysis_id == analysis_id,
            ]
            if user_id is not None:
                where.append(MessageTraceabilityRow.user_id == user_id)
            result = await session.execute(
                select(MessageTraceabilityRow.data).where(*where)
            )
            row = result.scalar_one_or_none()
        if row is None:
            return None
        return TraceabilityPayload.model_validate(row)