inquiry_robot/inquiry-agent/agent/runtime_db/session_store.py

"""
会话仓储壳：conversation_session 读写。

本文件职责：按 sender_id / thread_id 存取活动工单与打断栈。
禁止：写入主账六态；本壳不做询价业务分支。
线程：短连接；Memory 后端仅自检。
"""

from __future__ import annotations

import logging
import threading
from dataclasses import dataclass, field
from datetime import datetime, timezone
from typing import Any, Optional

from agent.runtime_db import assert_safe_database, connect_psycopg

logger = logging.getLogger(__name__)


def _utcnow() -> datetime:
    return datetime.now(timezone.utc)


@dataclass
class SessionRecord:
    """一条会话快照（非六态真相）。"""

    thread_id: str
    sender_id: str
    active_inquiry_no: str = ""
    interrupt_stack: list[Any] = field(default_factory=list)
    meta: dict[str, Any] = field(default_factory=dict)


class MemorySessionStore:
    """内存会话壳（自检 / 无 PG）。"""

    def __init__(self) -> None:
        self._by_thread: dict[str, SessionRecord] = {}
        self._lock = threading.Lock()

    def get(self, thread_id: str) -> Optional[SessionRecord]:
        with self._lock:
            rec = self._by_thread.get(thread_id)
            return None if rec is None else SessionRecord(
                thread_id=rec.thread_id,
                sender_id=rec.sender_id,
                active_inquiry_no=rec.active_inquiry_no,
                interrupt_stack=list(rec.interrupt_stack),
                meta=dict(rec.meta),
            )

    def upsert(self, record: SessionRecord) -> None:
        with self._lock:
            self._by_thread[record.thread_id] = SessionRecord(
                thread_id=record.thread_id,
                sender_id=record.sender_id,
                active_inquiry_no=record.active_inquiry_no,
                interrupt_stack=list(record.interrupt_stack),
                meta=dict(record.meta),
            )
        logger.info(
            "session.memory.upsert thread=%s sender=%s",
            record.thread_id,
            record.sender_id,
        )

    def find_by_sender(self, sender_id: str) -> list[SessionRecord]:
        with self._lock:
            return [
                SessionRecord(
                    thread_id=r.thread_id,
                    sender_id=r.sender_id,
                    active_inquiry_no=r.active_inquiry_no,
                    interrupt_stack=list(r.interrupt_stack),
                    meta=dict(r.meta),
                )
                for r in self._by_thread.values()
                if r.sender_id == sender_id
            ]


class PostgresSessionStore:
    """
    PG 会话壳。

    须已执行 001 迁移；读写 conversation_session，无业务推导。
    """

    def __init__(self, dsn: str) -> None:
        assert_safe_database(dsn)
        self._dsn = dsn

    def get(self, thread_id: str) -> Optional[SessionRecord]:
        with connect_psycopg(self._dsn) as conn:
            row = conn.execute(
                """
                SELECT thread_id, sender_id, active_inquiry_no, interrupt_stack, meta
                FROM conversation_session WHERE thread_id = %s
                """,
                (thread_id,),
            ).fetchone()
        if not row:
            return None
        return SessionRecord(
            thread_id=str(row[0]),
            sender_id=str(row[1]),
            active_inquiry_no=str(row[2] or ""),
            interrupt_stack=list(row[3] or []),
            meta=dict(row[4] or {}),
        )

    def upsert(self, record: SessionRecord) -> None:
        import json

        with connect_psycopg(self._dsn) as conn:
            conn.execute(
                """
                INSERT INTO conversation_session (
                    thread_id, sender_id, active_inquiry_no, interrupt_stack, meta, updated_at
                ) VALUES (%s, %s, %s, %s::jsonb, %s::jsonb, %s)
                ON CONFLICT (thread_id) DO UPDATE SET
                    sender_id = EXCLUDED.sender_id,
                    active_inquiry_no = EXCLUDED.active_inquiry_no,
                    interrupt_stack = EXCLUDED.interrupt_stack,
                    meta = EXCLUDED.meta,
                    updated_at = EXCLUDED.updated_at
                """,
                (
                    record.thread_id,
                    record.sender_id,
                    record.active_inquiry_no,
                    json.dumps(record.interrupt_stack, ensure_ascii=False),
                    json.dumps(record.meta, ensure_ascii=False),
                    _utcnow(),
                ),
            )
            conn.commit()
        logger.info("session.pg.upsert thread=%s", record.thread_id)

    def find_by_sender(self, sender_id: str) -> list[SessionRecord]:
        with connect_psycopg(self._dsn) as conn:
            rows = conn.execute(
                """
                SELECT thread_id, sender_id, active_inquiry_no, interrupt_stack, meta
                FROM conversation_session WHERE sender_id = %s
                """,
                (sender_id,),
            ).fetchall()
        return [
            SessionRecord(
                thread_id=str(r[0]),
                sender_id=str(r[1]),
                active_inquiry_no=str(r[2] or ""),
                interrupt_stack=list(r[3] or []),
                meta=dict(r[4] or {}),
            )
            for r in rows
        ]


def create_session_store(*, dsn: str = "", backend: str = "auto"):
    """
    工厂：有 DSN 且 backend!=memory → PG；否则 Memory。

    框架壳，不自动建业务会话。
    """
    if backend == "memory" or not dsn:
        return MemorySessionStore()
    return PostgresSessionStore(dsn)